Fix 0-d memref corner case for getMemRefRegion()

- fix crash on test/Transforms/canonicalize.mlir with
  -memref-bound-check

PiperOrigin-RevId: 228268486

diff --git a/mlir/lib/Analysis/Utils.cpp b/mlir/lib/Analysis/Utils.cpp
index d3beb78eb926326cd3e0ac24311bdd35d590cba2..8b241ede3646ed375d46f6b952b2e783bd331e5b 100644 (file)
--- a/mlir/lib/Analysis/Utils.cpp
+++ b/mlir/lib/Analysis/Utils.cpp
@@ -134,6 +134,15 @@ bool mlir::getMemRefRegion(OperationInst *opInst, unsigned loopDepth,
   // Build the constraints for this region.
   FlatAffineConstraints *regionCst = region->getConstraints();
 
+  if (rank == 0) {
+    // A rank 0 memref has a 0-d region.
+    SmallVector<ForInst *, 4> ivs;
+    getLoopIVs(*opInst, &ivs);
+    SmallVector<Value *, 4> regionSymbols(ivs.begin(), ivs.end());
+    regionCst->reset(0, loopDepth, 0, regionSymbols);
+    return true;
+  }
+
   FuncBuilder b(opInst);
   auto idMap = b.getMultiDimIdentityMap(rank);
 

diff --git a/mlir/test/Transforms/memref-bound-check.mlir b/mlir/test/Transforms/memref-bound-check.mlir
index 757844994276746fb9bade8533aa5dbf97f431fa..75bc25d326acd8e9ec4ed3677a55e398c219bd98 100644 (file)
--- a/mlir/test/Transforms/memref-bound-check.mlir
+++ b/mlir/test/Transforms/memref-bound-check.mlir
@@ -160,3 +160,11 @@ func @delinearize_mod_floordiv() {
   }
   return
 }
+
+// CHECK-LABEL: func @zero_d_memref
+func @zero_d_memref(%arg0: memref<i32>) {
+  %c0 = constant 0 : i32
+  // A 0-d memref always has in-bound accesses!
+  store %c0, %arg0[] : memref<i32>
+  return
+}