[mlir] Add MemRef filter to affine data copy optimization

[dcaballe]

Summary

This patch extends affine data copy optimization utility with an
optional memref filter argument. When the memref filter is used, data
copy optimization will only generate copies for such a memref.

Note: this patch is just porting the memref filter feature from Uday's
'hop' branch: https://github.com/bondhugula/llvm-project/tree/hop.

[bondhugula]

This looks great to me! Just a few minor comments.

mlir/include/mlir/Transforms/LoopUtils.h

18

I think we shouldn't need this.

174–175

Doc comment update needed for filter memref.

mlir/test/lib/Transforms/TestAffineDataCopy.cpp

10

-> "utility functions and options".

[mehdi_amini]

mlir/include/mlir/Transforms/LoopUtils.h

228

Seems like a private implementation for the one below, do we need to expose it publicly?
The API is leaking the fact that this is recursion helper...

mlir/lib/Transforms/Utils/LoopUtils.cpp

1797–1798

DenseMap is a surprising choice to me for a container that will be filled with monotonic indices starting at zero...
Can this be just a SmallVectorImpl<SmallVector<AffineForOp, 2>> ?

I think you just need to add this:

if (currLoopDepth >= depthToLoops.size()) depthToLoops.resize(currLoopDepth);

at the beginning of the function and change nothing else.

mlir/test/lib/Transforms/TestAffineDataCopy.cpp

30

Just a nit: can you use a pass option. Not that it matters much for a test pass but it'll display more consistently for mlir-opt --help

[rriddle]

mlir/lib/Transforms/Utils/LoopUtils.cpp

1639–1640

You can just do filterMemRef != loadOp.getMemRef()

1644–1646

Same here.

1804

Drop trivial braces.

[dcaballe]

• Address feedback

[dcaballe]

Thanks for the feedback!

mlir/include/mlir/Transforms/LoopUtils.h

174–175

Thanks Uday! I had added doc to the cpp file. I'll bring both docs into sync.

228

Good point! Thanks!

mlir/lib/Transforms/Utils/LoopUtils.cpp

1639–1640

We need the hasValue part. Otherwise, the condition would be true when filterMemRef has no value. I removed the getValue call, though :) Thanks!

1797–1798

I forgot to ask the same question to @andydavis1! I thought that maybe there was a use case that I wasn't taking into account. Even if we wanted to update the container after a loop transformation, we would be adding or removing loop levels to/from the back. ?

Another comment about the current implementation is that it adds an extra "empty" level. When we process the innermost loop and we do:

auto &loopsAtDepth = depthToLoops[currLoopDepth];

we create a new level that won't contain any loop.

I'll take care of both things if it makes sense to you.

if (currLoopDepth >= depthToLoops.size()) depthToLoops.resize(currLoopDepth);

Not sure I follow. Wouldn't this condition be always false and cause a realloc every time we add a new level?

[mehdi_amini]

mlir/lib/Transforms/Utils/LoopUtils.cpp

1797–1798

Not sure I follow. Wouldn't this condition be always false and cause a realloc every time we add a new level?

If this is a private helper yes you can skip the check.
If it is a public function you would resize by truncating if the depth was smaller than the vector size.

[dcaballe]

Rebasing on top of latest mlir-opt and TestLoopFusion changes.
If no more comments, I'll address 'gatherLoops' concerns in a
separate commit and keep this as a simple rafactoring of that
utility.

[bondhugula]

In D74342#1872854, @dcaballe wrote:

Rebasing on top of latest mlir-opt and TestLoopFusion changes.
If no more comments, I'll address 'gatherLoops' concerns in a
separate commit and keep this as a simple rafactoring of that
utility.

Yes, gatherLoops can be addressed separately; if better, you could keep gatherLoops private (and duplicated), and it could be addressed when made public. Either way is fine for me.

