[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

mlir/lib/Transforms/Utils/LoopUtils.cpp

@@ -2912,8 +2912,12 @@ LogicalResult mlir::generateCopyForMemRegion(
   if (failed(err))
     return err;
 
-  result.alloc =
+  const auto &en = fastBufferMap.find(memrefRegion.memref);
-      fastBufferMap.find(memrefRegion.memref)->second.getDefiningOp();
+  // 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();

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
+}