[llvm] [LoopVectorize] Fix crash with GC statepoint in epilogue vectorization (PR #179459)

[via llvm-commits]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-vectorizers

Author: Luohao Wang (Luohaothu)

<details>
<summary>Changes</summary>

Fixes #<!-- -->179407.

When vectorizing loops in functions with GC statepoint, `fixScalarResumeValuesFromBypass` assumed that `IVPhi->getIncomingValueForBlock(PH)` always returns a PHINode. However, in functions with GC statepoint, the preheader may be simplified and the incoming value could be a constant or other non-PHINode value, causing an assertion failure:
  "isa<> used on a null pointer"

The fix uses dyn_cast instead of cast to safely check the type, and when the incoming value is not a PHINode, creates a new PHI node in the preheader to merge the original value and the bypass value.

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


2 Files Affected:

- (modified) llvm/lib/Transforms/Vectorize/LoopVectorize.cpp (+21-3) 
- (added) llvm/test/Transforms/LoopVectorize/pr179407-gc-statepoint.ll (+40) 


``````````diff
diff --git a/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp b/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp
index abac45b265d10..2581353f900bf 100644
--- a/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp
+++ b/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp
@@ -9361,12 +9361,30 @@ static void fixScalarResumeValuesFromBypass(BasicBlock *BypassBlock, Loop *L,
   // Fix induction resume values from the additional bypass block.
   IRBuilder<> BypassBuilder(BypassBlock, BypassBlock->getFirstInsertionPt());
   for (const auto &[IVPhi, II] : LVL.getInductionVars()) {
-    auto *Inc = cast<PHINode>(IVPhi->getIncomingValueForBlock(PH));
+    Value *IncomingFromPH = IVPhi->getIncomingValueForBlock(PH);
     Value *V = createInductionAdditionalBypassValues(
         IVPhi, II, BypassBuilder, ExpandedSCEVs, MainVectorTripCount,
         LVL.getPrimaryInduction());
-    // TODO: Directly add as extra operand to the VPResumePHI recipe.
-    Inc->setIncomingValueForBlock(BypassBlock, V);
+
+    if (auto *Inc = dyn_cast<PHINode>(IncomingFromPH)) {
+      // TODO: Directly add as extra operand to the VPResumePHI recipe.
+      Inc->setIncomingValueForBlock(BypassBlock, V);
+    } else {
+      // If the incoming value from preheader is not a PHI node (e.g., a constant
+      // in functions with GC statepoint), we need to create a PHI node in the
+      // preheader that merges the original value and the bypass value, then update
+      // IVPhi to use this new PHI node.
+      PHINode *NewPhi = PHINode::Create(IncomingFromPH->getType(), 2,
+                                        IVPhi->getName() + ".phi.merge", PH->getFirstNonPHIIt());
+      for (BasicBlock *Pred : predecessors(PH)) {
+        if (Pred == BypassBlock)
+          NewPhi->addIncoming(V, BypassBlock);
+        else
+          NewPhi->addIncoming(IncomingFromPH, Pred);
+      }
+      // Update IVPhi to use the new PHI node for the incoming value from PH.
+      IVPhi->setIncomingValueForBlock(PH, NewPhi);
+    }
   }
 }
 
diff --git a/llvm/test/Transforms/LoopVectorize/pr179407-gc-statepoint.ll b/llvm/test/Transforms/LoopVectorize/pr179407-gc-statepoint.ll
new file mode 100644
index 0000000000000..a2c33b3f1a702
--- /dev/null
+++ b/llvm/test/Transforms/LoopVectorize/pr179407-gc-statepoint.ll
@@ -0,0 +1,40 @@
+; RUN: opt -passes=loop-vectorize -enable-epilogue-vectorization -force-vector-width=2 -epilogue-vectorization-force-VF=2 -S < %s | FileCheck %s
+;
+; Test case for issue #179407: LoopVectorize should not crash when
+; vectorizing loops in functions with GC statepoint.
+;
+; The issue was that fixScalarResumeValuesFromBypass assumed
+; IVPhi->getIncomingValueForBlock(PH) always returns a PHINode, but in
+; functions with GC statepoint, it could return a constant or other value.
+;
+; This test verifies that epilogue vectorization works correctly with
+; GC statepoint by checking that the loop is vectorized and the
+; bypass blocks are properly created.
+
+; CHECK: define void @wombat
+; CHECK-SAME: gc "statepoint-example"
+; CHECK: vector.body:
+; CHECK: vec.epilog.vector.body:
+; CHECK: vec.epilog.scalar.ph:
+; CHECK: bb1:
+
+define void @wombat(i64 %arg) gc "statepoint-example" {
+bb:
+  br label %bb1
+
+bb1:                                              ; preds = %bb1, %bb
+  %phi = phi i64 [ 0, %bb ], [ %add, %bb1 ]
+  %phi2 = phi i32 [ 0, %bb ], [ %or, %bb1 ]
+  %phi3 = phi i32 [ 0, %bb ], [ %select, %bb1 ]
+  %icmp = icmp eq i32 0, 0
+  %select = select i1 %icmp, i32 0, i32 %phi3
+  %or = or i32 %phi2, 0
+  %add = add i64 %phi, 1
+  %icmp4 = icmp ult i64 %phi, %arg
+  br i1 %icmp4, label %bb1, label %bb5
+
+bb5:                                              ; preds = %bb1
+  %phi6 = phi i32 [ %select, %bb1 ]
+  %phi7 = phi i32 [ %or, %bb1 ]
+  ret void
+}

``````````

</details>


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