[MLIR][GPU] Fix memref.dim folding with out-of-bound index

[joker-eph]

Fixes #118760

[llvmbot]

@llvm/pr-subscribers-mlir-gpu

Author: Mehdi Amini (joker-eph)

Changes

Fixes #118760

---

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

2 Files Affected:

• (modified) mlir/lib/Dialect/GPU/IR/GPUDialect.cpp (+2-1)
• (modified) mlir/test/Dialect/GPU/canonicalize.mlir (+11)

diff --git a/mlir/lib/Dialect/GPU/IR/GPUDialect.cpp b/mlir/lib/Dialect/GPU/IR/GPUDialect.cpp
index 2f96f10b662f87..ee00fbeb28b61d 100644
--- a/mlir/lib/Dialect/GPU/IR/GPUDialect.cpp
+++ b/mlir/lib/Dialect/GPU/IR/GPUDialect.cpp
@@ -2075,7 +2075,8 @@ struct SimplifyDimOfAllocOp : public OpRewritePattern<memref::DimOp> {
       return failure();
 
     auto memrefType = llvm::dyn_cast<MemRefType>(dimOp.getSource().getType());
-    if (!memrefType || !memrefType.isDynamicDim(index.value()))
+    if (!memrefType || index.value() >= memrefType.getRank() ||
+        !memrefType.isDynamicDim(index.value()))
       return failure();
 
     auto alloc = dimOp.getSource().getDefiningOp<AllocOp>();
diff --git a/mlir/test/Dialect/GPU/canonicalize.mlir b/mlir/test/Dialect/GPU/canonicalize.mlir
index d342ae9df10eea..33ce98e6da0ed2 100644
--- a/mlir/test/Dialect/GPU/canonicalize.mlir
+++ b/mlir/test/Dialect/GPU/canonicalize.mlir
@@ -152,6 +152,17 @@ func.func @gpu_dim_of_alloc(%size: index) -> index {
 
 // -----
 
+// CHECK-LABEL: func @out_of_bound_memref.dim
+//  CHECK:   %[[MEMREF:.[a-z0-9A-Z_]+]] = memref.dim
+//  CHECK:   return %[[MEMREF]] : index
+func.func @out_of_bound_memref.dim(%arg : memref<?xi8>, %size: index) -> index {
+  %c2 = arith.constant 2 : index
+  %1 = memref.dim %arg, %c2 : 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

[llvmbot]

@llvm/pr-subscribers-mlir

Author: Mehdi Amini (joker-eph)

Changes

Fixes #118760

---

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

2 Files Affected:

• (modified) mlir/lib/Dialect/GPU/IR/GPUDialect.cpp (+2-1)
• (modified) mlir/test/Dialect/GPU/canonicalize.mlir (+11)

diff --git a/mlir/lib/Dialect/GPU/IR/GPUDialect.cpp b/mlir/lib/Dialect/GPU/IR/GPUDialect.cpp
index 2f96f10b662f87..ee00fbeb28b61d 100644
--- a/mlir/lib/Dialect/GPU/IR/GPUDialect.cpp
+++ b/mlir/lib/Dialect/GPU/IR/GPUDialect.cpp
@@ -2075,7 +2075,8 @@ struct SimplifyDimOfAllocOp : public OpRewritePattern<memref::DimOp> {
       return failure();
 
     auto memrefType = llvm::dyn_cast<MemRefType>(dimOp.getSource().getType());
-    if (!memrefType || !memrefType.isDynamicDim(index.value()))
+    if (!memrefType || index.value() >= memrefType.getRank() ||
+        !memrefType.isDynamicDim(index.value()))
       return failure();
 
     auto alloc = dimOp.getSource().getDefiningOp<AllocOp>();
diff --git a/mlir/test/Dialect/GPU/canonicalize.mlir b/mlir/test/Dialect/GPU/canonicalize.mlir
index d342ae9df10eea..33ce98e6da0ed2 100644
--- a/mlir/test/Dialect/GPU/canonicalize.mlir
+++ b/mlir/test/Dialect/GPU/canonicalize.mlir
@@ -152,6 +152,17 @@ func.func @gpu_dim_of_alloc(%size: index) -> index {
 
 // -----
 
+// CHECK-LABEL: func @out_of_bound_memref.dim
+//  CHECK:   %[[MEMREF:.[a-z0-9A-Z_]+]] = memref.dim
+//  CHECK:   return %[[MEMREF]] : index
+func.func @out_of_bound_memref.dim(%arg : memref<?xi8>, %size: index) -> index {
+  %c2 = arith.constant 2 : index
+  %1 = memref.dim %arg, %c2 : 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