[Mlir-commits] [mlir] 5faa5f8 - [MLIR][Affine] Fix copy generation for missing memref definition depth check (#129187)

[llvmlistbot at llvm.org]

Author: Uday Bondhugula
Date: 2025-03-01T02:46:08+05:30
New Revision: 5faa5f848a35de13196f2f516f51aa970da942b4

URL: https://github.com/llvm/llvm-project/commit/5faa5f848a35de13196f2f516f51aa970da942b4
DIFF: https://github.com/llvm/llvm-project/commit/5faa5f848a35de13196f2f516f51aa970da942b4.diff

LOG: [MLIR][Affine] Fix copy generation for missing memref definition depth check (#129187)

Fixes: https://github.com/llvm/llvm-project/issues/122210

Added: 
    

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

Removed: 
    


################################################################################
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<IntegerAttr>(memrefType.getMemorySpace());
     if ((filterMemRef.has_value() && filterMemRef != memref) ||
-        (memorySpaceAttr &&
+        (isa_and_nonnull<IntegerAttr>(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<MemRefRegion>(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
+}