[mlir] Add a simplifying wrapper for generateCopy and expose it.

[timshen]

Summary

affineDataCopyGenerate is a monolithinc function that
combines several steps for good reasons, but it makes customizing
the behaivor even harder. The major two steps by affineDataCopyGenerate are:
a) Identify interesting memrefs and collect their uses.
b) Create new buffers to forward these uses.

Step (a) actually has requires tremendous customization options. One could see
that from the recently added filterMemRef parameter.

This patch adds a function that only does (b), in the hope that (a)
can be directly implemented by the callers. In fact, (a) is quite
simple if the caller has only one buffer to consider, or even one use.

[timshen]

Update comments.

[timshen]

I'm going to add some tests.

[rriddle]

mlir/include/mlir/Transforms/LoopUtils.h

18

Can you can forward declare the necessary things and move this header to the .cpp file?

mlir/lib/Transforms/Utils/LoopUtils.cpp

1812

Please drop trivial braces.

1820

Why not rework this?

assert(copyNests.size() <= 1 && ...);
result.copy_nest = copyNests.empty() ? nullptr : *copyNests.begin();

[timshen]

Add a test.

[timshen]

Address comments.

[timshen]

Address comments.

[bondhugula]

This looks like a good refactoring, thanks! Some (mostly minor) comments.

mlir/include/mlir/Transforms/LoopUtils.h

189

This doc comment should moved down to the method declaration below. Please mention that this performs data copy generation for a single memref region associated with a single op.

191

relevant

191

push -> pushed

198

Nit: Notice -> Note that

203

Doc comment.

207

A better name for 'where'?

mlir/lib/Transforms/Utils/LoopUtils.cpp

1818

"At most one copy nest expected" ?

mlir/test/lib/Transforms/TestAffineDataCopy.cpp

42

Test copy generation for a single op / memref region?

[timshen]

Address comments.

[bondhugula]

mlir/include/mlir/Transforms/LoopUtils.h

189

