[Mlir-commits] [mlir] 82e9160 - [MLIR][Affine] NFC: add convenience method for affine data copy for a loop body

[Mehdi Amini]

Author: Uday Bondhugula
Date: 2020-03-09T04:23:54Z
New Revision: 82e9160aabbb19cde1c6696331ebf0b8981ed494

URL: https://github.com/llvm/llvm-project/commit/82e9160aabbb19cde1c6696331ebf0b8981ed494
DIFF: https://github.com/llvm/llvm-project/commit/82e9160aabbb19cde1c6696331ebf0b8981ed494.diff

LOG: [MLIR][Affine] NFC: add convenience method for affine data copy for a loop body

add convenience method for affine data copy generation for a loop body

Signed-off-by: Uday Bondhugula <uday at polymagelabs.com>

Differential Revision: https://reviews.llvm.org/D75822

Added: 
    

Modified: 
    mlir/include/mlir/Transforms/LoopUtils.h
    mlir/lib/Transforms/Utils/LoopUtils.cpp
    mlir/test/lib/Transforms/TestAffineDataCopy.cpp

Removed: 
    


################################################################################
diff  --git a/mlir/include/mlir/Transforms/LoopUtils.h b/mlir/include/mlir/Transforms/LoopUtils.h
index 4748e13cf5eb..a82a06541a13 100644
--- a/mlir/include/mlir/Transforms/LoopUtils.h
+++ b/mlir/include/mlir/Transforms/LoopUtils.h
@@ -178,6 +178,13 @@ uint64_t affineDataCopyGenerate(Block::iterator begin, Block::iterator end,
                                 Optional<Value> filterMemRef,
                                 DenseSet<Operation *> &copyNests);
 
+/// 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);
+
 /// 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`.

diff  --git a/mlir/lib/Transforms/Utils/LoopUtils.cpp b/mlir/lib/Transforms/Utils/LoopUtils.cpp
index a8825c11bb12..c2cd2333f531 100644
--- a/mlir/lib/Transforms/Utils/LoopUtils.cpp
+++ b/mlir/lib/Transforms/Utils/LoopUtils.cpp
@@ -1784,6 +1784,17 @@ uint64_t mlir::affineDataCopyGenerate(Block::iterator begin,
   return totalCopyBuffersSizeInBytes;
 }
 
+// A convenience version of affineDataCopyGenerate for all ops in the body of
+// an AffineForOp.
+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);
+}
+
 /// Gathers all AffineForOps in 'block' at 'currLoopDepth' in 'depthToLoops'.
 static void
 gatherLoopsInBlock(Block *block, unsigned currLoopDepth,

diff  --git a/mlir/test/lib/Transforms/TestAffineDataCopy.cpp b/mlir/test/lib/Transforms/TestAffineDataCopy.cpp
index f5caa6922710..de7cdbd51b9d 100644
--- a/mlir/test/lib/Transforms/TestAffineDataCopy.cpp
+++ b/mlir/test/lib/Transforms/TestAffineDataCopy.cpp
@@ -73,9 +73,7 @@ void TestAffineDataCopy::runOnFunction() {
                                    /*tagMemorySpace=*/0,
                                    /*fastMemCapacityBytes=*/32 * 1024 * 1024UL};
   DenseSet<Operation *> copyNests;
-  affineDataCopyGenerate(loopNest.getBody()->begin(),
-                         std::prev(loopNest.getBody()->end()), copyOptions,
-                         memrefFilter, copyNests);
+  affineDataCopyGenerate(loopNest, copyOptions, memrefFilter, copyNests);
 }
 
 namespace mlir {