[DAGCombiner][VP] Add DAGCombine for VP_MUL.

[jacquesguan]

Summary

This patch uses visitMUL to combine VP_MUL, shares most logic of MUL with VP_MUL.

[RKSimon]

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

22804

(style) Don't use auto unless the type is obvious (i.e. a cast<>)

[craig.topper]

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

22804

Don't call isConstOrConstSplat(N1) twice

22806

What about mul by positive 1 or 0?

[jacquesguan]

Address comment.

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

22804

Done.

22806

Done.

[craig.topper]

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

22807

If one of the true/false values of the VP_SELECT is undef can we ignore the mask and EVL and return the other operand? @simoll @frasercrmck

[simoll]

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

22807

I believe so. Also for vp ops other than select or merge, if the mask is undef, we can scratch the operation entirely. If evl is undef, we can cut short to unreachable.

[jacquesguan]

Address comment.

[jacquesguan]

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

22807

Done, thanks.

[craig.topper]

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

22809

I think we should only use getSplatVector for scalable vectors. Fixed vectors should use getSplatBuildVector.

22809

You can just call DAG.getConstant(0, SDLoc(N), VT); It will do the right thing.

22820

DAG.getConstant(0, SDLoc(N), VT)

22828

Same comment as above.

[jacquesguan]

Address comment.

[jacquesguan]

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

22809

Done.

22820

Done.

22828

Done.

[simoll]

@jacquesguan Thanks for working on VP optimizations! We just discussed on the VP syncup call where we want to go with this.
The thing is that this patch is essentially replicating the existing MUL DAGCombiner patterns for VP_MUL. Ideally, we would lift the existing DAGCombiner logic to work on both VP and non-VP SDNodes - automatically.

I've implemented something like this in NEC's open source LLVM stack:
https://github.com/sx-aurora-dev/llvm-project/blob/develop/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

The idea is to templatize the pattern-matching and rewriting code and instantiate that for each VP and non-VP. In essence, https://reviews.llvm.org/D92086 but on the SelectionDAG.

Would you be interested in working on a more general approach? Please contact me directly so we can get you on the VP Discord server or arrange a call to talk about this.

[jacquesguan]

reuse visitMUL to combine VP_MUL.

[jacquesguan]

fix comment.

[craig.topper]

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

3870

IS -> Is

3876

What does it mean to support a VP op in FoldConstantArithmetic?

[jacquesguan]

Address comment.

[jacquesguan]

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

3876

I was thinking wrong. This TODO is removed now.

[jacquesguan]

ping.

[jacquesguan]

ping.

[jacquesguan]

ping.

[jacquesguan]

ping...

[craig.topper]

I'm concerned about the number of checks for IsVP. It looks like maybe this function will be hard to maintain in the future. Every new change will have to think about whether it works for VP. Most LLVM developers don't care about VP so may miss this detail.

[jacquesguan]

In D121187#4582381, @craig.topper wrote:

I'm concerned about the number of checks for IsVP. It looks like maybe this function will be hard to maintain in the future. Every new change will have to think about whether it works for VP. Most LLVM developers don't care about VP so may miss this detail.

Yes, it is a problem, and even we remove all IsVP in the function, a developper that wants to change this function should also think about whether the change works on VP version too. Actually I used a new function visitVP_MUL to implement the combination of VP_MUL at begining, and @simoll suggested to make the existed logic work on both no-vp and vp ops, so I switch to this way. Do you think which one is a better solution, reuse the visitMUL or create a new visitVP_MUL with some repeat code?

[jacquesguan]

ping

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

private:
     SDValue visitMLOAD(SDNode *N);
     SDValue visitMSTORE(SDNode *N);
     SDValue visitMGATHER(SDNode *N);
     SDValue visitMSCATTER(SDNode *N);
     SDValue visitFP_TO_FP16(SDNode *N);
     SDValue visitFP16_TO_FP(SDNode *N);
     SDValue visitVECREDUCE(SDNode *N);
     SDValue visitVPOp(SDNode *N);
