[WebAssembly] Optimize vector shift using a splat value from outside block

[YolandaCY]

Summary

The vector shift operation in WebAssembly uses an i32 shift amount type,
while the LLVM IR requires binary operator uses the same type of operands.
When the shift amount operand is splated from a different block, the splat source
will not be exported and the vector shift will be unrolled to scalar shifts.
This patch enables the vector shift to identify the splat source value from
the other block, and generate expected WebAssembly bytecode when lowering.

[YolandaCY]

This is to resolve a WebAssembly codegen issue when vector shift is used in a loop, while the shift amount is initialized outside the loop. Could you help take a look? Thanks!

[tlively]

Thanks for the patch! It looks like this will be a nice improvement.

@craig.topper, it would be great to get your comments as well, as someone more familiar with the target-independent infrastructure here.

llvm/include/llvm/CodeGen/SelectionDAG.h

2459–2460

It would be good to add a comment describing the contents of this map.

llvm/lib/Target/WebAssembly/WebAssemblyISelLowering.cpp

840

Why allow null instruction pointers here? It would seem simpler to ensure that callers pass a valid instruction and to assume we have a valid instruction here.

844

What is the benefit of including isShiftAmountScalar() here, given that we know it is always true?

llvm/test/CodeGen/WebAssembly/simd-shift-in-loop.ll

2

What do you think about using the auto-update script for this test? The output would be more verbose, but it would also be easy to update if anything changes, and I think it would be helpful to see that the whole function is emitted correctly.

17

Can we add a test where the vshift is a phi, just to show that that still works correctly?

[YolandaCY]

Add test and comments

[YolandaCY]

Thank you Thomas for the comments! Please see my updates in the new revision.

llvm/lib/Target/WebAssembly/WebAssemblyISelLowering.cpp

840

This is a fast check on the outside block when visit the splat vector, and don't know yet if the splat vector will be used by a shift op.
To identify the instruction I need to iterate all uses of the splat vector until we find the vector shift. Since this is only needed for WebAssembly target, I add a quick check here to reduce the cost for other platforms. Seems a little confusion, do you think we need to seperate it to two functions?

844

This will be called in SelectionDAGBuilder to skip the optimizaiton for other platforms when visit shift.

llvm/test/CodeGen/WebAssembly/simd-shift-in-loop.ll

2

Sure. Previously I make it simple to avoid mismatch on unrelated changes. But if we have an auto-update script that would be helpful to verify the whole function directly.

17

Sure. I have added one more test.

[tlively]

llvm/lib/Target/WebAssembly/WebAssemblyISelLowering.cpp

840

Oh I see, that makes sense.

844

We know that the isShiftAmountScalar() here will be the WebAssemblyTargetLowering version of isShiftAmountScalar(), so it will always be true, right? So this line could be:

return I->isShift() && I->getOperand(1) == Splat;

[craig.topper]

Could this use the TargetLowering::shouldSinkOperands hook to get CodeGenPrepare to move the splat into the loop. ARM, X86, and RISC-V all do that.

[YolandaCY]

Use the shouldSinkOperands hook in CodeGenPrepare

[YolandaCY]

In D158399#4607933, @craig.topper wrote:

Could this use the TargetLowering::shouldSinkOperands hook to get CodeGenPrepare to move the splat into the loop. ARM, X86, and RISC-V all do that.

Thanks for the suggestion! I have revised the code to use this existing hook. Please help take a look again. @craig.topper @tlively Thanks!

[tlively]

Nice, this LGTM. Thanks for the tip, @craig.topper!

[YolandaCY]

In D158399#4613902, @tlively wrote:

Nice, this LGTM. Thanks for the tip, @craig.topper!

Thanks Thomas. Could you help me commit this change?

[tlively]

Sure thing. I'll use your author email from the other patches in your phabricator profile.

[YolandaCY]

In D158399#4617129, @tlively wrote:

Sure thing. I'll use your author email from the other patches in your phabricator profile.

OK, Thank you!

llvm/include/llvm/CodeGen/SelectionDAG.h

