[Mlir-commits] [mlir] [mlir][llvm] Fix negative GEP crash in type consistency (PR #74859)

[Rik Huijzer]

https://github.com/rikhuijzer created https://github.com/llvm/llvm-project/pull/74859

Fixes https://github.com/llvm/llvm-project/issues/74453.

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

>From 77cea9416e1cf480c02bd7dfda98d5d8741f19d3 Mon Sep 17 00:00:00 2001
From: Rik Huijzer <github at huijzer.xyz>
Date: Fri, 8 Dec 2023 17:17:08 +0100
Subject: [PATCH] [mlir][llvm] Fix negative GEP crash in type consistency

---
 .../Dialect/LLVMIR/Transforms/TypeConsistency.cpp  |  5 ++++-
 mlir/test/Dialect/LLVMIR/type-consistency.mlir     | 14 ++++++++++++++
 2 files changed, 18 insertions(+), 1 deletion(-)

diff --git a/mlir/lib/Dialect/LLVMIR/Transforms/TypeConsistency.cpp b/mlir/lib/Dialect/LLVMIR/Transforms/TypeConsistency.cpp
index b094c650ff193..3e85682725514 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 1504a98e6f8cc..3a1ab924ebdac 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) {