[Mlir-commits] [mlir] [mlir][gpu] Fix memref.dim folding with negative index (PR #205338)

[Longsheng Mou]

https://github.com/CoTinker created https://github.com/llvm/llvm-project/pull/205338

Fixes #205073.

>From be94cbcb19c8b95582dfe0c82e53f6070c306d7b Mon Sep 17 00:00:00 2001
From: Longsheng Mou <longshengmou at gmail.com>
Date: Tue, 23 Jun 2026 11:26:01 +0800
Subject: [PATCH] [mlir][gpu] Fix memref.dim folding with negative index

---
 mlir/lib/Dialect/GPU/IR/GPUDialect.cpp  |  7 ++++---
 mlir/test/Dialect/GPU/canonicalize.mlir | 11 +++++++++++
 2 files changed, 15 insertions(+), 3 deletions(-)

diff --git a/mlir/lib/Dialect/GPU/IR/GPUDialect.cpp b/mlir/lib/Dialect/GPU/IR/GPUDialect.cpp
index a178bb453d86b..eaea0142e0438 100644
--- a/mlir/lib/Dialect/GPU/IR/GPUDialect.cpp
+++ b/mlir/lib/Dialect/GPU/IR/GPUDialect.cpp
@@ -2323,9 +2323,10 @@ struct SimplifyDimOfAllocOp : public OpRewritePattern<memref::DimOp> {
     if (!index)
       return failure();
 
+    int64_t indexVal = index.value();
     auto memrefType = llvm::dyn_cast<MemRefType>(dimOp.getSource().getType());
-    if (!memrefType || index.value() >= memrefType.getRank() ||
-        !memrefType.isDynamicDim(index.value()))
+    if (!memrefType || indexVal < 0 || indexVal >= memrefType.getRank() ||
+        !memrefType.isDynamicDim(indexVal))
       return failure();
 
     auto alloc = dimOp.getSource().getDefiningOp<AllocOp>();
@@ -2333,7 +2334,7 @@ struct SimplifyDimOfAllocOp : public OpRewritePattern<memref::DimOp> {
       return failure();
 
     Value substituteOp = *(alloc.getDynamicSizes().begin() +
-                           memrefType.getDynamicDimIndex(index.value()));
+                           memrefType.getDynamicDimIndex(indexVal));
     rewriter.replaceOp(dimOp, substituteOp);
     return success();
   }
diff --git a/mlir/test/Dialect/GPU/canonicalize.mlir b/mlir/test/Dialect/GPU/canonicalize.mlir
index 7627af11c636c..9943c8b3a572e 100644
--- a/mlir/test/Dialect/GPU/canonicalize.mlir
+++ b/mlir/test/Dialect/GPU/canonicalize.mlir
@@ -252,6 +252,17 @@ func.func @out_of_bound_memref.dim(%arg : memref<?xi8>, %size: index) -> index {
 
 // -----
 
+// CHECK-LABEL: func @negative_memref_dim
+//  CHECK:   %[[MEMREF:.*]] = memref.dim
+//  CHECK:   return %[[MEMREF]] : index
+func.func @negative_memref_dim(%arg: memref<?xi8>) -> index {
+  %c-2 = arith.constant -2 : index
+  %1 = memref.dim %arg, %c-2 : memref<?xi8>
+  return %1 : index
+}
+
+// -----
+
 // CHECK-LABEL: func @simplify_gpu_launch
 func.func @simplify_gpu_launch() attributes {llvm.emit_c_interface} {
   %cst = arith.constant 0.000000e+00 : f32