mlir/include/mlir/Transforms/LoopUtils.h

 // methods: these are not passes by themselves but are used either by passes,
 // optimization sequences, or in turn by other transformation utilities.
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_TRANSFORMS_LOOP_UTILS_H
 #define MLIR_TRANSFORMS_LOOP_UTILS_H
 
 #include "mlir/IR/Block.h"
 #include "mlir/Support/LLVM.h"
 #include "mlir/Support/LogicalResult.h"
 
 namespace mlir {
 class AffineForOp;
 class FuncOp;
 class OpBuilder;
 class Value;

 /// block iterator range [`begin', `end'), where `end' can't be past the
 /// terminator of the block (since additional operations are potentially
 /// inserted right before `end`. Returns the total size of fast memory space
 /// buffers used. `copyOptions` provides various parameters, and the output
 /// argument `copyNests` is the set of all copy nests inserted, each represented
 /// by its root affine.for. Since we generate alloc's and dealloc's for all fast
 /// buffers (before and after the range of operations resp. or at a hoisted
 /// position), all of the fast memory capacity is assumed to be available for
-/// processing this block range.
+/// processing this block range. When 'filterMemRef' is specified, copies are
+/// only generated for the provided MemRef.
 uint64_t affineDataCopyGenerate(Block::iterator begin, Block::iterator end,
                                 const AffineCopyOptions &copyOptions,
+                                Optional<Value> filterMemRef,
                                 DenseSet<Operation *> &copyNests);
 
 /// Tile a nest of standard for loops rooted at `rootForOp` by finding such
 /// parametric tile sizes that the outer loops have a fixed number of iterations
 /// as defined in `sizes`.
 TileLoops extractFixedOuterLoops(loop::ForOp rootFOrOp,
                                  ArrayRef<int64_t> sizes);
 

 /// ```
 ///    loop.for %i = %lb + %step * (threadIdx.x + blockIdx.x * blockDim.x)
 ///       to %ub step %gridDim.x * blockDim.x * %step {
 ///      ...
 ///    }
 /// ```
 void mapLoopToProcessorIds(loop::ForOp forOp, ArrayRef<Value> processorId,
                            ArrayRef<Value> numProcessors);
+
+/// Gathers all AffineForOps in 'func' grouped by loop depth.
+void gatherLoops(FuncOp func,
+                 DenseMap<unsigned, SmallVector<AffineForOp, 2>> &depthToLoops);
+
 } // end namespace mlir
 
 #endif // MLIR_TRANSFORMS_LOOP_UTILS_H

mlir/lib/Transforms/AffineDataCopyGeneration.cpp

