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

[Tim Shen via Phabricator via llvm-commits]

timshen updated this revision to Diff 249514.
timshen added a comment.

Update comments.


Repository:
  rG LLVM Github Monorepo

CHANGES SINCE LAST ACTION
  https://reviews.llvm.org/D75965/new/

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/9071d91c/attachment.bin>