[llvm] [VectorCombine] Check load offsets in APInt before narrowing shuffle indices (PR #210904)

[via llvm-commits]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-vectorizers

Author: hanbeom (ParkHanbum)

<details>
<summary>Changes</summary>

vectorizeLoadInsert narrowed the APInt element offset to unsigned before
checking its range. Large offsets could wrap and select the wrong vector
element.

Check the offset range as an APInt before converting it to the shuffle index.
This prevents invalid load widening and is covered by a regression test.

Fixes #<!-- -->210903

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


2 Files Affected:

- (modified) llvm/lib/Transforms/Vectorize/VectorCombine.cpp (+3-2) 
- (modified) llvm/test/Transforms/VectorCombine/X86/load-inseltpoison.ll (+16) 


``````````diff
diff --git a/llvm/lib/Transforms/Vectorize/VectorCombine.cpp b/llvm/lib/Transforms/Vectorize/VectorCombine.cpp
index ef467b4c91aff..f623a20d067ca 100644
--- a/llvm/lib/Transforms/Vectorize/VectorCombine.cpp
+++ b/llvm/lib/Transforms/Vectorize/VectorCombine.cpp
@@ -294,9 +294,10 @@ bool VectorCombine::vectorizeLoadInsert(Instruction &I) {
       return false;
 
     // If we load MinVecNumElts, will our target element still be loaded?
-    OffsetEltIndex = Offset.udiv(ScalarSizeInBytes).getZExtValue();
-    if (OffsetEltIndex >= MinVecNumElts)
+    APInt OffsetEltIndexAP = Offset.udiv(ScalarSizeInBytes);
+    if (OffsetEltIndexAP.uge(MinVecNumElts))
       return false;
+    OffsetEltIndex = OffsetEltIndexAP.getZExtValue();
 
     if (!isSafeToLoadUnconditionally(SrcPtr, MinVecTy, Align(1), *DL, Load,
                                      SQ.AC, SQ.DT))
diff --git a/llvm/test/Transforms/VectorCombine/X86/load-inseltpoison.ll b/llvm/test/Transforms/VectorCombine/X86/load-inseltpoison.ll
index e8381d1b206e2..8cd99bbf31a8c 100644
--- a/llvm/test/Transforms/VectorCombine/X86/load-inseltpoison.ll
+++ b/llvm/test/Transforms/VectorCombine/X86/load-inseltpoison.ll
@@ -336,6 +336,22 @@ define <4 x i32> @gep013_bitcast_load_i32_insert_v4i32(ptr align 1 dereferenceab
   ret <4 x i32> %r
 }
 
+; The element index is 2^32 and does not fit in unsigned. It must be rejected
+; before conversion to the shuffle-mask index type.
+
+define <4 x i32> @load_insert_large_offset(ptr align 16 dereferenceable(17179869188) %p) {
+; CHECK-LABEL: @load_insert_large_offset(
+; CHECK-NEXT:    [[GEP:%.*]] = getelementptr inbounds i8, ptr [[P:%.*]], i64 17179869184
+; CHECK-NEXT:    [[S:%.*]] = load i32, ptr [[GEP]], align 4
+; CHECK-NEXT:    [[R:%.*]] = insertelement <4 x i32> poison, i32 [[S]], i64 0
+; CHECK-NEXT:    ret <4 x i32> [[R]]
+;
+  %gep = getelementptr inbounds i8, ptr %p, i64 17179869184
+  %s = load i32, ptr %gep, align 4
+  %r = insertelement <4 x i32> poison, i32 %s, i64 0
+  ret <4 x i32> %r
+}
+
 ; If there are enough dereferenceable bytes, we can offset the vector load.
 
 define <8 x i16> @gep10_load_i16_insert_v8i16(ptr align 16 dereferenceable(32) %p) nofree nosync {

``````````

</details>


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