AffineDataCopyGeneration::runOnBlock(Block *block,
   // Create [begin, end) ranges.
   auto it = curBegin;
   while (it != block->end()) {
     AffineForOp forOp;
     // If you hit a non-copy for loop, we will split there.
     if ((forOp = dyn_cast<AffineForOp>(&*it)) && copyNests.count(forOp) == 0) {
       // Perform the copying up unti this 'for' op first.
       affineDataCopyGenerate(/*begin=*/curBegin, /*end=*/it, copyOptions,
-                             copyNests);
+                             /*filterMemRef=*/llvm::None, copyNests);
 
       // Returns true if the footprint is known to exceed capacity.
       auto exceedsCapacity = [&](AffineForOp forOp) {
         Optional<int64_t> footprint =
             getMemoryFootprintBytes(forOp,
                                     /*memorySpace=*/0);
         return (footprint.hasValue() &&
                 static_cast<uint64_t>(footprint.getValue()) >
if ((forOp = dyn_cast<AffineForOp>(&*it)) && copyNests.count(forOp) == 0) {
         // portion of the block until 'it', and for 'it', which is 'forOp'. Note
         // that for the latter, the copies are placed just before this loop (for
         // incoming copies) and right after (for outgoing ones).
 
         // Inner loop copies have their own scope - we don't thus update
         // consumed capacity. The footprint check above guarantees this inner
         // loop's footprint fits.
         affineDataCopyGenerate(/*begin=*/it, /*end=*/std::next(it), copyOptions,
-                               copyNests);
+                               /*filterMemRef=*/llvm::None, copyNests);
       }
       // Get to the next load or store op after 'forOp'.
       curBegin = std::find_if(std::next(it), block->end(), [&](Operation &op) {
         return (isa<AffineLoadOp>(op) || isa<AffineStoreOp>(op) ||
                 isa<AffineForOp>(op)) &&
                copyNests.count(&op) == 0;
       });
       it = curBegin;
     } else {
       assert(copyNests.count(&*it) == 0 &&
              "all copy nests generated should have been skipped above");
       // We simply include this op in the current range and continue for more.
       ++it;
     }
   }
 
   // Generate the copy for the final block range.
   if (curBegin != block->end()) {
     // Can't be a terminator because it would have been skipped above.
     assert(!curBegin->isKnownTerminator() && "can't be a terminator");
     // Exclude the affine terminator - hence, the std::prev.
     affineDataCopyGenerate(/*begin=*/curBegin, /*end=*/std::prev(block->end()),
-                           copyOptions, copyNests);
+                           copyOptions, /*filterMemRef=*/llvm::None, copyNests);
   }
 
   return success();
 }
 
 void AffineDataCopyGeneration::runOnFunction() {
   FuncOp f = getFunction();
   OpBuilder topBuilder(f.getBody());

mlir/lib/Transforms/Utils/LoopUtils.cpp

for (unsigned d = 0; d < rank; d++) {
     auto dimSize = memRefType.getDimSize(d);
     assert(dimSize > 0 && "filtered dynamic shapes above");
     regionCst->addConstantLowerBound(d, 0);
     regionCst->addConstantUpperBound(d, dimSize - 1);
   }
   return true;
 }
 
