[llvm] [LSV] Added check for mismatched GEP strides in getConstantOffsetComplexAddrs (PR #186671)

[via llvm-commits]

llvmbot wrote:


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

@llvm/pr-subscribers-vectorizers

Author: Weiwen He (he-weiwen)

<details>
<summary>Changes</summary>

When two GEPs A & B with variable indices have source element types of different sizes (e.g. i8 vs i32), their offsets cannot be constants. Currently the `getConstantOffsetComplexAddrs` does not check this and simply use the stride of GEP A assuming they are equal, resulting in miscompile.

Add a check of the stride against the stride of B and bailout if they mismatch. An early check on source element type should also work but the semantics is slightly different (different types may have the same stride), though I'm not sure if an example can be constructed to demonstrate the difference.

Fixes #<!-- -->186327

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


2 Files Affected:

- (modified) llvm/lib/Transforms/Vectorize/LoadStoreVectorizer.cpp (+2) 
- (added) llvm/test/Transforms/LoadStoreVectorizer/X86/mismatched-gep-stride.ll (+32) 


``````````diff
diff --git a/llvm/lib/Transforms/Vectorize/LoadStoreVectorizer.cpp b/llvm/lib/Transforms/Vectorize/LoadStoreVectorizer.cpp
index e6bb7bd076e4d..7cffebfc502f7 100644
--- a/llvm/lib/Transforms/Vectorize/LoadStoreVectorizer.cpp
+++ b/llvm/lib/Transforms/Vectorize/LoadStoreVectorizer.cpp
@@ -1468,6 +1468,8 @@ std::optional<APInt> Vectorizer::getConstantOffsetComplexAddrs(
     return std::nullopt;
 
   uint64_t Stride = GTIA.getSequentialElementStride(DL);
+  if (Stride != GTIB.getSequentialElementStride(DL))
+    return std::nullopt;
 
   // Only look through a ZExt/SExt.
   if (!isa<SExtInst>(OpA) && !isa<ZExtInst>(OpA))
diff --git a/llvm/test/Transforms/LoadStoreVectorizer/X86/mismatched-gep-stride.ll b/llvm/test/Transforms/LoadStoreVectorizer/X86/mismatched-gep-stride.ll
new file mode 100644
index 0000000000000..ecde58f104942
--- /dev/null
+++ b/llvm/test/Transforms/LoadStoreVectorizer/X86/mismatched-gep-stride.ll
@@ -0,0 +1,32 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 6
+; RUN: opt -mtriple=x86_64-- -passes=load-store-vectorizer -S -o - %s | FileCheck %s
+
+; Two GEPs with different source element types and different strides
+; must not be vectorized. The byte offset cannot be constant.
+define void @mismatched_stride(ptr %p, ptr %idx_ptr, ptr %out) {
+; CHECK-LABEL: define void @mismatched_stride(
+; CHECK-SAME: ptr [[P:%.*]], ptr [[IDX_PTR:%.*]], ptr [[OUT:%.*]]) {
+; CHECK-NEXT:    [[X:%.*]] = load i32, ptr [[IDX_PTR]], align 4
+; CHECK-NEXT:    [[Y:%.*]] = add nsw i32 [[X]], 1
+; CHECK-NEXT:    [[SX:%.*]] = sext i32 [[X]] to i64
+; CHECK-NEXT:    [[SY:%.*]] = sext i32 [[Y]] to i64
+; CHECK-NEXT:    [[A:%.*]] = getelementptr i8, ptr [[P]], i64 [[SX]]
+; CHECK-NEXT:    [[B:%.*]] = getelementptr i32, ptr [[P]], i64 [[SY]]
+; CHECK-NEXT:    [[VA:%.*]] = load i8, ptr [[A]], align 2
+; CHECK-NEXT:    [[VB:%.*]] = load i8, ptr [[B]], align 2
+; CHECK-NEXT:    [[SUM:%.*]] = add i8 [[VA]], [[VB]]
+; CHECK-NEXT:    store i8 [[SUM]], ptr [[OUT]], align 1
+; CHECK-NEXT:    ret void
+;
+  %x = load i32, ptr %idx_ptr
+  %y = add nsw i32 %x, 1
+  %sx = sext i32 %x to i64
+  %sy = sext i32 %y to i64
+  %a = getelementptr i8, ptr %p, i64 %sx
+  %b = getelementptr i32, ptr %p, i64 %sy
+  %va = load i8, ptr %a, align 2
+  %vb = load i8, ptr %b, align 2
+  %sum = add i8 %va, %vb
+  store i8 %sum, ptr %out
+  ret void
+}

``````````

</details>


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