[llvm] [LV] Fix crash in replaceWithFinalIfReductionStore for constant-folded reductions (PR #200983)

[via llvm-commits]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-llvm-transforms

Author: 陈子昂 (Michael-Chen-NJU)

<details>
<summary>Changes</summary>

When a reduction's intermediate value is simplified to a constant (e.g. or x, -1 -> -1), the VPReductionPHIRecipe is removed as dead code. However, replaceWithFinalIfReductionStore still attempts to find the reduction phi via findUserOf, which returns null, causing a segfault.

Handle this case by checking for null before dereferencing.

Fixes #<!-- -->200742.

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


2 Files Affected:

- (modified) llvm/lib/Transforms/Vectorize/LoopVectorize.cpp (+8-2) 
- (added) llvm/test/Transforms/LoopVectorize/reduction-invariant-store-const-folded.ll (+24) 


``````````diff
diff --git a/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp b/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp
index abb19324b76eb..57bdedd38acbc 100644
--- a/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp
+++ b/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp
@@ -6392,8 +6392,14 @@ bool VPRecipeBuilder::replaceWithFinalIfReductionStore(
       // if tail folded.
       if (auto *Blend = VPlanPatternMatch::findUserOf<VPBlendRecipe>(Val))
         Val = Blend;
-      assert(VPlanPatternMatch::findUserOf<VPReductionPHIRecipe>(Val)
-                     ->getBackedgeValue() == Val &&
+      // The reduction may have been folded to a constant, so the
+      // VPReductionPHIRecipe may no longer exist.
+      auto *RedPhi = VPlanPatternMatch::findUserOf<VPReductionPHIRecipe>(Val);
+      if (!RedPhi) {
+        VPI->eraseFromParent();
+        return true;
+      }
+      assert(RedPhi->getBackedgeValue() == Val &&
              "Store isn't backedge value?");
       auto *Recipe = new VPReplicateRecipe(
           SI, {Val, Addr}, true /* IsUniform */, nullptr /*Mask*/, *VPI, *VPI,
diff --git a/llvm/test/Transforms/LoopVectorize/reduction-invariant-store-const-folded.ll b/llvm/test/Transforms/LoopVectorize/reduction-invariant-store-const-folded.ll
new file mode 100644
index 0000000000000..6e2775eb82917
--- /dev/null
+++ b/llvm/test/Transforms/LoopVectorize/reduction-invariant-store-const-folded.ll
@@ -0,0 +1,24 @@
+; RUN: opt < %s -passes="loop-vectorize" -S | FileCheck %s
+;
+; Verify that loop-vectorize does not crash when a reduction with an invariant
+; store is simplified to a constant (or x, -1 -> -1).
+
+define void @reduction_store_const_folded() {
+; CHECK-LABEL: define void @reduction_store_const_folded()
+; CHECK-NOT: vector.body
+; CHECK: ret void
+entry:
+  br label %loop
+
+loop:
+  %j = phi i32 [ 0, %entry ], [ %add, %loop ]
+  %or810 = phi i32 [ 0, %entry ], [ %or, %loop ]
+  %or = or i32 %or810, -1
+  store i32 %or, ptr null, align 4
+  %add = add i32 %j, 1
+  %cmp = icmp slt i32 %j, 1
+  br i1 %cmp, label %loop, label %exit
+
+exit:
+  ret void
+}

``````````

</details>


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