#endif
   /// either a BUILD_VECTOR or SPLAT_VECTOR depending on the
   /// scalability of the desired vector type.
   SDValue getSplat(EVT VT, const SDLoc &DL, SDValue Op) {
     assert(VT.isVector() && "Can't splat to non-vector type");
     return VT.isScalableVector() ?
       getSplatVector(VT, DL, Op) : getSplatBuildVector(VT, DL, Op);
   }
 
+  /// Returns an exported splat source from another block. This helps the
+  /// WebAssembly target lowering for vector shift operation where i32 is used
+  /// as shift amount value type.
+  const Value *getExportedSplatSource(const SDNode *N, Register &Reg) const {
+    auto I = ExportedSplatValueMap.find(N);
+    if (I != ExportedSplatValueMap.end()) {
+      Reg = I->second.second;
+      return I->second.first;
+    }
+    return nullptr;
+  }
+
+  /// Set exported splat source mapping to the splat node.
+  void addExportedSplatSource(const SDNode *N, const Value *V,
+                              const unsigned Reg) {
+    ExportedSplatValueMap[N] = {V, Reg};
+  }
+
   /// Returns a vector of type ResVT whose elements contain the linear sequence
   ///   <0, Step, Step * 2, Step * 3, ...>
   SDValue getStepVector(const SDLoc &DL, EVT ResVT, APInt StepVal);
 
   /// Returns a vector of type ResVT whose elements contain the linear sequence
   ///   <0, 1, 2, 3, ...>
   SDValue getStepVector(const SDLoc &DL, EVT ResVT);
 
private:
   std::vector<CondCodeSDNode*> CondCodeNodes;
 
   std::vector<SDNode*> ValueTypeNodes;
   std::map<EVT, SDNode*, EVT::compareRawBits> ExtendedValueTypeNodes;
   StringMap<SDNode*> ExternalSymbols;
 
   std::map<std::pair<std::string, unsigned>, SDNode *> TargetExternalSymbols;
   DenseMap<MCSymbol *, SDNode *> MCSymbols;
 