+    SDValue visitVP_MUL(SDNode *N);
 
     SDValue visitFADDForFMACombine(SDNode *N);
     SDValue visitFSUBForFMACombine(SDNode *N);
     SDValue visitFMULForFMADistributiveCombine(SDNode *N);
 
     SDValue XformToShuffleWithZero(SDNode *N);
     bool reassociationCanBreakAddressingModePattern(unsigned Opc,
                                                     const SDLoc &DL, SDValue N0,
SDValue DAGCombiner::visitMULFIX(SDNode *N) {
 
   return SDValue();
 }
 
 SDValue DAGCombiner::visitMUL(SDNode *N) {
   SDValue N0 = N->getOperand(0);
   SDValue N1 = N->getOperand(1);
   EVT VT = N0.getValueType();
 
   // fold (mul x, undef) -> 0
   if (N0.isUndef() || N1.isUndef())
     return DAG.getConstant(0, SDLoc(N), VT);
 
   // fold (mul c1, c2) -> c1*c2
   if (SDValue C = DAG.FoldConstantArithmetic(ISD::MUL, SDLoc(N), VT, {N0, N1}))
     return C;
 
   // canonicalize constant to RHS (vector doesn't have to splat)
   if (DAG.isConstantIntBuildVectorOrConstantInt(N0) &&
       !DAG.isConstantIntBuildVectorOrConstantInt(N1))
     return DAG.getNode(ISD::MUL, SDLoc(N), VT, N1, N0);
 
   bool N1IsConst = false;
if ((Opcode == ISD::VECREDUCE_OR &&
         (Opcode == ISD::VECREDUCE_AND &&
          (N0.getOperand(0).isUndef() || isAllOnesOrAllOnesSplat(Vec))))
       return DAG.getNode(Opcode, SDLoc(N), N->getValueType(0), Subvec);
   }
 
   return SDValue();
 }
 
+SDValue DAGCombiner::visitVP_MUL(SDNode *N) {
+  SDValue N0 = N->getOperand(0);
+  SDValue N1 = N->getOperand(1);
+  SDValue Mask = N->getOperand(2);
+  SDValue EVL = N->getOperand(3);
+  EVT VT = N0.getValueType();
+
+  // canonicalize constant splat to RHS
+  if (isConstOrConstSplat(N0) && !isConstOrConstSplat(N1))
+    return DAG.getNode(ISD::VP_MUL, SDLoc(N), VT, N1, N0, Mask, EVL);
+
+  auto ConstValue1 = isConstOrConstSplat(N1);
+
+  // fold vp_mul X, (splat -1), Mask, EVL -> vp_sub (splat 0), X, Mask, EVL
+  if (ConstValue1 && ConstValue1->isAllOnes()) {
+    auto ZeroSplat = DAG.getSplatVector(
+        VT, SDLoc(N),
+        DAG.getConstant(0, SDLoc(N), ConstValue1->getValueType(0)));
+    return DAG.getNode(ISD::VP_SUB, SDLoc(N), VT, ZeroSplat, N0, Mask, EVL);
+  }
+
+  // fold vp_mul X, (1 << C), Mask, EVL -> vp_shl X, C, Mask, EVL
+  if (ConstValue1 && ConstValue1->getAPIntValue().isPowerOf2()) {
+    auto Splat = DAG.getSplatVector(
+        VT, SDLoc(N),
+        DAG.getConstant(ConstValue1->getAPIntValue().countTrailingZeros(),
+                        SDLoc(N), ConstValue1->getValueType(0)));
+    return DAG.getNode(ISD::VP_SHL, SDLoc(N), VT, N0, Splat, Mask, EVL);
+  }
+
+  return SDValue();
+}
+
 SDValue DAGCombiner::visitVPOp(SDNode *N) {
   // VP operations in which all vector elements are disabled - either by
   // determining that the mask is all false or that the EVL is 0 - can be
   // eliminated.
   bool AreAllEltsDisabled = false;
   if (auto EVLIdx = ISD::getVPExplicitVectorLengthIdx(N->getOpcode()))
     AreAllEltsDisabled |= isNullConstant(N->getOperand(*EVLIdx));
   if (auto MaskIdx = ISD::getVPMaskIdx(N->getOpcode()))
     AreAllEltsDisabled |=
         ISD::isConstantSplatVectorAllZeros(N->getOperand(*MaskIdx).getNode());
 
   // This is the only generic VP combine we support for now.
-  if (!AreAllEltsDisabled)
+  if (!AreAllEltsDisabled) {
+    switch (N->getOpcode()) {
+    case ISD::VP_MUL:
+      return visitVP_MUL(N);
+    // TODO: Support more VP op.
+    default:
+      break;
+    }
     return SDValue();
+  }
 
   // Binary operations can be replaced by UNDEF.
   if (ISD::isVPBinaryOp(N->getOpcode()))
     return DAG.getUNDEF(N->getValueType(0));
 
   // VP Memory operations can be replaced by either the chain (stores) or the
   // chain + undef (loads).
   if (const auto *MemSD = dyn_cast<MemSDNode>(N)) {

llvm/test/CodeGen/RISCV/rvv/fixed-vectors-vmul-vp.ll

 ; RV64-NEXT:    ret
   %elt.head = insertelement <16 x i64> poison, i64 %b, i32 0
   %vb = shufflevector <16 x i64> %elt.head, <16 x i64> poison, <16 x i32> zeroinitializer
   %head = insertelement <16 x i1> poison, i1 true, i32 0
   %m = shufflevector <16 x i1> %head, <16 x i1> poison, <16 x i32> zeroinitializer
   %v = call <16 x i64> @llvm.vp.mul.v16i64(<16 x i64> %va, <16 x i64> %vb, <16 x i1> %m, i32 %evl)
   ret <16 x i64> %v
 }
+
+define <16 x i64> @vmul_vx_1_v16i64(<16 x i64> %va, <16 x i1> %m, i32 zeroext %evl) {
+; CHECK-LABEL: vmul_vx_1_v16i64:
+; CHECK:       # %bb.0:
+; CHECK-NEXT:    vsetvli zero, a0, e64, m8, ta, mu
+; CHECK-NEXT:    vrsub.vi v8, v8, 0, v0.t
+; CHECK-NEXT:    ret
+  %elt.head = insertelement <16 x i64> poison, i64 -1, i32 0
+  %vb = shufflevector <16 x i64> %elt.head, <16 x i64> poison, <16 x i32> zeroinitializer
+  %v = call <16 x i64> @llvm.vp.mul.v16i64(<16 x i64> %va, <16 x i64> %vb, <16 x i1> %m, i32 %evl)
+  ret <16 x i64> %v
+}
+
+define <16 x i64> @vmul_vx_1_v16i64_unmasked(<16 x i64> %va, i32 zeroext %evl) {
+; CHECK-LABEL: vmul_vx_1_v16i64_unmasked:
+; CHECK:       # %bb.0:
+; CHECK-NEXT:    vsetvli zero, a0, e64, m8, ta, mu
+; CHECK-NEXT:    vrsub.vi v8, v8, 0
+; CHECK-NEXT:    ret
+  %elt.head = insertelement <16 x i64> poison, i64 -1, i32 0
+  %vb = shufflevector <16 x i64> %elt.head, <16 x i64> poison, <16 x i32> zeroinitializer
+  %head = insertelement <16 x i1> poison, i1 true, i32 0
+  %m = shufflevector <16 x i1> %head, <16 x i1> poison, <16 x i32> zeroinitializer
+  %v = call <16 x i64> @llvm.vp.mul.v16i64(<16 x i64> %va, <16 x i64> %vb, <16 x i1> %m, i32 %evl)
+  ret <16 x i64> %v
+}
+
+define <16 x i64> @vmul_vx_2_v16i64(<16 x i64> %va, <16 x i1> %m, i32 zeroext %evl) {
+; CHECK-LABEL: vmul_vx_2_v16i64:
+; CHECK:       # %bb.0:
+; CHECK-NEXT:    vsetvli zero, a0, e64, m8, ta, mu
+; CHECK-NEXT:    vsll.vi v8, v8, 6, v0.t
+; CHECK-NEXT:    ret
+  %elt.head = insertelement <16 x i64> poison, i64 64, i32 0
+  %vb = shufflevector <16 x i64> %elt.head, <16 x i64> poison, <16 x i32> zeroinitializer
+  %v = call <16 x i64> @llvm.vp.mul.v16i64(<16 x i64> %va, <16 x i64> %vb, <16 x i1> %m, i32 %evl)
+  ret <16 x i64> %v
+}
+
+define <16 x i64> @vmul_vx_2_v16i64_unmasked(<16 x i64> %va, i32 zeroext %evl) {
+; CHECK-LABEL: vmul_vx_2_v16i64_unmasked:
+; CHECK:       # %bb.0:
+; CHECK-NEXT:    vsetvli zero, a0, e64, m8, ta, mu
+; CHECK-NEXT:    vsll.vi v8, v8, 6
+; CHECK-NEXT:    ret
+  %elt.head = insertelement <16 x i64> poison, i64 64, i32 0
+  %vb = shufflevector <16 x i64> %elt.head, <16 x i64> poison, <16 x i32> zeroinitializer
+  %head = insertelement <16 x i1> poison, i1 true, i32 0
+  %m = shufflevector <16 x i1> %head, <16 x i1> poison, <16 x i32> zeroinitializer
+  %v = call <16 x i64> @llvm.vp.mul.v16i64(<16 x i64> %va, <16 x i64> %vb, <16 x i1> %m, i32 %evl)
+  ret <16 x i64> %v
+}

llvm/test/CodeGen/RISCV/rvv/vmul-vp.ll

 ; RV64-NEXT:    ret
   %elt.head = insertelement <vscale x 8 x i64> poison, i64 %b, i32 0
   %vb = shufflevector <vscale x 8 x i64> %elt.head, <vscale x 8 x i64> poison, <vscale x 8 x i32> zeroinitializer
   %head = insertelement <vscale x 8 x i1> poison, i1 true, i32 0
   %m = shufflevector <vscale x 8 x i1> %head, <vscale x 8 x i1> poison, <vscale x 8 x i32> zeroinitializer
   %v = call <vscale x 8 x i64> @llvm.vp.mul.nxv8i64(<vscale x 8 x i64> %va, <vscale x 8 x i64> %vb, <vscale x 8 x i1> %m, i32 %evl)
   ret <vscale x 8 x i64> %v
 }
+
+define <vscale x 8 x i64> @vmul_vx_1_nxv8i64(<vscale x 8 x i64> %va, <vscale x 8 x i1> %m, i32 zeroext %evl) {
+; CHECK-LABEL: vmul_vx_1_nxv8i64:
+; CHECK:       # %bb.0:
+; CHECK-NEXT:    vsetvli zero, a0, e64, m8, ta, mu
+; CHECK-NEXT:    vrsub.vi v8, v8, 0, v0.t
+; CHECK-NEXT:    ret
+  %elt.head = insertelement <vscale x 8 x i64> poison, i64 -1, i32 0
+  %vb = shufflevector <vscale x 8 x i64> %elt.head, <vscale x 8 x i64> poison, <vscale x 8 x i32> zeroinitializer
+  %v = call <vscale x 8 x i64> @llvm.vp.mul.nxv8i64(<vscale x 8 x i64> %va, <vscale x 8 x i64> %vb, <vscale x 8 x i1> %m, i32 %evl)
+  ret <vscale x 8 x i64> %v
+}
+
+define <vscale x 8 x i64> @vmul_vx_1_nxv8i64_unmasked(<vscale x 8 x i64> %va, i32 zeroext %evl) {
+; CHECK-LABEL: vmul_vx_1_nxv8i64_unmasked:
+; CHECK:       # %bb.0:
+; CHECK-NEXT:    vsetvli zero, a0, e64, m8, ta, mu
+; CHECK-NEXT:    vrsub.vi v8, v8, 0
+; CHECK-NEXT:    ret
+  %elt.head = insertelement <vscale x 8 x i64> poison, i64 -1, i32 0
+  %vb = shufflevector <vscale x 8 x i64> %elt.head, <vscale x 8 x i64> poison, <vscale x 8 x i32> zeroinitializer
+  %head = insertelement <vscale x 8 x i1> poison, i1 true, i32 0
+  %m = shufflevector <vscale x 8 x i1> %head, <vscale x 8 x i1> poison, <vscale x 8 x i32> zeroinitializer
+  %v = call <vscale x 8 x i64> @llvm.vp.mul.nxv8i64(<vscale x 8 x i64> %va, <vscale x 8 x i64> %vb, <vscale x 8 x i1> %m, i32 %evl)
+  ret <vscale x 8 x i64> %v
+}
+
+define <vscale x 8 x i64> @vmul_vx_2_nxv8i64(<vscale x 8 x i64> %va, <vscale x 8 x i1> %m, i32 zeroext %evl) {
+; CHECK-LABEL: vmul_vx_2_nxv8i64:
+; CHECK:       # %bb.0:
+; CHECK-NEXT:    vsetvli zero, a0, e64, m8, ta, mu
+; CHECK-NEXT:    vsll.vi v8, v8, 6, v0.t
+; CHECK-NEXT:    ret
+  %elt.head = insertelement <vscale x 8 x i64> poison, i64 64, i32 0
+  %vb = shufflevector <vscale x 8 x i64> %elt.head, <vscale x 8 x i64> poison, <vscale x 8 x i32> zeroinitializer
+  %v = call <vscale x 8 x i64> @llvm.vp.mul.nxv8i64(<vscale x 8 x i64> %va, <vscale x 8 x i64> %vb, <vscale x 8 x i1> %m, i32 %evl)
+  ret <vscale x 8 x i64> %v
+}
+
+define <vscale x 8 x i64> @vmul_vx_2_nxv8i64_unmasked(<vscale x 8 x i64> %va, i32 zeroext %evl) {
+; CHECK-LABEL: vmul_vx_2_nxv8i64_unmasked:
+; CHECK:       # %bb.0:
+; CHECK-NEXT:    vsetvli zero, a0, e64, m8, ta, mu
+; CHECK-NEXT:    vsll.vi v8, v8, 6
+; CHECK-NEXT:    ret
+  %elt.head = insertelement <vscale x 8 x i64> poison, i64 64, i32 0
+  %vb = shufflevector <vscale x 8 x i64> %elt.head, <vscale x 8 x i64> poison, <vscale x 8 x i32> zeroinitializer
+  %head = insertelement <vscale x 8 x i1> poison, i1 true, i32 0
+  %m = shufflevector <vscale x 8 x i1> %head, <vscale x 8 x i1> poison, <vscale x 8 x i32> zeroinitializer
+  %v = call <vscale x 8 x i64> @llvm.vp.mul.nxv8i64(<vscale x 8 x i64> %va, <vscale x 8 x i64> %vb, <vscale x 8 x i1> %m, i32 %evl)
+  ret <vscale x 8 x i64> %v
+}