[Mlir-commits] [mlir] 071d26f - [MLIR] Fix generateCopyForMemRefRegion

[Uday Bondhugula]

Author: Uday Bondhugula
Date: 2021-06-30T10:24:10+05:30
New Revision: 071d26f8082391612f3a3f71b1135cbdceb0a30a

URL: https://github.com/llvm/llvm-project/commit/071d26f8082391612f3a3f71b1135cbdceb0a30a
DIFF: https://github.com/llvm/llvm-project/commit/071d26f8082391612f3a3f71b1135cbdceb0a30a.diff

LOG: [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

Added: 
    

Modified: 
    mlir/lib/Transforms/Utils/LoopUtils.cpp
    mlir/test/Dialect/Affine/affine-data-copy.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Transforms/Utils/LoopUtils.cpp b/mlir/lib/Transforms/Utils/LoopUtils.cpp
index 5e71e706ff7f3..ac3f87ea25f41 100644
--- 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 11288784a4e91..243f9d0b65319 100644
--- 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
+}