+  /// Maps the splat source value to splat vector when they are exported from
+  /// external block. This information is needed for later WebAssembly lowering.
+  DenseMap<const SDNode *, std::pair<const Value *, unsigned>>
+      ExportedSplatValueMap;
+
   FlagInserter *Inserter = nullptr;
 };
 
 template <> struct GraphTraits<SelectionDAG*> : public GraphTraits<SDNode*> {
   using nodes_iterator = pointer_iterator<SelectionDAG::allnodes_iterator>;
 
   static nodes_iterator nodes_begin(SelectionDAG *G) {
     return nodes_iterator(G->allnodes_begin());

llvm/include/llvm/CodeGen/TargetLowering.h

public:
   /// Returns true if the target can instruction select the specified FP
   /// immediate natively. If false, the legalizer will materialize the FP
   /// immediate as a load from a constant pool.
   virtual bool isFPImmLegal(const APFloat & /*Imm*/, EVT /*VT*/,
                             bool ForCodeSize = false) const {
     return false;
   }
 
+  virtual bool isShiftAmountScalar() const { return false; }
+
+  virtual bool hasSplatValueUseForVectorOp(const Instruction *I = nullptr,
+                                           const Value *Splat = nullptr) const {
+    return false;
+  }
+
   /// Targets can use this to indicate that they only support *some*
   /// VECTOR_SHUFFLE operations, those with specific masks.  By default, if a
   /// target supports the VECTOR_SHUFFLE node, all mask values are assumed to be
   /// legal.
   virtual bool isShuffleMaskLegal(ArrayRef<int> /*Mask*/, EVT /*VT*/) const {
     return true;
   }
 

llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp

void SelectionDAG::clear() {
   OperandRecycler.clear(OperandAllocator);
   OperandAllocator.Reset();
   CSEMap.clear();
 
   ExtendedValueTypeNodes.clear();
   ExternalSymbols.clear();
   TargetExternalSymbols.clear();
   MCSymbols.clear();
+  ExportedSplatValueMap.clear();
   SDEI.clear();
   std::fill(CondCodeNodes.begin(), CondCodeNodes.end(),
             static_cast<CondCodeSDNode*>(nullptr));
   std::fill(ValueTypeNodes.begin(), ValueTypeNodes.end(),
             static_cast<SDNode*>(nullptr));
 
   EntryNode.UseList = nullptr;
   InsertNode(&EntryNode);

llvm/lib/CodeGen/SelectionDAG/SelectionDAGBuilder.h

void EmitBranchForMergedCondition(const Value *Cond, MachineBasicBlock *TBB,
                                     MachineBasicBlock *FBB,
                                     MachineBasicBlock *CurBB,
                                     MachineBasicBlock *SwitchBB,
                                     BranchProbability TProb, BranchProbability FProb,
                                     bool InvertCond);
   bool ShouldEmitAsBranches(const std::vector<SwitchCG::CaseBlock> &Cases);
   bool isExportableFromCurrentBlock(const Value *V, const BasicBlock *FromBB);
   void CopyToExportRegsIfNeeded(const Value *V);
+  Register GetExportReg(const Value *V);
   void ExportFromCurrentBlock(const Value *V);
   void LowerCallTo(const CallBase &CB, SDValue Callee, bool IsTailCall,
                    bool IsMustTailCall, const BasicBlock *EHPadBB = nullptr);
 
   // Lower range metadata from 0 to N to assert zext to an integer of nearest
   // floor power of two.
   SDValue lowerRangeToAssertZExt(SelectionDAG &DAG, const Instruction &I,
                                  SDValue Op);

llvm/lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp

RegsForValue::RegsForValue(const SmallVector<unsigned, 4> &regs, MVT regvt,
                            EVT valuevt, std::optional<CallingConv::ID> CC)
     : ValueVTs(1, valuevt), RegVTs(1, regvt), Regs(regs),
       RegCount(1, regs.size()), CallConv(CC) {}
 
 RegsForValue::RegsForValue(LLVMContext &Context, const TargetLowering &TLI,
                            const DataLayout &DL, unsigned Reg, Type *Ty,
                            std::optional<CallingConv::ID> CC) {
   ComputeValueVTs(TLI, DL, Ty, ValueVTs);
-
   CallConv = CC;
 
   for (EVT ValueVT : ValueVTs) {
     unsigned NumRegs =
         isABIMangled()
             ? TLI.getNumRegistersForCallingConv(Context, *CC, ValueVT)
             : TLI.getNumRegisters(Context, ValueVT);
     MVT RegisterVT =
void SelectionDAGBuilder::visitRet(const ReturnInst &I) {
   // Update the DAG with the new chain value resulting from return lowering.
   DAG.setRoot(Chain);
 }
 
 /// CopyToExportRegsIfNeeded - If the given value has virtual registers
 /// created for it, emit nodes to copy the value into the virtual
 /// registers.
 void SelectionDAGBuilder::CopyToExportRegsIfNeeded(const Value *V) {
+  if (Register reg = GetExportReg(V)) {
+    assert((!V->use_empty() || isa<CallBrInst>(V)) &&
+           "Unused value assigned virtual registers!");
+    CopyValueToVirtualRegister(V, reg);
+  }
+}
+
+Register SelectionDAGBuilder::GetExportReg(const Value *V) {
   // Skip empty types
   if (V->getType()->isEmptyTy())
-    return;
+    return Register();
 
   DenseMap<const Value *, Register>::iterator VMI = FuncInfo.ValueMap.find(V);
-  if (VMI != FuncInfo.ValueMap.end()) {
-    assert((!V->use_empty() || isa<CallBrInst>(V)) &&
-           "Unused value assigned virtual registers!");
-    CopyValueToVirtualRegister(V, VMI->second);
-  }
+  return (VMI != FuncInfo.ValueMap.end()) ? VMI->second : Register();
 }
 
 /// ExportFromCurrentBlock - If this condition isn't known to be exported from
 /// the current basic block, add it to ValueMap now so that we'll get a
 /// CopyTo/FromReg.
 void SelectionDAGBuilder::ExportFromCurrentBlock(const Value *V) {
   // No need to export constants.
   if (!isa<Instruction>(V) && !isa<Argument>(V)) return;
SDValue BinNodeValue = DAG.getNode(Opcode, getCurSDLoc(), Op1.getValueType(),
                                      Op1, Op2, Flags);
   setValue(&I, BinNodeValue);
 }
 
 void SelectionDAGBuilder::visitShift(const User &I, unsigned Opcode) {
   SDValue Op1 = getValue(I.getOperand(0));
   SDValue Op2 = getValue(I.getOperand(1));
 
-  EVT ShiftTy = DAG.getTargetLoweringInfo().getShiftAmountTy(
-      Op1.getValueType(), DAG.getDataLayout());
+  auto &TLI = DAG.getTargetLoweringInfo();
+  Value *V2 = I.getOperand(1);
+  if (I.getType()->isVectorTy() && TLI.isShiftAmountScalar() &&
+      !(NodeMap[V2].getNode())) {
+    const Value *splat = getSplatValue(V2);
+    if (splat && !isa<Constant>(splat)) {
+      assert(FuncInfo.isExportedInst(splat) && "Splat value is not exported");
+      // TODO: It's possible to be mapped to multiple splat vectors.
+      DenseMap<const Value *, Register>::iterator It =
+          FuncInfo.ValueMap.find(splat);
+      if (It != FuncInfo.ValueMap.end()) {
+        DAG.addExportedSplatSource(Op2.getNode(), splat, It->second);
+      }
+    }
+  }
+
+  EVT ShiftTy = TLI.getShiftAmountTy(Op1.getValueType(), DAG.getDataLayout());
 
   // Coerce the shift amount to the right type if we can. This exposes the
   // truncate or zext to optimization early.
   if (!I.getType()->isVectorTy() && Op2.getValueType() != ShiftTy) {
     assert(ShiftTy.getSizeInBits() >= Log2_32_Ceil(Op1.getValueSizeInBits()) &&
            "Unexpected shift type");
     Op2 = DAG.getZExtOrTrunc(Op2, getCurSDLoc(), ShiftTy);
   }
void SelectionDAGBuilder::visitShuffleVector(const User &I) {
   // The DAGCombiner will perform a BUILD_VECTOR -> SPLAT_VECTOR transformation
   // for targets that support a SPLAT_VECTOR for non-scalable vector types.
   assert(!VT.isScalableVector() && "Unsupported scalable vector shuffle");
 
   unsigned SrcNumElts = SrcVT.getVectorNumElements();
   unsigned MaskNumElts = Mask.size();
 
   if (SrcNumElts == MaskNumElts) {
+    if (TLI.hasSplatValueUseForVectorOp() && GetExportReg(&I)) {
+      if (const Value *splat = getSplatValue(&I)) {
+        for (const Use &U : I.uses()) {
+          Instruction *UserI = cast<Instruction>(U.getUser());
+          if (UserI->getType()->isVectorTy() &&
+              TLI.hasSplatValueUseForVectorOp(UserI, &I)) {
+            ExportFromCurrentBlock(splat);
+          }
+        }
+      }
+    }
     setValue(&I, DAG.getVectorShuffle(VT, DL, Src1, Src2, Mask));
     return;
   }
 
   // Normalize the shuffle vector since mask and vector length don't match.
   if (SrcNumElts < MaskNumElts) {
     // Mask is longer than the source vectors. We can use concatenate vector to
     // make the mask and vectors lengths match.

llvm/lib/Target/WebAssembly/WebAssemblyISelLowering.h

bool isLegalAddressingMode(const DataLayout &DL, const AddrMode &AM, Type *Ty,
                              unsigned AS,
                              Instruction *I = nullptr) const override;
   bool allowsMisalignedMemoryAccesses(EVT, unsigned AddrSpace, Align Alignment,
                                       MachineMemOperand::Flags Flags,
                                       unsigned *Fast) const override;
   bool isIntDivCheap(EVT VT, AttributeList Attr) const override;
   bool isVectorLoadExtDesirable(SDValue ExtVal) const override;
   bool isOffsetFoldingLegal(const GlobalAddressSDNode *GA) const override;
+  bool isShiftAmountScalar() const override;
+  bool hasSplatValueUseForVectorOp(const Instruction *I = nullptr,
+                                   const Value *Splat = nullptr) const override;
   EVT getSetCCResultType(const DataLayout &DL, LLVMContext &Context,
                          EVT VT) const override;
   bool getTgtMemIntrinsic(IntrinsicInfo &Info, const CallInst &I,
                           MachineFunction &MF,
                           unsigned Intrinsic) const override;
 
   void computeKnownBitsForTargetNode(const SDValue Op, KnownBits &Known,
                                      const APInt &DemandedElts,

llvm/lib/Target/WebAssembly/WebAssemblyISelLowering.cpp

 
 bool WebAssemblyTargetLowering::isOffsetFoldingLegal(
     const GlobalAddressSDNode *GA) const {
   // Wasm doesn't support function addresses with offsets
   const GlobalValue *GV = GA->getGlobal();
   return isa<Function>(GV) ? false : TargetLowering::isOffsetFoldingLegal(GA);
 }
 
+bool WebAssemblyTargetLowering::isShiftAmountScalar() const { return true; }
+
+bool WebAssemblyTargetLowering::hasSplatValueUseForVectorOp(
+    const Instruction *I, const Value *Splat) const {
+  if (!I) {
+    return isShiftAmountScalar();
+  }
+
+  return I->isShift() && isShiftAmountScalar() && I->getOperand(1) == Splat;
+}
+
 EVT WebAssemblyTargetLowering::getSetCCResultType(const DataLayout &DL,
                                                   LLVMContext &C,
                                                   EVT VT) const {
   if (VT.isVector())
     return VT.changeVectorElementTypeToInteger();
 
   // So far, all branch instructions in Wasm take an I32 condition.
   // The default TargetLowering::getSetCCResultType returns the pointer size,
if (MaskOp.getValueType().isVector()) {
         MaskOp = LHS;
     }
 
     return MaskOp;
   };
 
   // Skip vector and operation
   ShiftVal = SkipImpliedMask(ShiftVal, LaneBits - 1);
+  if (ShiftVal.getValueType().isVector()) {
+    auto SavedShiftVal = ShiftVal;
     ShiftVal = DAG.getSplatValue(ShiftVal);
-  if (!ShiftVal)
+    if (!ShiftVal) {
+      Register InReg;
+      if (auto splat =
+              DAG.getExportedSplatSource(SavedShiftVal.getNode(), InReg)) {
+        EVT RegisterVT = getRegisterType(
+            splat->getContext(),
+            getValueType(DAG.getDataLayout(), splat->getType()));
+        ShiftVal =
+            DAG.getCopyFromReg(DAG.getEntryNode(), DL, InReg, RegisterVT);
+      }
+
+      if (!ShiftVal) {
         return unrollVectorShift(Op, DAG);
+      }
+    }
+  }
 
   // Skip scalar and operation
   ShiftVal = SkipImpliedMask(ShiftVal, LaneBits - 1);
   // Use anyext because none of the high bits can affect the shift
   ShiftVal = DAG.getAnyExtOrTrunc(ShiftVal, DL, MVT::i32);
 
   unsigned Opcode;
   switch (Op.getOpcode()) {

llvm/test/CodeGen/WebAssembly/simd-shift-in-loop.ll

+; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py UTC_ARGS: --version 2
+; RUN: llc < %s -verify-machineinstrs -mattr=+simd128 | FileCheck %s
+
+; Test that SIMD shifts can be lowered correctly even when shift
+; values are exported from outside blocks.
+
+target triple = "wasm32-unknown-unknown"
+
+define void @shl_loop(ptr %a, i8 %shift, i32 %count) {
+; CHECK-LABEL: shl_loop:
+; CHECK:         .functype shl_loop (i32, i32, i32) -> ()
+; CHECK-NEXT:  # %bb.0: # %entry
+; CHECK-NEXT:  .LBB0_1: # %body
+; CHECK-NEXT:    # =>This Inner Loop Header: Depth=1
+; CHECK-NEXT:    loop # label0:
+; CHECK-NEXT:    local.get 0
+; CHECK-NEXT:    local.get 0
+; CHECK-NEXT:    v128.load 0:p2align=0
+; CHECK-NEXT:    local.get 1
+; CHECK-NEXT:    i8x16.shl
+; CHECK-NEXT:    v128.store 16
+; CHECK-NEXT:    local.get 0
+; CHECK-NEXT:    i32.const 16
+; CHECK-NEXT:    i32.add
+; CHECK-NEXT:    local.set 0
+; CHECK-NEXT:    local.get 2
+; CHECK-NEXT:    i32.const -1
+; CHECK-NEXT:    i32.add
+; CHECK-NEXT:    local.tee 2
+; CHECK-NEXT:    i32.eqz
+; CHECK-NEXT:    br_if 0 # 0: up to label0
+; CHECK-NEXT:  # %bb.2: # %exit
+; CHECK-NEXT:    end_loop
+; CHECK-NEXT:    # fallthrough-return
+entry:
+ %t1 = insertelement <16 x i8> undef, i8 %shift, i32 0
+ %vshift = shufflevector <16 x i8> %t1, <16 x i8> undef, <16 x i32> zeroinitializer
+ br label %body
+body:
+ %out = phi ptr [%a, %entry], [%b, %body]
+ %i = phi i32 [0, %entry], [%next, %body]
+ %v = load <16 x i8>, ptr %out, align 1
+ %r = shl <16 x i8> %v, %vshift
+ %b = getelementptr inbounds i8, ptr %out, i32 16
+ store <16 x i8> %r, ptr %b
+ %next = add i32 %i, 1
+ %i.cmp = icmp eq i32 %next, %count
+ br i1 %i.cmp, label %body, label %exit
+exit:
+ ret void
+}
+
+; Test that SIMD shifts can be lowered correctly when shift value
+; is a phi inside loop body.
+
+define void @shl_phi_loop(ptr %a, i8 %shift, i32 %count) {
+; CHECK-LABEL: shl_phi_loop:
+; CHECK:         .functype shl_phi_loop (i32, i32, i32) -> ()
+; CHECK-NEXT:  # %bb.0: # %entry
+; CHECK-NEXT:  .LBB1_1: # %body
+; CHECK-NEXT:    # =>This Inner Loop Header: Depth=1
+; CHECK-NEXT:    loop # label1:
+; CHECK-NEXT:    local.get 0
+; CHECK-NEXT:    local.get 0
+; CHECK-NEXT:    v128.load 0:p2align=0
+; CHECK-NEXT:    local.get 1
+; CHECK-NEXT:    i8x16.shl
+; CHECK-NEXT:    v128.store 16
+; CHECK-NEXT:    local.get 1
+; CHECK-NEXT:    i32.const 1
+; CHECK-NEXT:    i32.and
+; CHECK-NEXT:    local.set 1
+; CHECK-NEXT:    local.get 0
+; CHECK-NEXT:    i32.const 16
+; CHECK-NEXT:    i32.add
+; CHECK-NEXT:    local.set 0
+; CHECK-NEXT:    local.get 2
+; CHECK-NEXT:    i32.const -1
+; CHECK-NEXT:    i32.add
+; CHECK-NEXT:    local.tee 2
+; CHECK-NEXT:    i32.eqz
+; CHECK-NEXT:    br_if 0 # 0: up to label1
+; CHECK-NEXT:  # %bb.2: # %exit
+; CHECK-NEXT:    end_loop
+; CHECK-NEXT:    # fallthrough-return
+entry:
+ br label %body
+body:
+ %out = phi ptr [%a, %entry], [%b, %body]
+ %i = phi i32 [0, %entry], [%next, %body]
+ %t1 = phi i8 [%shift, %entry], [%sand, %body]
+ %t2 = insertelement <16 x i8> undef, i8 %t1, i32 0
+ %vshift = shufflevector <16 x i8> %t2, <16 x i8> undef, <16 x i32> zeroinitializer
+ %v = load <16 x i8>, ptr %out, align 1
+ %r = shl <16 x i8> %v, %vshift
+ %b = getelementptr inbounds i8, ptr %out, i32 16
+ store <16 x i8> %r, ptr %b
+ %sand = and i8 %t1, 1
+ %next = add i32 %i, 1
+ %i.cmp = icmp eq i32 %next, %count
+ br i1 %i.cmp, label %body, label %exit
+exit:
+ ret void
+}