From 82e9160aabbb19cde1c6696331ebf0b8981ed494 Mon Sep 17 00:00:00 2001
From: Uday Bondhugula <uday@polymagelabs.com>
Date: Mon, 9 Mar 2020 04:21:40 +0000
Subject: [PATCH] [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@polymagelabs.com>

Differential Revision: https://reviews.llvm.org/D75822
---
 mlir/include/mlir/Transforms/LoopUtils.h        |  7 +++++++
 mlir/lib/Transforms/Utils/LoopUtils.cpp         | 11 +++++++++++
 mlir/test/lib/Transforms/TestAffineDataCopy.cpp |  4 +---
 3 files changed, 19 insertions(+), 3 deletions(-)

diff --git a/mlir/include/mlir/Transforms/LoopUtils.h b/mlir/include/mlir/Transforms/LoopUtils.h
index 4748e13..a82a065 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 a8825c1..c2cd233 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 f5caa69..de7cdbd 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 {
-- 
2.7.4