[llvm] [SelectionDAG] Fold VECREDUCE_ADD of a constant BUILD_VECTOR (PR #207560)

[via llvm-commits]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-llvm-selectiondag

Author: Ankit Kumar Tiwari (ankit-cybertron)

<details>
<summary>Changes</summary>


This PR implements constant folding for ISD::VECREDUCE_ADD when the input is a BUILD_VECTOR of integer constants. The fold computes the sum at compile time using APInt arithmetic and returns a folded constant scalar, instead of emitting real vector materialization and reduction instructions for a value that's already known.

Folding is skipped if any element is undef or opaque, since the result can't be assumed at compile time in that case.

#<!-- -->206743

---
Full diff: https://github.com/llvm/llvm-project/pull/207560.diff


3 Files Affected:

- (modified) llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp (+18) 
- (modified) llvm/test/CodeGen/RISCV/rvv/fold-binary-reduce.ll (+1-7) 
- (added) llvm/test/CodeGen/RISCV/rvv/vecreduce-add-constant-fold.ll (+53) 


``````````diff
diff --git a/llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp b/llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp
index c542f8e7cc20b..a6a5c29e1b716 100644
--- a/llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp
+++ b/llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp
@@ -7114,6 +7114,7 @@ SDValue SelectionDAG::getNode(unsigned Opcode, const SDLoc &DL, EVT VT,
   case ISD::CTTZ_ZERO_POISON:
   case ISD::CTPOP:
   case ISD::CTLS:
+  case ISD::VECREDUCE_ADD:
   case ISD::STEP_VECTOR: {
     SDValue Ops = {N1};
     if (SDValue Fold = FoldConstantArithmetic(Opcode, DL, VT, Ops))
@@ -7785,6 +7786,23 @@ SDValue SelectionDAG::FoldConstantArithmetic(unsigned Opcode, const SDLoc &DL,
     // Early-out if we failed to constant fold a bitcast.
     if (Opcode == ISD::BITCAST)
       return SDValue();
+
+    // Constant fold VECREDUCE_ADD with a BUILD_VECTOR of integer constants.
+    if (Opcode == ISD::VECREDUCE_ADD && N1.getOpcode() == ISD::BUILD_VECTOR &&
+        ISD::isBuildVectorOfConstantSDNodes(N1.getNode())) {
+      unsigned EltBits = N1.getValueType().getScalarSizeInBits();
+      APInt Acc = APInt::getZero(EltBits);
+      bool HasUndef = false;
+      for (SDValue Elt : N1->op_values()) {
+        if (Elt.isUndef() || cast<ConstantSDNode>(Elt)->isOpaque()) {
+          HasUndef = true;
+          break;
+        }
+        Acc += cast<ConstantSDNode>(Elt)->getAPIntValue().trunc(EltBits);
+      }
+      if (!HasUndef)
+        return getConstant(Acc.zextOrTrunc(VT.getSizeInBits()), DL, VT);
+    }
   }
 
   // Handle binops special cases.
diff --git a/llvm/test/CodeGen/RISCV/rvv/fold-binary-reduce.ll b/llvm/test/CodeGen/RISCV/rvv/fold-binary-reduce.ll
index 134bea04b58fc..5a1ce80eb9468 100644
--- a/llvm/test/CodeGen/RISCV/rvv/fold-binary-reduce.ll
+++ b/llvm/test/CodeGen/RISCV/rvv/fold-binary-reduce.ll
@@ -327,13 +327,7 @@ entry:
 define void @crash(<2 x i32> %0) {
 ; CHECK-LABEL: crash:
 ; CHECK:       # %bb.0: # %entry
-; CHECK-NEXT:    vsetivli zero, 4, e16, mf2, ta, ma
-; CHECK-NEXT:    vmv.v.i v9, 0
-; CHECK-NEXT:    vsetvli zero, zero, e32, m1, ta, ma
-; CHECK-NEXT:    vmv.x.s a0, v8
-; CHECK-NEXT:    vmv.s.x v8, a0
-; CHECK-NEXT:    vsetvli zero, zero, e16, mf2, ta, ma
-; CHECK-NEXT:    vredsum.vs v8, v9, v8
+; CHECK-NEXT:    vsetivli zero, 1, e32, m1, ta, ma
 ; CHECK-NEXT:    vmv.x.s a0, v8
 ; CHECK-NEXT:    sb a0, 0(zero)
 ; CHECK-NEXT:    ret
diff --git a/llvm/test/CodeGen/RISCV/rvv/vecreduce-add-constant-fold.ll b/llvm/test/CodeGen/RISCV/rvv/vecreduce-add-constant-fold.ll
new file mode 100644
index 0000000000000..c45ead8152995
--- /dev/null
+++ b/llvm/test/CodeGen/RISCV/rvv/vecreduce-add-constant-fold.ll
@@ -0,0 +1,53 @@
+; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py
+; RUN: llc -mtriple=riscv64 -mattr=+v -verify-machineinstrs < %s | FileCheck %s
+
+; Case 1: all-constant vector -> should fold to a plain constant (10)
+define i32 @test_const() {
+; CHECK-LABEL: test_const:
+; CHECK:       # %bb.0:
+; CHECK-NEXT:    li a0, 10
+; CHECK-NEXT:    ret
+  %r = call i32 @llvm.vector.reduce.add.v4i32(<4 x i32> <i32 1, i32 2, i32 3, i32 4>)
+  ret i32 %r
+}
+
+; Case 2: vector containing an undef element -> should NOT fold
+define i32 @test_undef() {
+; CHECK-LABEL: test_undef:
+; CHECK:       # %bb.0:
+; CHECK-NEXT:    vsetivli zero, 4, e32, m1, ta, ma
+; CHECK-NEXT:    vmv.s.x v8, zero
+; CHECK-NEXT:    vid.v v9
+; CHECK-NEXT:    vadd.vi v9, v9, 1
+; CHECK-NEXT:    vredsum.vs v8, v9, v8
+; CHECK-NEXT:    vmv.x.s a0, v8
+; CHECK-NEXT:    ret
+  %r = call i32 @llvm.vector.reduce.add.v4i32(<4 x i32> <i32 1, i32 2, i32 undef, i32 4>)
+  ret i32 %r
+}
+
+; Case 3: 8-element vector, sum = 360
+define i32 @test_const_wide() {
+; CHECK-LABEL: test_const_wide:
+; CHECK:       # %bb.0:
+; CHECK-NEXT:    li a0, 360
+; CHECK-NEXT:    ret
+  %r = call i32 @llvm.vector.reduce.add.v8i32(<8 x i32> <i32 10, i32 20, i32 30, i32 40, i32 50, i32 60, i32 70, i32 80>)
+  ret i32 %r
+}
+
+; Case 4: non-constant input -> should NOT fold
+define i32 @test_nonconst(<4 x i32> %v) {
+; CHECK-LABEL: test_nonconst:
+; CHECK:       # %bb.0:
+; CHECK-NEXT:    vsetivli zero, 4, e32, m1, ta, ma
+; CHECK-NEXT:    vmv.s.x v9, zero
+; CHECK-NEXT:    vredsum.vs v8, v8, v9
+; CHECK-NEXT:    vmv.x.s a0, v8
+; CHECK-NEXT:    ret
+  %r = call i32 @llvm.vector.reduce.add.v4i32(<4 x i32> %v)
+  ret i32 %r
+}
+
+declare i32 @llvm.vector.reduce.add.v4i32(<4 x i32>)
+declare i32 @llvm.vector.reduce.add.v8i32(<8 x i32>)

``````````

</details>


https://github.com/llvm/llvm-project/pull/207560