diff --git a/mlir/lib/Transforms/Utils/LoopUtils.cpp b/mlir/lib/Transforms/Utils/LoopUtils.cpp --- a/mlir/lib/Transforms/Utils/LoopUtils.cpp +++ b/mlir/lib/Transforms/Utils/LoopUtils.cpp @@ -2912,8 +2912,12 @@ 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 --- a/mlir/test/Dialect/Affine/affine-data-copy.mlir +++ b/mlir/test/Dialect/Affine/affine-data-copy.mlir @@ -270,3 +270,16 @@ // 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 +}