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

[Rik Huijzer]

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

>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 1/2] [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 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) {

>From 3ec4bef689932b6bf671d37a61b353e661fb87ad Mon Sep 17 00:00:00 2001
From: Rik Huijzer <github at huijzer.xyz>
Date: Fri, 8 Dec 2023 17:38:39 +0100
Subject: [PATCH 2/2] Update
 mlir/lib/Dialect/LLVMIR/Transforms/TypeConsistency.cpp

Co-authored-by: Tobias Gysi <tobias.gysi at nextsilicon.com>
---
 mlir/lib/Dialect/LLVMIR/Transforms/TypeConsistency.cpp | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/mlir/lib/Dialect/LLVMIR/Transforms/TypeConsistency.cpp b/mlir/lib/Dialect/LLVMIR/Transforms/TypeConsistency.cpp
index 3e856827255149..cf900ac0be8fd2 100644
--- a/mlir/lib/Dialect/LLVMIR/Transforms/TypeConsistency.cpp
+++ b/mlir/lib/Dialect/LLVMIR/Transforms/TypeConsistency.cpp
@@ -164,7 +164,7 @@ static std::optional<uint64_t> gepToByteOffset(DataLayout &layout, GEPOp gep) {
     int32_t gepIndex = indexInt.getInt();
     if (gepIndex < 0)
       return std::nullopt;
-    indices.push_back((uint32_t)gepIndex);
+    indices.push_back(static_cast<uint32_t>(gepIndex));
   }
 
   uint64_t offset = indices[0] * layout.getTypeSize(gep.getElemType());