[MLIR] Fix generateCopyForMemRefRegion

[bondhugula]

Summary

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.

[bondhugula]

Add a comment.

mlir/lib/Transforms/Utils/LoopUtils.cpp

LogicalResult mlir::generateCopyForMemRegion(
   DenseSet<Operation *> copyNests;
 
   auto err = generateCopy(memrefRegion, block, begin, end, block, begin, end,
                           copyOptions, fastBufferMap, copyNests,
                           &result.sizeInBytes, &begin, &end);
   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();
 }
 
 /// Gathers all AffineForOps in 'block' at 'currLoopDepth' in 'depthToLoops'.
 static void
 gatherLoopsInBlock(Block *block, unsigned currLoopDepth,

mlir/test/Dialect/Affine/affine-data-copy.mlir

 // CHECK-NEXT:   }
 // CHECK-NEXT: }
 // CHECK-NEXT: affine.for %[[ii_:.*]] = 0 to 2048 {
 // CHECK-NEXT:   affine.for %[[jj_:.*]] = max #[[$LB]]()[%{{.*}}, %{{.*}}] to min #[[$UB]]()[%{{.*}}, %{{.*}}] {
 // CHECK-NEXT:     affine.load %[[BUF]][%[[ii_]], %[[jj_]] - symbol(%[[j]]) * 6] : memref<2048x6xf64>
 // 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
+}