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

[Tim Shen via Phabricator via llvm-commits]

timshen created this revision.
timshen added reviewers: bondhugula, ftynse.
Herald added subscribers: llvm-commits, Joonsoo, liufengdb, aartbik, lucyrfox, mgester, arpith-jacob, nicolasvasilache, antiagainst, shauheen, burmako, jpienaar, rriddle, mehdi_amini, sanjoy.google, bixia.
Herald added a project: LLVM.
timshen updated this revision to Diff 249514.
timshen added a comment.

Update comments.


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.


Repository:
  rG LLVM Github Monorepo

https://reviews.llvm.org/D75965

Files:
  mlir/include/mlir/Transforms/LoopUtils.h
  mlir/lib/Transforms/Utils/LoopUtils.cpp


Index: mlir/lib/Transforms/Utils/LoopUtils.cpp
===================================================================
--- mlir/lib/Transforms/Utils/LoopUtils.cpp
+++ mlir/lib/Transforms/Utils/LoopUtils.cpp
@@ -1797,6 +1797,33 @@
                                 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,
Index: mlir/include/mlir/Transforms/LoopUtils.h
===================================================================
--- mlir/include/mlir/Transforms/LoopUtils.h
+++ mlir/include/mlir/Transforms/LoopUtils.h
@@ -15,6 +15,7 @@
 #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"
@@ -185,6 +186,28 @@
                                 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`.


-------------- next part --------------
A non-text attachment was scrubbed...
Name: D75965.249514.patch
Type: text/x-patch
Size: 3088 bytes
Desc: not available
URL: <http://lists.llvm.org/pipermail/llvm-commits/attachments/20200310/0f804306/attachment.bin>