[llvm] [VectorCombine] Fix duplicate freeze scheduling for shared indexes (PR #213461)

[via llvm-commits]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-llvm-transforms

@llvm/pr-subscribers-vectorizers

Author: Aditya prakash jha (AdityaOP007)

<details>
<summary>Changes</summary>

## Summary
#<!-- -->213423 

This PR fixes a crash in `VectorCombine::scalarizeLoadExtract()` when multiple `extractelement` instructions share the same index instruction.

The issue was caused by `NeedFreeze` tracking freeze requests per `ExtractElementInst`, which could schedule duplicate `freeze()` operations for the same index. After the first freeze updated the shared index, processing the second request used stale state and triggered the `is_contained(ToFreeze->users(), &UserI)` assertion.

The fix tracks freeze requests using the shared index instruction instead, ensuring each index is frozen only once while preserving the existing assertion and use-def invariants.

A regression test has also been added to cover the shared-index case and prevent future regressions.


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


4 Files Affected:

- (modified) llvm/lib/Transforms/Vectorize/VectorCombine.cpp (+9-6) 
- (added) llvm/test/Transforms/VectorCombine/X86/pr213423.ll (+33) 
- (added) test.ll (+5) 
- (added) test2.ll (+7) 


``````````diff
diff --git a/llvm/lib/Transforms/Vectorize/VectorCombine.cpp b/llvm/lib/Transforms/Vectorize/VectorCombine.cpp
index 726f564b1aad9..2c1143b0957a6 100644
--- a/llvm/lib/Transforms/Vectorize/VectorCombine.cpp
+++ b/llvm/lib/Transforms/Vectorize/VectorCombine.cpp
@@ -1929,7 +1929,7 @@ static ScalarizationResult canScalarizeAccess(VectorType *VecTy, Value *Idx,
 
   // If the index may be poison, check if we can insert a freeze before the
   // range of the index is restricted.
-  Value *IdxBase;
+  Value *IdxBase = nullptr;
   ConstantInt *CI;
   if (match(Idx, m_And(m_Value(IdxBase), m_ConstantInt(CI)))) {
     IdxRange = IdxRange.binaryAnd(CI->getValue());
@@ -1937,7 +1937,7 @@ static ScalarizationResult canScalarizeAccess(VectorType *VecTy, Value *Idx,
     IdxRange = IdxRange.urem(CI->getValue());
   }
 
-  if (ValidIndices.contains(IdxRange))
+  if (IdxBase && ValidIndices.contains(IdxRange))
     return ScalarizationResult::safeWithFreeze(IdxBase);
   return ScalarizationResult::unsafe();
 }
@@ -2086,7 +2086,7 @@ bool VectorCombine::scalarizeLoadExtract(LoadInst *LI, VectorType *VecTy,
   if (!TTI.allowVectorElementIndexingUsingGEP())
     return false;
 
-  DenseMap<ExtractElementInst *, ScalarizationResult> NeedFreeze;
+  DenseMap<Instruction *, ScalarizationResult> NeedFreeze;
   llvm::scope_exit FailureGuard([&]() {
     // If the transform is aborted, discard the ScalarizationResults.
     for (auto &Pair : NeedFreeze)
@@ -2106,7 +2106,8 @@ bool VectorCombine::scalarizeLoadExtract(LoadInst *LI, VectorType *VecTy,
     if (ScalarIdx.isUnsafe())
       return false;
     if (ScalarIdx.isSafeWithFreeze()) {
-      NeedFreeze.try_emplace(UI, ScalarIdx);
+      NeedFreeze.try_emplace(cast<Instruction>(UI->getIndexOperand()),
+                             ScalarIdx);
       ScalarIdx.discard();
     }
 
@@ -2140,9 +2141,11 @@ bool VectorCombine::scalarizeLoadExtract(LoadInst *LI, VectorType *VecTy,
     Value *Idx = EI->getIndexOperand();
 
     // Insert 'freeze' for poison indexes.
-    auto It = NeedFreeze.find(EI);
-    if (It != NeedFreeze.end())
+    auto It = NeedFreeze.find(cast<Instruction>(Idx));
+    if (It != NeedFreeze.end()) {
       It->second.freeze(Builder, *cast<Instruction>(Idx));
+      NeedFreeze.erase(It);
+    }
 
     Builder.SetInsertPoint(EI);
     Value *GEP =
diff --git a/llvm/test/Transforms/VectorCombine/X86/pr213423.ll b/llvm/test/Transforms/VectorCombine/X86/pr213423.ll
new file mode 100644
index 0000000000000..2b932bca422d6
--- /dev/null
+++ b/llvm/test/Transforms/VectorCombine/X86/pr213423.ll
@@ -0,0 +1,33 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
+; RUN: opt < %s -passes=vector-combine -S | FileCheck %s
+
+; Issue #213423
+; Check that we do not crash when multiple extracts share the same index instruction
+; and it needs to be frozen.
+define void @test_shared_index(ptr %p, i32 %idx.base) {
+; CHECK-LABEL: @test_shared_index(
+; CHECK-NEXT:    [[IDX:%.*]] = and i32 [[IDX_BASE:%.*]], 1
+; CHECK-NEXT:    [[V:%.*]] = load <2 x i32>, ptr [[P:%.*]], align 8
+; CHECK-NEXT:    [[E1:%.*]] = extractelement <2 x i32> [[V]], i32 [[IDX]]
+; CHECK-NEXT:    [[E2:%.*]] = extractelement <2 x i32> [[V]], i32 [[IDX]]
+; CHECK-NEXT:    ret void
+;
+  %idx = and i32 %idx.base, 1
+  %v = load <2 x i32>, ptr %p
+  %e1 = extractelement <2 x i32> %v, i32 %idx
+  %e2 = extractelement <2 x i32> %v, i32 %idx
+  ret void
+}
+
+; Check that we do not use an uninitialized pointer when the index doesn't match and/urem
+; but is naturally in bounds.
+define void @test_uninitialized_base(ptr %p, i1 %idx) {
+; CHECK-LABEL: @test_uninitialized_base(
+; CHECK-NEXT:    [[V:%.*]] = load <2 x i32>, ptr [[P:%.*]], align 8
+; CHECK-NEXT:    [[E:%.*]] = extractelement <2 x i32> [[V]], i1 [[IDX:%.*]]
+; CHECK-NEXT:    ret void
+;
+  %v = load <2 x i32>, ptr %p
+  %e = extractelement <2 x i32> %v, i1 %idx
+  ret void
+}
diff --git a/test.ll b/test.ll
new file mode 100644
index 0000000000000..cc7d1e1111b31
--- /dev/null
+++ b/test.ll
@@ -0,0 +1,5 @@
+define void @test(ptr %p, i1 %idx) {
+  %v = load <2 x i32>, ptr %p
+  %e = extractelement <2 x i32> %v, i1 %idx
+  ret void
+}
diff --git a/test2.ll b/test2.ll
new file mode 100644
index 0000000000000..3b5e2b0a42912
--- /dev/null
+++ b/test2.ll
@@ -0,0 +1,7 @@
+define void @test(ptr %p, i32 %idx.base) {
+  %idx = and i32 %idx.base, 1
+  %v = load <2 x i32>, ptr %p
+  %e1 = extractelement <2 x i32> %v, i32 %idx
+  %e2 = extractelement <2 x i32> %v, i32 %idx
+  ret void
+}

``````````

</details>


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