[mlir][llvm] Fix negative GEP crash in type consistency

[rikhuijzer]

Fixes #74453.

The gepToByteOffset was implicitly casting an signed integer to an unsigned integer even though negative dimensions are valid for llvm.getelementptr.

[llvmbot]

@llvm/pr-subscribers-mlir

@llvm/pr-subscribers-mlir-llvm

Author: Rik Huijzer (rikhuijzer)

Changes

Fixes #74453.

The gepToByteOffset was implicitly casting an signed integer to an unsigned integer even though negative dimensions are valid for llvm.getelementptr.

---

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

2 Files Affected:

• (modified) mlir/lib/Dialect/LLVMIR/Transforms/TypeConsistency.cpp (+4-1)
• (modified) mlir/test/Dialect/LLVMIR/type-consistency.mlir (+14)

diff --git a/mlir/lib/Dialect/LLVMIR/Transforms/TypeConsistency.cpp b/mlir/lib/Dialect/LLVMIR/Transforms/TypeConsistency.cpp
index b094c650ff1932..3e856827255149 100644
--- a/mlir/lib/Dialect/LLVMIR/Transforms/TypeConsistency.cpp
+++ b/mlir/lib/Dialect/LLVMIR/Transforms/TypeConsistency.cpp
@@ -161,7 +161,10 @@ static std::optional<uint64_t> gepToByteOffset(DataLayout &layout, GEPOp gep) {
     IntegerAttr indexInt = llvm::dyn_cast_if_present<IntegerAttr>(index);
     if (!indexInt)
       return std::nullopt;
-    indices.push_back(indexInt.getInt());
+    int32_t gepIndex = indexInt.getInt();
+    if (gepIndex < 0)
+      return std::nullopt;
+    indices.push_back((uint32_t)gepIndex);
   }
 
   uint64_t offset = indices[0] * layout.getTypeSize(gep.getElemType());
diff --git a/mlir/test/Dialect/LLVMIR/type-consistency.mlir b/mlir/test/Dialect/LLVMIR/type-consistency.mlir
index 1504a98e6f8cca..3a1ab924ebdacb 100644
--- a/mlir/test/Dialect/LLVMIR/type-consistency.mlir
+++ b/mlir/test/Dialect/LLVMIR/type-consistency.mlir
@@ -151,6 +151,20 @@ llvm.func @index_to_struct(%arg: i32) {
 
 // -----
 
+// CHECK-LABEL: llvm.func @no_crash_on_negative_gep_index
+llvm.func @no_crash_on_negative_gep_index() {
+  %0 = llvm.mlir.constant(1.000000e+00 : f16) : f16
+  %1 = llvm.mlir.constant(1 : i32) : i32
+  // CHECK: %[[ALLOCA:.*]] = llvm.alloca %{{.*}} x !llvm.struct<"foo", (i32, i32, i32)>
+  %2 = llvm.alloca %1 x !llvm.struct<"foo", (i32, i32, i32)> : (i32) -> !llvm.ptr
+  // CHECK: llvm.getelementptr %[[ALLOCA]][-1] : (!llvm.ptr) -> !llvm.ptr, f32
+  %3 = llvm.getelementptr %2[-1] : (!llvm.ptr) -> !llvm.ptr, f32
+  llvm.store %0, %3 : f16, !llvm.ptr
+  llvm.return
+}
+
+// -----
+
 // CHECK-LABEL: llvm.func @coalesced_store_ints
 // CHECK-SAME: %[[ARG:.*]]: i64
 llvm.func @coalesced_store_ints(%arg: i64) {

[gysit]

LGTM modulo the C style cast.

[rikhuijzer]

Thanks for the review @gysit! I'll merge once CI on main is green again.