Doc comment here please (even if it's brief).

194

buffer -> buffer allocation

196

alloc above is an operation. alloc -> allocated buffer?

201

generateCopyFromMemRefRegion -> generateCopyForMemRegion ?

215

Camel case.

215

insertionPoint would be used typically for iterator types. Shouldn't this just be a Block::iterator? At the call site, you could pass Block::iterator(op).

mlir/lib/Transforms/Utils/LoopUtils.cpp

1803–1804

Block::iterator(insertion_point) would have given you the iterator. But you don't need this anyway if you pass a Block::iterator itself to this method.

[bondhugula]

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

10

Do you need the '=1'? It could be confusing unless it's '=true'.

mlir/test/lib/Transforms/TestAffineDataCopy.cpp

85

/*loopDepth = */0

[timshen]

Update documentation.

[bondhugula]

mlir/include/mlir/Transforms/LoopUtils.h

201

This doc comment appears to provides detail before the actual information. Perhaps just rephrase to:

/// Similar to affineDataCopyGenerate, but works on a single memref region.

You could move "The logic of "find ..." as a comment on the definition of the method.

206

-> a single memref region denoted ...

212

Nit: isn't -> aren't

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

143

// MEMREF-REGION-LABEL

here to be safe.

202

Perhaps a comment here on what is being checked for.

[timshen]

Address all comments.

mlir/include/mlir/Transforms/LoopUtils.h

201

Simplified further. The "The logic of" part looks like design rationales rather than function descriptions, which is repeating after the API itself.

215

This op is more than an iterator insertion point. It's also the op that's analyzed by the memref region. It also indicate that something next to the iterator is going to change (std::next(insertionPoint)). Change the name to analyzedOp, and adjusted comments for that.

mlir/lib/Transforms/Utils/LoopUtils.cpp

1803–1804

N/A.

[bondhugula]

Looks good!

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/Analysis/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;

 
 /// A convenience version of affineDataCopyGenerate for all ops in the body of
 /// an AffineForOp.
 uint64_t affineDataCopyGenerate(AffineForOp forOp,
                                 const AffineCopyOptions &copyOptions,
                                 Optional<Value> filterMemRef,
                                 DenseSet<Operation *> &copyNests);
 
+/// generateDataCopyAroundOp is similar to affineDataCopyGenerate, but with some
+/// simplifications:
+/// * The logic of "find relavant memrefs and their uses" is de-coupled and push
+/// back to the users. It focuses on generating fast buffers and associated
+/// loops/DMAs.
+/// * It handles a single memref per call.
+/// * The prologue and epilogue always surround `op`, not in potentially
+/// arbitrary places.
+///
+/// Notice that certain options in `copyOptions` isn't looked at anymore, like
+/// slowMemorySpace.
+struct CopyGenerateResult {
+  uint64_t sizeInBytes;
+  Operation *alloc;
+  Operation *copyNest;
+};
+
+LogicalResult generateDataCopyAroundOp(const MemRefRegion &memrefRegion,
+                                       Operation *where,
+                                       const AffineCopyOptions &copyOptions,
+                                       CopyGenerateResult &result);
+
 /// 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);
 
 /// Replace a perfect nest of "for" loops with a single linearized loop. Assumes
 /// `loops` contains a list of perfectly nested loops with bounds and steps
 /// independent of any loop induction variable involved in the nest.
 void coalesceLoops(MutableArrayRef<loop::ForOp> loops);
 
 /// Maps `forOp` for execution on a parallel grid of virtual `processorIds` of
 /// size given by `numProcessors`. This is achieved by embedding the SSA values

mlir/lib/Transforms/Utils/LoopUtils.cpp

uint64_t mlir::affineDataCopyGenerate(AffineForOp forOp,
                                       const AffineCopyOptions &copyOptions,
                                       Optional<Value> filterMemRef,
                                       DenseSet<Operation *> &copyNests) {
   return affineDataCopyGenerate(forOp.getBody()->begin(),
                                 std::prev(forOp.getBody()->end()), copyOptions,
                                 filterMemRef, copyNests);
 }
 
+LogicalResult mlir::generateDataCopyAroundOp(
+    const MemRefRegion &memrefRegion, Operation *where,
+    const AffineCopyOptions &copyOptions, CopyGenerateResult &result) {
+  Block *block = where->getBlock();
+  auto begin = where->getIterator();
+  auto end = std::next(begin);
+  DenseMap<Value, Value> fastBufferMap;
+  DenseSet<Operation *> copyNests;
+
+  auto err = generateCopy(memrefRegion, block, begin, end, block, begin, end,
+                          copyOptions, fastBufferMap, copyNests,
+                          &result.sizeInBytes, &begin, &end);
+  if (failed(err)) {
+    return err;
+  }
+  result.alloc =
+      fastBufferMap.find(memrefRegion.memref)->second.getDefiningOp();
+  if (copyNests.empty()) {
+    result.copyNest = nullptr;
+  } else {
+    assert(copyNests.size() == 1 &&
+           "Multiple copy nests generated appear for a single memref.");
+    result.copyNest = *copyNests.begin();
+  }
+  return success();
+}
+
 /// Gathers all AffineForOps in 'block' at 'currLoopDepth' in 'depthToLoops'.
 static void
 gatherLoopsInBlock(Block *block, unsigned currLoopDepth,
                    std::vector<SmallVector<AffineForOp, 2>> &depthToLoops) {
   // Add a new empty level to output if it doesn't exist level already.
   assert(currLoopDepth <= depthToLoops.size() && "Unexpected currLoopDepth");
   if (currLoopDepth == depthToLoops.size())
     depthToLoops.push_back(SmallVector<AffineForOp, 2>());

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
+// RUN: mlir-opt %s -split-input-file -test-affine-data-copy='test-generate-data-copy-around-op=1' | FileCheck %s --check-prefix=COPY_AROUND_OP
 
 // -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)>

 
 // -----
 
 //
 // This test case will lead to single element buffers. These are eventually
 // expected to be turned into registers via alloca and mem2reg.
 //
 // 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>

 //  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
+
+//      COPY_AROUND_OP: alloc() : memref<1024x1024xf32>
+//  COPY_AROUND_OP-NOT: alloc()
+//      COPY_AROUND_OP: affine.for %{{.*}} = 0 to 1024 {
+//      COPY_AROUND_OP:   affine.for %{{.*}} = 0 to 1024 {
+//      COPY_AROUND_OP: affine.for %{{.*}} = 0 to 1024 {
+// COPY_AROUND_OP-NEXT:   affine.for %{{.*}} = 0 to 1024 {
+// COPY_AROUND_OP-NEXT:     affine.for %{{.*}} = 0 to 1024 {
+//      COPY_AROUND_OP: dealloc %{{.*}} : memref<1024x1024xf32>
+//  COPY_AROUND_OP-NOT: dealloc

mlir/test/lib/Transforms/TestAffineDataCopy.cpp

struct TestAffineDataCopy : public FunctionPass<TestAffineDataCopy> {
   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)};
+  Option<bool> clTestGenerateDataCopyAroundOp{
+      *this, "test-generate-data-copy-around-op",
+      llvm::cl::desc("Test through generateCopyAroundOp"),
+      llvm::cl::init(false)};
 };
 
 } // end anonymous namespace
 
 void TestAffineDataCopy::runOnFunction() {
   // Gather all AffineForOps by loop depth.
   std::vector<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() - 1;
   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) {
+  AffineLoadOp load;
+  if (clMemRefFilter || clTestGenerateDataCopyAroundOp) {
     // 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();
+      if (auto ld = dyn_cast<AffineLoadOp>(op)) {
+        load = ld;
         break;
       }
     }
   }
 
   AffineCopyOptions copyOptions = {/*generateDma=*/false,
                                    /*slowMemorySpace=*/0,
                                    /*fastMemorySpace=*/0,
                                    /*tagMemorySpace=*/0,
                                    /*fastMemCapacityBytes=*/32 * 1024 * 1024UL};
+  if (clMemRefFilter) {
     DenseSet<Operation *> copyNests;
-  affineDataCopyGenerate(loopNest, copyOptions, memrefFilter, copyNests);
+    affineDataCopyGenerate(loopNest, copyOptions, load.getMemRef(), copyNests);
+  } else if (clTestGenerateDataCopyAroundOp) {
+    CopyGenerateResult result;
+    MemRefRegion region(loopNest.getLoc());
+    region.compute(load, 0);
+    generateDataCopyAroundOp(region, loopNest, copyOptions, result);
+  }
 }
 
 namespace mlir {
 void registerTestAffineDataCopyPass() {
   PassRegistration<TestAffineDataCopy>(
       PASS_NAME, "Tests affine data copy utility functions.");
 }
 } // namespace mlir