[llvm] [VectorCombine] Fix bailing out of the foldSingleElementStore with a freeze-pending index (PR #207702)

[via llvm-commits]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-llvm-transforms

Author: Kacper Doga (varev-dev)

<details>
<summary>Changes</summary>

`foldSingleElementStore` calls `canScalarizeAccess` to decide whether a single-element store can be scalarized. When the insert index is not a constant, is provably in-bounds and may be poison, `canScalarizeAccess` returns `SafeWithFreeze` and stashes a `ToFreeze` value that must later be consumed by `freeze()` or cleared by `discard()`.

The bail-out for memory being modified between the load and store was folded into same condition as the scalarizability check:

```cpp
auto ScalarizableIdx = canScalarizeAccess(...);
if (ScalarizableIdx.isUnsafe() || isMemModifiedBetween(...))
  return false;
```

When the index took the `SafeWithFreeze` path and memory was modified between the load and the store, `isMemModifiedBetween` returns true and hit `return false` while ScalarizableIdx still holds a live `ToFreeze`.

Its destructor then raises:
> Assertion !ToFreeze && "freeze() not called with ToFreeze being set"

### Fix
Move the `isMemModifiedBetween` bail-out above the `canScalarizeAccess` call so the freeze-pending result is never constructed on the path where we give up.
### Coverage
- Added `@<!-- -->insert_store_mem_modify_nonconst_index` to `llvm/test/Transforms/VectorCombine/load-insert-store.ll`

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


2 Files Affected:

- (modified) llvm/lib/Transforms/Vectorize/VectorCombine.cpp (+4-3) 
- (modified) llvm/test/Transforms/VectorCombine/load-insert-store.ll (+20) 


``````````diff
diff --git a/llvm/lib/Transforms/Vectorize/VectorCombine.cpp b/llvm/lib/Transforms/Vectorize/VectorCombine.cpp
index 009ecf89f3d04..91ed59a605a4c 100644
--- a/llvm/lib/Transforms/Vectorize/VectorCombine.cpp
+++ b/llvm/lib/Transforms/Vectorize/VectorCombine.cpp
@@ -1989,11 +1989,12 @@ bool VectorCombine::foldSingleElementStore(Instruction &I) {
         SrcAddr != SI->getPointerOperand()->stripPointerCasts())
       return false;
 
+    if (isMemModifiedBetween(Load->getIterator(), SI->getIterator(),
+                             MemoryLocation::get(SI), AA))
+      return false;
     auto ScalarizableIdx =
         canScalarizeAccess(VecTy, Idx, SQ.getWithInstruction(Load));
-    if (ScalarizableIdx.isUnsafe() ||
-        isMemModifiedBetween(Load->getIterator(), SI->getIterator(),
-                             MemoryLocation::get(SI), AA))
+    if (ScalarizableIdx.isUnsafe())
       return false;
 
     // Ensure we add the load back to the worklist BEFORE its users so they can
diff --git a/llvm/test/Transforms/VectorCombine/load-insert-store.ll b/llvm/test/Transforms/VectorCombine/load-insert-store.ll
index 93565c1a708eb..897b113c197cf 100644
--- a/llvm/test/Transforms/VectorCombine/load-insert-store.ll
+++ b/llvm/test/Transforms/VectorCombine/load-insert-store.ll
@@ -713,6 +713,26 @@ entry:
   ret void
 }
 
+; Non-constant index makes the access "safe with freeze", but memory is
+; modified between the load and store, so the fold must bail out without
+; leaving the pending ToFreeze value set.
+define void @insert_store_mem_modify_nonconst_index(ptr %p, i8 %s) {
+; CHECK-LABEL: @insert_store_mem_modify_nonconst_index(
+; CHECK-NEXT:    [[LD:%.*]] = load <2 x i8>, ptr [[P:%.*]], align 2
+; CHECK-NEXT:    [[A:%.*]] = and i8 [[S:%.*]], 1
+; CHECK-NEXT:    [[INS:%.*]] = insertelement <2 x i8> [[LD]], i8 0, i8 [[A]]
+; CHECK-NEXT:    store i32 0, ptr [[P]], align 4
+; CHECK-NEXT:    store <2 x i8> [[INS]], ptr [[P]], align 2
+; CHECK-NEXT:    ret void
+;
+  %ld = load <2 x i8>, ptr %p
+  %a = and i8 %s, 1
+  %ins = insertelement <2 x i8> %ld, i8 0, i8 %a
+  store i32 0, ptr %p
+  store <2 x i8> %ins, ptr %p
+  ret void
+}
+
 ; Check cases when calls may modify memory
 define void @insert_store_with_call(ptr %p, ptr %q, i8 %s) {
 ; CHECK-LABEL: @insert_store_with_call(

``````````

</details>


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