[MLIR] Fix generateCopyForMemRefRegion

Fix generateCopyForMemRefRegion for a missing check: in some cases, when
the thing to generate copies for itself is empty, no fast buffer/copy
loops would have been allocated/generated. Add an extra assertion there
while at this.

Differential Revision: https://reviews.llvm.org/D105170

diff --git a/mlir/lib/Transforms/Utils/LoopUtils.cpp b/mlir/lib/Transforms/Utils/LoopUtils.cpp
index 5e71e70..ac3f87e 100644 (file)
--- a/mlir/lib/Transforms/Utils/LoopUtils.cpp
+++ b/mlir/lib/Transforms/Utils/LoopUtils.cpp
@@ -2912,8 +2912,12 @@ LogicalResult mlir::generateCopyForMemRegion(
   if (failed(err))
     return err;
 
-  result.alloc =
-      fastBufferMap.find(memrefRegion.memref)->second.getDefiningOp();
+  const auto &en = fastBufferMap.find(memrefRegion.memref);
+  // In some cases (empty loops), no copy generation would have happened.
+  if (en == fastBufferMap.end())
+    return failure();
+  result.alloc = en->second.getDefiningOp();
+  assert(result.alloc && "fast buffer expected to be locally allocated");
   assert(copyNests.size() <= 1 && "At most one copy nest is expected.");
   result.copyNest = copyNests.empty() ? nullptr : *copyNests.begin();
   return success();

diff --git a/mlir/test/Dialect/Affine/affine-data-copy.mlir b/mlir/test/Dialect/Affine/affine-data-copy.mlir
index 1128878..243f9d0 100644 (file)
--- a/mlir/test/Dialect/Affine/affine-data-copy.mlir
+++ b/mlir/test/Dialect/Affine/affine-data-copy.mlir
@@ -270,3 +270,16 @@ func @max_lower_bound(%M: memref<2048x516xf64>, %i : index, %j : index) {
 // CHECK-NEXT:    }
 // CHECK-NEXT: }
 // CHECK-NEXT: memref.dealloc %[[BUF]] : memref<2048x6xf64>
+
+// -----
+
+// CHECK-LABEL: func @empty_loop
+func @empty_loop(%arg0: memref<1024x1024xf64>) {
+  // Empty loop - so no copy generation happens.
+  affine.for %i = 0 to 0 {
+    affine.load %arg0[0, %i] : memref<1024x1024xf64>
+  }
+  return
+  // CHECK-NOT:    memref.alloc
+  // CHECK:        return
+}