-/// Generates copies for a contiguous sequence of operations in `block` in the
-/// iterator range [`begin', `end'), where `end' can't be past the terminator of
-/// the block (since additional operations are potentially inserted right before
-/// `end'. Returns the total size of the fast buffers used.
-//  Since we generate alloc's and dealloc's for all fast buffers (before and
-//  after the range of operations resp.), all of the fast memory capacity is
-//  assumed to be available for processing this block range.
+/// Performs explicit copying for the contiguous sequence of operations in the
+/// block iterator range [`begin', `end'), where `end' can't be past the
+/// terminator of the block (since additional operations are potentially
+/// inserted right before `end`. Returns the total size of fast memory space
+/// buffers used. `copyOptions` provides various parameters, and the output
+/// argument `copyNests` is the set of all copy nests inserted, each represented
+/// by its root affine.for. Since we generate alloc's and dealloc's for all fast
+/// buffers (before and after the range of operations resp. or at a hoisted
+/// position), all of the fast memory capacity is assumed to be available for
+/// processing this block range. When 'filterMemRef' is specified, copies are
+/// only generated for the provided MemRef.
 uint64_t mlir::affineDataCopyGenerate(Block::iterator begin,
                                       Block::iterator end,
                                       const AffineCopyOptions &copyOptions,
+                                      Optional<Value> filterMemRef,
                                       DenseSet<Operation *> &copyNests) {
   if (begin == end)
     return 0;
 
   assert(begin->getBlock() == std::prev(end)->getBlock() &&
          "Inconsistent block begin/end args");
   assert(end != end->getBlock()->end() && "end can't be the block terminator");
 
uint64_t mlir::affineDataCopyGenerate(Block::iterator begin,
 
   // To check for errors when walking the block.
   bool error = false;
 
   // Walk this range of operations  to gather all memory regions.
   block->walk(begin, end, [&](Operation *opInst) {
     // Gather regions to allocate to buffers in faster memory space.
     if (auto loadOp = dyn_cast<AffineLoadOp>(opInst)) {
-      if ((loadOp.getMemRefType().getMemorySpace() !=
+      if ((filterMemRef.hasValue() && filterMemRef != loadOp.getMemRef()) ||
+          (loadOp.getMemRefType().getMemorySpace() !=
            copyOptions.slowMemorySpace))
         return;
     } else if (auto storeOp = dyn_cast<AffineStoreOp>(opInst)) {
-      if (storeOp.getMemRefType().getMemorySpace() !=
+      if ((filterMemRef.hasValue() && filterMemRef != storeOp.getMemRef()) ||
+          storeOp.getMemRefType().getMemorySpace() !=
               copyOptions.slowMemorySpace)
         return;
     } else {
       // Neither load nor a store op.
       return;
     }
 
     // Compute the MemRefRegion accessed.
     auto region = std::make_unique<MemRefRegion>(opInst->getLoc());
uint64_t mlir::affineDataCopyGenerate(Block::iterator begin,
   if (totalCopyBuffersSizeInBytes > copyOptions.fastMemCapacityBytes) {
     StringRef str = "Total size of all copy buffers' for this block "
                     "exceeds fast memory capacity\n";
     block->getParentOp()->emitError(str);
   }
 
   return totalCopyBuffersSizeInBytes;
 }
+
+/// Gathers all AffineForOps in 'block' at 'currLoopDepth' in 'depthToLoops'.
+static void gatherLoopsInBlock(
+    Block *block, unsigned currLoopDepth,
+    DenseMap<unsigned, SmallVector<AffineForOp, 2>> &depthToLoops) {
+  auto &loopsAtDepth = depthToLoops[currLoopDepth];
+  for (auto &op : *block) {
+    if (auto forOp = dyn_cast<AffineForOp>(op)) {
+      loopsAtDepth.push_back(forOp);
+      gatherLoopsInBlock(forOp.getBody(), currLoopDepth + 1, depthToLoops);
+    }
+  }
+}
+
+/// Gathers all AffineForOps in 'func' grouped by loop depth.
+void mlir::gatherLoops(
+    FuncOp func,
+    DenseMap<unsigned, SmallVector<AffineForOp, 2>> &depthToLoops) {
+  for (auto &block : func)
+    gatherLoopsInBlock(&block, /*currLoopDepth=*/0, depthToLoops);
+}

mlir/test/Transforms/affine-data-copy.mlir

 // RUN: mlir-opt %s -split-input-file -affine-data-copy-generate -affine-data-copy-generate-dma=false -affine-data-copy-generate-fast-mem-space=0 -affine-data-copy-generate-skip-non-unit-stride-loops | FileCheck %s
 // Small buffer size to trigger fine copies.
 // RUN: mlir-opt %s -affine-data-copy-generate -affine-data-copy-generate-dma=false -affine-data-copy-generate-fast-mem-space=0 -affine-data-copy-generate-fast-mem-capacity=1 | FileCheck --check-prefix=CHECK-SMALL %s
 
+// Test affine data copy with a memref filter. We use a test pass that invokes
+// affine data copy utility on the input loop nest.
+// '-test-affine-data-copy-memref-filter' passes the first memref found in an
+// affine.load op in the innermost loop as a filter.
+// RUN: mlir-opt %s -split-input-file -test-affine-data-copy='memref-filter=1' | FileCheck %s --check-prefix=FILTER
+
 // -copy-skip-non-stride-loops forces the copies to be placed right inside the
 // tile space loops, avoiding the sensitivity of copy placement depth to memory
 // footprint -- so that one could write a definite test case and not have to
 // update it each time something related to the cost functions change.
 
 #map0 = affine_map<(d0) -> (d0)>
 #map1 = affine_map<(d0) -> (d0 + 128)>
 
 // Map used to index the original memref while copying.
 // CHECK-DAG: [[MEM_IDX_MAP:map[0-9]+]] = affine_map<(d0, d1) -> (d0 + d1)>
 // Map used to index the buffer while computing.
 // CHECK-DAG: [[BUF_IDX_MAP:map[0-9]+]] = affine_map<(d0, d1, d2, d3) -> (-d0 + d2, -d1 + d3)>
 
 // CHECK-LABEL: func @matmul
+// FILTER-LABEL: func @matmul
 func @matmul(%A: memref<4096x4096xf32>, %B: memref<4096x4096xf32>, %C: memref<4096x4096xf32>) -> memref<4096x4096xf32> {
   affine.for %i = 0 to 4096 step 128 {
     affine.for %j = 0 to 4096 step 128 {
       affine.for %k = 0 to 4096 step 128 {
         affine.for %ii = #map0(%i) to #map1(%i) {
           affine.for %jj = #map0(%j) to #map1(%j) {
             affine.for %kk = #map0(%k) to #map1(%k) {
               %5 = affine.load %A[%ii, %kk] : memref<4096x4096xf32>

 // CHECK:         [[BUFA]] = affine.load [[BUFC]][%{{.*}}, %{{.*}}] : memref<128x128xf32>
 // CHECK:         store [[BUFA]], %{{.*}}[%{{.*}}, %{{.*}}] : memref<4096x4096xf32>
 // CHECK:       }
 // CHECK:     }
 // CHECK:     dealloc [[BUFC]] : memref<128x128xf32>
 // CHECK:   }
 // CHECK: }
 
+// Check that only one memref is copied when memref filter is used.
+
+//      FILTER: affine.for %{{.*}} = 0 to 4096 step 128 {
+//      FILTER:   alloc() : memref<128x4096xf32>
+//  FILTER-NOT:   alloc()
+//      FILTER:   affine.for %{{.*}} = 0 to 128 {
+//      FILTER:     affine.for %{{.*}} = 0 to 4096 {
+//      FILTER:     affine.for %{{.*}} = 0 to 4096 step 128 {
+// FILTER-NEXT:       affine.for %{{.*}} = 0 to 4096 step 128 {
+// FILTER-NEXT:         affine.for %{{.*}} = #map{{.*}}(%{{.*}}) to #map{{.*}}(%{{.*}}) {
+// FILTER-NEXT:           affine.for %{{.*}} = #map{{.*}}(%{{.*}}) to #map{{.*}}(%{{.*}}) {
+// FILTER-NEXT:             affine.for %{{.*}} = #map{{.*}}(%{{.*}}) to #map{{.*}}(%{{.*}}) {
+//      FILTER:   dealloc %1 : memref<128x4096xf32>
+//  FILTER-NOT:   dealloc %1 : memref<128x4096xf32>
+
+// -----
+
 //
 // This test case will lead to single element buffers. These are eventually
 // expected to be turned into registers via alloca and mem2reg.
 //
-// CHECK-SMALL: func @foo
+// CHECK-SMALL-LABEL: func @foo
+// FILTER-LABEL: func @foo
 func @foo(%arg0: memref<1024x1024xf32>, %arg1: memref<1024x1024xf32>, %arg2: memref<1024x1024xf32>) -> memref<1024x1024xf32> {
   affine.for %i = 0 to 1024 {
     affine.for %j = 0 to 1024 {
       affine.for %k = 0 to 1024 {
         %6 = affine.load %arg1[%k, %j] : memref<1024x1024xf32>
         %7 = affine.load %arg2[%i, %j] : memref<1024x1024xf32>
         %9 = addf %6, %7 : f32
         affine.store %9, %arg2[%i, %j] : memref<1024x1024xf32>

 // CHECK-SMALL:     %{{.*}} = affine.apply #map{{.*}}(%arg{{.*}}, %c0{{.*}})
 // CHECK-SMALL:     %{{.*}} = affine.apply #map{{.*}}(%arg{{.*}}, %c0{{.*}})
 // CHECK-SMALL:     %{{.*}} = affine.load %{{.*}}[%c0{{.*}}, %c0{{.*}}] : memref<1x1xf32>
 // CHECK-SMALL:     affine.store %{{.*}}, %arg{{.*}}[%{{.*}}, %{{.*}}] : memref<1024x1024xf32>
 // CHECK-SMALL:     dealloc %{{.*}} : memref<1x1xf32>
 // CHECK-SMALL:   }
 // CHECK-SMALL: }
 // CHECK-SMALL: return
+
+// Check that only one memref is copied when memref filter is used.
+
+//      FILTER: alloc() : memref<1024x1024xf32>
+//  FILTER-NOT: alloc()
+//      FILTER: affine.for %{{.*}} = 0 to 1024 {
+//      FILTER:   affine.for %{{.*}} = 0 to 1024 {
+//      FILTER: affine.for %{{.*}} = 0 to 1024 {
+// FILTER-NEXT:   affine.for %{{.*}} = 0 to 1024 {
+// FILTER-NEXT:     affine.for %{{.*}} = 0 to 1024 {
+//      FILTER: dealloc %{{.*}} : memref<1024x1024xf32>
+//  FILTER-NOT: dealloc

mlir/test/Transforms/dma-generate.mlir

func @test_analysis_util(%arg0: memref<4x4x16x1xf32>, %arg1: memref<144x9xf32>, %arg2: memref<2xf32>) -> (memref<144x9xf32>, memref<2xf32>) {
   return %arg1, %arg2 : memref<144x9xf32>, memref<2xf32>
 }
 // CHECK:       affine.for %{{.*}} = 0 to 9 step 3 {
 // CHECK:         [[BUF:%[0-9]+]] = alloc() : memref<2xf32, 2>
 // CHECK:         affine.dma_start %{{.*}}[%{{.*}} floordiv 8], [[BUF]]
 // CHECK:         affine.dma_wait %{{.*}}[%{{.*}}], %{{.*}} : memref<1xi32>
 // CHECK:         affine.for %{{.*}} =
 
-// ----
+// -----
 
 #map3 = affine_map<(d0) -> (d0)>
 #map12 = affine_map<(d0) -> (d0 + 3)>
 #map14 = affine_map<(d0, d1) -> ((d0 + d1 * 72) floordiv 2304 + ((((d0 + d1 * 72) mod 2304) mod 1152) mod 9) floordiv 3)>
 #map15 = affine_map<(d0, d1) -> ((d0 + d1 * 72) mod 2304 - (((d0 + d1 * 72) mod 2304) floordiv 1152) * 1151 - ((((d0 + d1 * 72) mod 2304) mod 1152) floordiv 9) * 9 - (((((d0 + d1 * 72) mod 2304) mod 1152) mod 9) floordiv 3) * 3)>
 #map16 = affine_map<(d0, d1) -> (((((d0 + d1 * 72) mod 2304) mod 1152) floordiv 9) floordiv 8)>
 // Test for test case in b/128303048 #4.
+// CHECK-LABEL: func @test_memref_bounds
 func @test_memref_bounds(%arg0: memref<4x4x16x1xvector<8x128xf32>>, %arg1: memref<144x9xvector<8x128xf32>>, %arg2: memref<2xvector<8x128xf32>>) -> (memref<144x9xvector<8x128xf32>>, memref<2xvector<8x128xf32>>) {
   %c0 = constant 0 : index
   affine.for %i8 = 0 to 9 step 3 {
     affine.for %i9 = #map3(%i8) to #map12(%i8) {
       affine.for %i10 = 0 to 64 {
         %10 = affine.apply #map14(%i9, %i10)
         %11 = affine.apply #map15(%i9, %i10)
         %12 = affine.apply #map16(%i9, %i10)

mlir/test/lib/Transforms/CMakeLists.txt

 add_llvm_library(MLIRTestTransforms
+  TestAffineDataCopy.cpp
   TestAllReduceLowering.cpp
   TestCallGraph.cpp
   TestConstantFold.cpp
   TestLoopFusion.cpp
   TestGpuMemoryPromotion.cpp
   TestInlining.cpp
   TestLinalgTransforms.cpp
   TestLiveness.cpp

mlir/test/lib/Transforms/TestAffineDataCopy.cpp

+//===- TestAffineDataCopy.cpp - Test affine data copy utility -------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// This file implements a pass to test affine data copy utility functions and
+// options.
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Analysis/Passes.h"
+#include "mlir/Dialect/AffineOps/AffineOps.h"
+#include "mlir/Pass/Pass.h"
+#include "mlir/Transforms/LoopUtils.h"
+#include "mlir/Transforms/Passes.h"
+
+#define PASS_NAME "test-affine-data-copy"
+
+using namespace mlir;
+
+static llvm::cl::OptionCategory clOptionsCategory(PASS_NAME " options");
+
+namespace {
+
+struct TestAffineDataCopy : public FunctionPass<TestAffineDataCopy> {
+  TestAffineDataCopy() = default;
+  TestAffineDataCopy(const TestAffineDataCopy &pass){};
+
+  void runOnFunction() override;
+
+private:
+  Option<bool> clMemRefFilter{
+      *this, "memref-filter",
+      llvm::cl::desc(
+          "Enable memref filter testing in affine data copy optimization"),
+      llvm::cl::init(false)};
+};
+
+} // end anonymous namespace
+
+void TestAffineDataCopy::runOnFunction() {
+  // Gather all AffineForOps by loop depth.
+  DenseMap<unsigned, SmallVector<AffineForOp, 2>> depthToLoops;
+  gatherLoops(getFunction(), depthToLoops);
+  assert(depthToLoops.size() && "Loop nest not found");
+
+  // Only support tests with a single loop nest and a single innermost loop
+  // for now.
+  unsigned innermostLoopIdx = depthToLoops.size() - 2;
+  if (depthToLoops[0].size() != 1 || depthToLoops[innermostLoopIdx].size() != 1)
+    return;
+
+  auto loopNest = depthToLoops[0][0];
+  auto innermostLoop = depthToLoops[innermostLoopIdx][0];
+  Optional<Value> memrefFilter;
+  if (clMemRefFilter) {
+    // Gather MemRef filter. For simplicity, we use the first loaded memref
+    // found in the innermost loop.
+    for (auto &op : *innermostLoop.getBody()) {
+      if (auto load = dyn_cast<AffineLoadOp>(op)) {
+        memrefFilter = load.getMemRef();
+        break;
+      }
+    }
+  }
+
+  AffineCopyOptions copyOptions = {/*generateDma=*/false,
+                                   /*slowMemorySpace=*/0,
+                                   /*fastMemorySpace=*/0,
+                                   /*tagMemorySpace=*/0,
+                                   /*fastMemCapacityBytes=*/32 * 1024 * 1024UL};
+  DenseSet<Operation *> copyNests;
+  affineDataCopyGenerate(loopNest.getBody()->begin(),
+                         std::prev(loopNest.getBody()->end()), copyOptions,
+                         memrefFilter, copyNests);
+}
+
+namespace mlir {
+void registerTestAffineDataCopyPass() {
+  PassRegistration<TestAffineDataCopy>(
+      PASS_NAME, "Tests affine data copy utility functions.");
+}
+} // namespace mlir

mlir/test/lib/Transforms/TestLoopFusion.cpp

 #include "mlir/Analysis/AffineStructures.h"
 #include "mlir/Analysis/Passes.h"
 #include "mlir/Analysis/Utils.h"
 #include "mlir/Dialect/AffineOps/AffineOps.h"
 #include "mlir/Dialect/StandardOps/Ops.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/Pass/Pass.h"
 #include "mlir/Transforms/LoopFusionUtils.h"
+#include "mlir/Transforms/LoopUtils.h"
 #include "mlir/Transforms/Passes.h"
 
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/Debug.h"
 
 #define DEBUG_TYPE "test-loop-fusion"
 

 namespace {
 
 struct TestLoopFusion : public FunctionPass<TestLoopFusion> {
   void runOnFunction() override;
 };
 
 } // end anonymous namespace
 
-// Gathers all AffineForOps in 'block' at 'currLoopDepth' in 'depthToLoops'.
-static void
-gatherLoops(Block *block, unsigned currLoopDepth,
-            DenseMap<unsigned, SmallVector<AffineForOp, 2>> &depthToLoops) {
-  auto &loopsAtDepth = depthToLoops[currLoopDepth];
-  for (auto &op : *block) {
-    if (auto forOp = dyn_cast<AffineForOp>(op)) {
-      loopsAtDepth.push_back(forOp);
-      gatherLoops(forOp.getBody(), currLoopDepth + 1, depthToLoops);
-    }
-  }
-}
-
 // Run fusion dependence check on 'loops[i]' and 'loops[j]' at loop depths
 // in range ['loopDepth' + 1, 'maxLoopDepth'].
 // Emits a remark on 'loops[i]' if a fusion-preventing dependence exists.
 // Returns false as IR is not transformed.
 static bool testDependenceCheck(AffineForOp srcForOp, AffineForOp dstForOp,
                                 unsigned i, unsigned j, unsigned loopDepth,
                                 unsigned maxLoopDepth) {
   mlir::ComputationSliceState sliceUnion;

 
 void TestLoopFusion::runOnFunction() {
   DenseMap<unsigned, SmallVector<AffineForOp, 2>> depthToLoops;
   if (clTestLoopFusionTransformation) {
     // Run loop fusion until a fixed point is reached.
     do {
       depthToLoops.clear();
       // Gather all AffineForOps by loop depth.
-      for (auto &block : getFunction())
-        gatherLoops(&block, /*currLoopDepth=*/0, depthToLoops);
+      gatherLoops(getFunction(), depthToLoops);
 
       // Try to fuse all combinations of src/dst loop nests in 'depthToLoops'.
     } while (iterateLoops(depthToLoops, testLoopFusionTransformation,
                           /*return_on_change=*/true));
     return;
   }
 
   // Gather all AffineForOps by loop depth.
-  for (Block &block : getFunction())
-    gatherLoops(&block, /*currLoopDepth=*/0, depthToLoops);
+  gatherLoops(getFunction(), depthToLoops);
 
   // Run tests on all combinations of src/dst loop nests in 'depthToLoops'.
   if (clTestDependenceCheck)
     iterateLoops(depthToLoops, testDependenceCheck);
   if (clTestSliceComputation)
     iterateLoops(depthToLoops, testSliceComputation);
 }
 
 namespace mlir {
 void registerTestLoopFusion() {
   PassRegistration<TestLoopFusion>("test-loop-fusion",
                                    "Tests loop fusion utility functions.");
 }
 } // namespace mlir

mlir/tools/mlir-opt/mlir-opt.cpp

 // Defined in the test directory, no public header.
 void registerConvertToTargetEnvPass();
 void registerInliner();
 void registerMemRefBoundCheck();
 void registerPassManagerTestPass();
 void registerPatternsTestPass();
 void registerSimpleParametricTilingPass();
 void registerSymbolTestPasses();
+void registerTestAffineDataCopyPass();
 void registerTestAllReduceLoweringPass();
 void registerTestCallGraphPass();
 void registerTestConstantFold();
 void registerTestFunc();
 void registerTestGpuMemoryPromotionPass();
 void registerTestLinalgTransforms();
 void registerTestLivenessPass();
 void registerTestLoopFusion();

 void registerTestPasses() {
   registerConvertToTargetEnvPass();
   registerInliner();
   registerMemRefBoundCheck();
   registerPassManagerTestPass();
   registerPatternsTestPass();
   registerSimpleParametricTilingPass();
   registerSymbolTestPasses();
+  registerTestAffineDataCopyPass();
   registerTestAllReduceLoweringPass();
   registerTestCallGraphPass();
   registerTestConstantFold();
   registerTestFunc();
   registerTestGpuMemoryPromotionPass();
   registerTestLinalgTransforms();
   registerTestLivenessPass();
   registerTestLoopFusion();