diff --git a/mlir/lib/Dialect/Affine/Utils/LoopUtils.cpp b/mlir/lib/Dialect/Affine/Utils/LoopUtils.cpp index a8c24e1423425..c7e412b2b0fd9 100644 --- a/mlir/lib/Dialect/Affine/Utils/LoopUtils.cpp +++ b/mlir/lib/Dialect/Affine/Utils/LoopUtils.cpp @@ -2332,17 +2332,21 @@ mlir::affine::affineDataCopyGenerate(Block::iterator begin, Block::iterator end, memref = storeOp.getMemRef(); memrefType = storeOp.getMemRefType(); } - // Neither load nor a store op. + // Not an affine.load/store op. if (!memref) return; - auto memorySpaceAttr = - dyn_cast_or_null(memrefType.getMemorySpace()); if ((filterMemRef.has_value() && filterMemRef != memref) || - (memorySpaceAttr && + (isa_and_nonnull(memrefType.getMemorySpace()) && memrefType.getMemorySpaceAsInt() != copyOptions.slowMemorySpace)) return; + if (!memref.getParentRegion()->isAncestor(block->getParent())) { + LLVM_DEBUG(llvm::dbgs() << "memref definition is inside of the depth at " + "which copy-in/copy-out would happen\n"); + return; + } + // Compute the MemRefRegion accessed. auto region = std::make_unique(opInst->getLoc()); if (failed(region->compute(opInst, copyDepth, /*sliceState=*/nullptr, diff --git a/mlir/test/Dialect/Affine/affine-data-copy.mlir b/mlir/test/Dialect/Affine/affine-data-copy.mlir index 26eef0a7925a7..453a0eabc4fdd 100644 --- a/mlir/test/Dialect/Affine/affine-data-copy.mlir +++ b/mlir/test/Dialect/Affine/affine-data-copy.mlir @@ -419,3 +419,15 @@ func.func @scalar_memref_copy_in_loop(%3:memref<480xi1>) { // CHECK: memref.dealloc %[[FAST_MEMREF]] : memref<480xi1> return } + +// CHECK-LABEL: func @memref_def_inside +func.func @memref_def_inside(%arg0: index) { + %0 = llvm.mlir.constant(1.000000e+00 : f32) : f32 + // No copy generation can happen at this depth given the definition inside. + affine.for %arg1 = 0 to 29 { + %alloc_7 = memref.alloc() : memref<1xf32> + // CHECK: affine.store {{.*}} : memref<1xf32> + affine.store %0, %alloc_7[0] : memref<1xf32> + } + return +}