[RISCV] Sink vp.splat operands of VP intrinsic.

Author: NexMing

llvm/lib/Target/RISCV/RISCVTargetTransformInfo.cpp

@@ -2837,39 +2837,43 @@ bool RISCVTTIImpl::isProfitableToSinkOperands(
   if (!ST->sinkSplatOperands())
     return false;
 
-  for (auto OpIdx : enumerate(I->operands())) {
-    if (!canSplatOperand(I, OpIdx.index()))
-      continue;
-
-    Instruction *Op = dyn_cast<Instruction>(OpIdx.value().get());
+  for (auto &U : I->operands()) {
+    auto *Op = dyn_cast<Instruction>(U.get());
     // Make sure we are not already sinking this operand
     if (!Op || any_of(Ops, [&](Use *U) { return U->get() == Op; }))
       continue;
 
-    // We are looking for a splat that can be sunk.
-    if (!match(Op, m_Shuffle(m_InsertElt(m_Undef(), m_Value(), m_ZeroInt()),
+    // We are looking for a splat/vp.splat that can be sunk.
+    // FIXME: Should we care about the poison value of inactive elements in
+    // vp.splat?
+    bool IsVPSplat = match(Op, m_Intrinsic<Intrinsic::experimental_vp_splat>(
+                                   m_Value(), m_Value(), m_Value()));
+    if (!IsVPSplat &&
+        !match(Op, m_Shuffle(m_InsertElt(m_Undef(), m_Value(), m_ZeroInt()),
                              m_Undef(), m_ZeroMask())))
       continue;
 
     // Don't sink i1 splats.
     if (cast<VectorType>(Op->getType())->getElementType()->isIntegerTy(1))
       continue;
 
-    // All uses of the shuffle should be sunk to avoid duplicating it across gpr
-    // and vector registers
-    for (Use &U : Op->uses()) {
-      Instruction *Insn = cast<Instruction>(U.getUser());
-      if (!canSplatOperand(Insn, U.getOperandNo()))
-        return false;
-    }
+    // All uses of the splat/vp.splat should be sunk to avoid duplicating it
+    // across gpr and vector registers
+    if (any_of(Op->uses(), [this](Use &U) {
+          return !canSplatOperand(cast<Instruction>(U.getUser()),
+                                  U.getOperandNo());
+        }))
+      continue;
 
-    Use *InsertEltUse = &Op->getOperandUse(0);
-    // Sink any fpexts since they might be used in a widening fp pattern.
-    auto *InsertElt = cast<InsertElementInst>(InsertEltUse);
-    if (isa<FPExtInst>(InsertElt->getOperand(1)))
-      Ops.push_back(&InsertElt->getOperandUse(1));
-    Ops.push_back(InsertEltUse);
-    Ops.push_back(&OpIdx.value());
+    if (!IsVPSplat) {
+      Use *InsertEltUse = &Op->getOperandUse(0);
+      // Sink any fpexts since they might be used in a widening fp pattern.
+      auto *InsertElt = cast<InsertElementInst>(InsertEltUse);
+      if (isa<FPExtInst>(InsertElt->getOperand(1)))
+        Ops.push_back(&InsertElt->getOperandUse(1));
+      Ops.push_back(InsertEltUse);
+    }
+    Ops.push_back(&U);
   }
   return true;
 }

llvm/test/CodeGen/RISCV/rvv/sink-splat-operands.ll

@@ -5896,30 +5896,30 @@ define void @sink_vp_splat(ptr nocapture %out, ptr nocapture %in) {
 ; CHECK:       # %bb.0: # %entry
 ; CHECK-NEXT:    li a2, 0
 ; CHECK-NEXT:    li a3, 1024
-; CHECK-NEXT:    lui a4, 1
+; CHECK-NEXT:    li a4, 3
+; CHECK-NEXT:    lui a5, 1
 ; CHECK-NEXT:  .LBB129_1: # %vector.body
 ; CHECK-NEXT:    # =>This Loop Header: Depth=1
 ; CHECK-NEXT:    # Child Loop BB129_2 Depth 2
-; CHECK-NEXT:    vsetvli a5, a3, e32, m4, ta, ma
-; CHECK-NEXT:    slli a6, a2, 2
+; CHECK-NEXT:    vsetvli a6, a3, e32, m4, ta, ma
+; CHECK-NEXT:    slli a7, a2, 2
 ; CHECK-NEXT:    vmv.v.i v8, 0
-; CHECK-NEXT:    vmv.v.i v12, 3
-; CHECK-NEXT:    add a7, a1, a6
-; CHECK-NEXT:    li t0, 1024
+; CHECK-NEXT:    add t0, a1, a7
+; CHECK-NEXT:    li t1, 1024
 ; CHECK-NEXT:  .LBB129_2: # %for.body424
 ; CHECK-NEXT:    # Parent Loop BB129_1 Depth=1
 ; CHECK-NEXT:    # => This Inner Loop Header: Depth=2
-; CHECK-NEXT:    vle32.v v16, (a7)
-; CHECK-NEXT:    addi t0, t0, -1
-; CHECK-NEXT:    vmacc.vv v8, v16, v12
-; CHECK-NEXT:    add a7, a7, a4
-; CHECK-NEXT:    bnez t0, .LBB129_2
+; CHECK-NEXT:    vle32.v v12, (t0)
+; CHECK-NEXT:    addi t1, t1, -1
+; CHECK-NEXT:    vmacc.vx v8, a4, v12
+; CHECK-NEXT:    add t0, t0, a5
+; CHECK-NEXT:    bnez t1, .LBB129_2
 ; CHECK-NEXT:  # %bb.3: # %vector.latch
 ; CHECK-NEXT:    # in Loop: Header=BB129_1 Depth=1
-; CHECK-NEXT:    add a6, a0, a6
-; CHECK-NEXT:    sub a3, a3, a5
-; CHECK-NEXT:    vse32.v v8, (a6)
-; CHECK-NEXT:    add a2, a2, a5
+; CHECK-NEXT:    add a7, a0, a7
+; CHECK-NEXT:    sub a3, a3, a6
+; CHECK-NEXT:    vse32.v v8, (a7)
+; CHECK-NEXT:    add a2, a2, a6
 ; CHECK-NEXT:    bnez a3, .LBB129_1
 ; CHECK-NEXT:  # %bb.4: # %for.cond.cleanup
 ; CHECK-NEXT:    ret