From 071d26f8082391612f3a3f71b1135cbdceb0a30a Mon Sep 17 00:00:00 2001 From: Uday Bondhugula Date: Wed, 30 Jun 2021 09:50:33 +0530 Subject: [PATCH] [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 | 8 ++++++-- mlir/test/Dialect/Affine/affine-data-copy.mlir | 13 +++++++++++++ 2 files changed, 19 insertions(+), 2 deletions(-) diff --git a/mlir/lib/Transforms/Utils/LoopUtils.cpp b/mlir/lib/Transforms/Utils/LoopUtils.cpp index 5e71e706ff7f..ac3f87ea25f4 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 11288784a4e9..243f9d0b6531 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 +}