[NVPTX] support VECREDUCE_SEQ ops and remove option

Author: Prince781

llvm/lib/Target/NVPTX/NVPTXISelLowering.cpp

@@ -876,6 +876,7 @@ NVPTXTargetLowering::NVPTXTargetLowering(const NVPTXTargetMachine &TM,
     if (EltVT == MVT::f16 || EltVT == MVT::bf16 || EltVT == MVT::f32 ||
         EltVT == MVT::f64) {
       setOperationAction({ISD::VECREDUCE_FADD, ISD::VECREDUCE_FMUL,
+                          ISD::VECREDUCE_SEQ_FADD, ISD::VECREDUCE_SEQ_FMUL,
                           ISD::VECREDUCE_FMAX, ISD::VECREDUCE_FMIN,
                           ISD::VECREDUCE_FMAXIMUM, ISD::VECREDUCE_FMINIMUM},
                          VT, Custom);
@@ -2270,12 +2271,19 @@ static SDValue BuildTreeReduction(
 /// max3/min3 when the target supports them.
 SDValue NVPTXTargetLowering::LowerVECREDUCE(SDValue Op,
                                             SelectionDAG &DAG) const {
-  if (DisableFOpTreeReduce)
-    return SDValue();
-
   SDLoc DL(Op);
   const SDNodeFlags Flags = Op->getFlags();
-  const SDValue &Vector = Op.getOperand(0);
+  SDValue Vector;
+  SDValue Accumulator;
+  if (Op->getOpcode() == ISD::VECREDUCE_SEQ_FADD ||
+      Op->getOpcode() == ISD::VECREDUCE_SEQ_FMUL) {
+    // special case with accumulator as first arg
+    Accumulator = Op.getOperand(0);
+    Vector = Op.getOperand(1);
+  } else {
+    // default case
+    Vector = Op.getOperand(0);
+  }
   EVT EltTy = Vector.getValueType().getVectorElementType();
   const bool CanUseMinMax3 = EltTy == MVT::f32 && STI.getSmVersion() >= 100 &&
                              STI.getPTXVersion() >= 88;
@@ -2287,10 +2295,12 @@ SDValue NVPTXTargetLowering::LowerVECREDUCE(SDValue Op,
 
   switch (Op->getOpcode()) {
   case ISD::VECREDUCE_FADD:
+  case ISD::VECREDUCE_SEQ_FADD:
     ScalarOps = {{ISD::FADD, 2}};
     IsReassociatable = false;
     break;
   case ISD::VECREDUCE_FMUL:
+  case ISD::VECREDUCE_SEQ_FMUL:
     ScalarOps = {{ISD::FMUL, 2}};
     IsReassociatable = false;
     break;
@@ -2369,11 +2379,13 @@ SDValue NVPTXTargetLowering::LowerVECREDUCE(SDValue Op,
   }
 
   // Lower to tree reduction.
-  if (IsReassociatable || Flags.hasAllowReassociation())
+  if (IsReassociatable || Flags.hasAllowReassociation()) {
+    // we don't expect an accumulator for reassociatable vector reduction ops
+    assert(!Accumulator && "unexpected accumulator");
     return BuildTreeReduction(Elements, EltTy, ScalarOps, DL, Flags, DAG);
+  }
 
   // Lower to sequential reduction.
-  SDValue Accumulator;
   for (unsigned OpIdx = 0, I = 0; I < NumElts; ++OpIdx) {
     assert(OpIdx < ScalarOps.size() && "no smaller operators for reduction");
     const auto [DefaultScalarOp, DefaultGroupSize] = ScalarOps[OpIdx];
@@ -3234,6 +3246,8 @@ NVPTXTargetLowering::LowerOperation(SDValue Op, SelectionDAG &DAG) const {
     return LowerCONCAT_VECTORS(Op, DAG);
   case ISD::VECREDUCE_FADD:
   case ISD::VECREDUCE_FMUL:
+  case ISD::VECREDUCE_SEQ_FADD:
+  case ISD::VECREDUCE_SEQ_FMUL:
   case ISD::VECREDUCE_FMAX:
   case ISD::VECREDUCE_FMIN:
   case ISD::VECREDUCE_FMAXIMUM:

llvm/test/CodeGen/NVPTX/reduction-intrinsics.ll

@@ -23,19 +23,19 @@ define half @reduce_fadd_half(<8 x half> %in) {
 ; CHECK-EMPTY:
 ; CHECK-NEXT:  // %bb.0:
 ; CHECK-NEXT:    ld.param.v4.b32 {%r1, %r2, %r3, %r4}, [reduce_fadd_half_param_0];
-; CHECK-NEXT:    mov.b32 {%rs1, %rs2}, %r4;
-; CHECK-NEXT:    mov.b32 {%rs3, %rs4}, %r3;
-; CHECK-NEXT:    mov.b32 {%rs5, %rs6}, %r2;
-; CHECK-NEXT:    mov.b32 {%rs7, %rs8}, %r1;
-; CHECK-NEXT:    mov.b16 %rs9, 0x0000;
-; CHECK-NEXT:    add.rn.f16 %rs10, %rs7, %rs9;
-; CHECK-NEXT:    add.rn.f16 %rs11, %rs10, %rs8;
-; CHECK-NEXT:    add.rn.f16 %rs12, %rs11, %rs5;
-; CHECK-NEXT:    add.rn.f16 %rs13, %rs12, %rs6;
-; CHECK-NEXT:    add.rn.f16 %rs14, %rs13, %rs3;
-; CHECK-NEXT:    add.rn.f16 %rs15, %rs14, %rs4;
-; CHECK-NEXT:    add.rn.f16 %rs16, %rs15, %rs1;
-; CHECK-NEXT:    add.rn.f16 %rs17, %rs16, %rs2;
+; CHECK-NEXT:    mov.b32 {%rs1, %rs2}, %r1;
+; CHECK-NEXT:    mov.b16 %rs3, 0x0000;
+; CHECK-NEXT:    add.rn.f16 %rs4, %rs1, %rs3;
+; CHECK-NEXT:    add.rn.f16 %rs5, %rs4, %rs2;
+; CHECK-NEXT:    mov.b32 {%rs6, %rs7}, %r2;
+; CHECK-NEXT:    add.rn.f16 %rs8, %rs5, %rs6;
+; CHECK-NEXT:    add.rn.f16 %rs9, %rs8, %rs7;
+; CHECK-NEXT:    mov.b32 {%rs10, %rs11}, %r3;
+; CHECK-NEXT:    add.rn.f16 %rs12, %rs9, %rs10;
+; CHECK-NEXT:    add.rn.f16 %rs13, %rs12, %rs11;
+; CHECK-NEXT:    mov.b32 {%rs14, %rs15}, %r4;
+; CHECK-NEXT:    add.rn.f16 %rs16, %rs13, %rs14;
+; CHECK-NEXT:    add.rn.f16 %rs17, %rs16, %rs15;
 ; CHECK-NEXT:    st.param.b16 [func_retval0], %rs17;
 ; CHECK-NEXT:    ret;
   %res = call half @llvm.vector.reduce.fadd(half 0.0, <8 x half> %in)
@@ -174,17 +174,17 @@ define half @reduce_fmul_half(<8 x half> %in) {
 ; CHECK-EMPTY:
 ; CHECK-NEXT:  // %bb.0:
 ; CHECK-NEXT:    ld.param.v4.b32 {%r1, %r2, %r3, %r4}, [reduce_fmul_half_param_0];
-; CHECK-NEXT:    mov.b32 {%rs1, %rs2}, %r4;
-; CHECK-NEXT:    mov.b32 {%rs3, %rs4}, %r3;
-; CHECK-NEXT:    mov.b32 {%rs5, %rs6}, %r2;
-; CHECK-NEXT:    mov.b32 {%rs7, %rs8}, %r1;
-; CHECK-NEXT:    mul.rn.f16 %rs9, %rs7, %rs8;
-; CHECK-NEXT:    mul.rn.f16 %rs10, %rs9, %rs5;
-; CHECK-NEXT:    mul.rn.f16 %rs11, %rs10, %rs6;
-; CHECK-NEXT:    mul.rn.f16 %rs12, %rs11, %rs3;
-; CHECK-NEXT:    mul.rn.f16 %rs13, %rs12, %rs4;
-; CHECK-NEXT:    mul.rn.f16 %rs14, %rs13, %rs1;
-; CHECK-NEXT:    mul.rn.f16 %rs15, %rs14, %rs2;
+; CHECK-NEXT:    mov.b32 {%rs1, %rs2}, %r2;
+; CHECK-NEXT:    mov.b32 {%rs3, %rs4}, %r1;
+; CHECK-NEXT:    mul.rn.f16 %rs5, %rs3, %rs4;
+; CHECK-NEXT:    mul.rn.f16 %rs6, %rs5, %rs1;
+; CHECK-NEXT:    mul.rn.f16 %rs7, %rs6, %rs2;
+; CHECK-NEXT:    mov.b32 {%rs8, %rs9}, %r3;
+; CHECK-NEXT:    mul.rn.f16 %rs10, %rs7, %rs8;
+; CHECK-NEXT:    mul.rn.f16 %rs11, %rs10, %rs9;
+; CHECK-NEXT:    mov.b32 {%rs12, %rs13}, %r4;
+; CHECK-NEXT:    mul.rn.f16 %rs14, %rs11, %rs12;
+; CHECK-NEXT:    mul.rn.f16 %rs15, %rs14, %rs13;
 ; CHECK-NEXT:    st.param.b16 [func_retval0], %rs15;
 ; CHECK-NEXT:    ret;
   %res = call half @llvm.vector.reduce.fmul(half 1.0, <8 x half> %in)