[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

43 ↗

(On Diff #249528)

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 ↗

(On Diff #249710)

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

mlir/test/lib/Transforms/TestAffineDataCopy.cpp

86 ↗

(On Diff #249710)

/*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 ↗

(On Diff #249710)

// MEMREF-REGION-LABEL

here to be safe.

202 ↗

(On Diff #249710)

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 = op->getBlock();
+  auto begin = op->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>());