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

[bondhugula]

Summary

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

Signed-off-by: Uday Bondhugula <[e-mail]>

[bondhugula]

@mehdi_amini could you please commit this for me? Thanks.

mlir/include/mlir/Transforms/LoopUtils.h

 /// position), all of the fast memory capacity is assumed to be available for
 /// processing this block range. When 'filterMemRef' is specified, copies are
 /// only generated for the provided MemRef.
 uint64_t affineDataCopyGenerate(Block::iterator begin, Block::iterator end,
                                 const AffineCopyOptions &copyOptions,
                                 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`.
 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

mlir/lib/Transforms/Utils/LoopUtils.cpp

if (totalCopyBuffersSizeInBytes > copyOptions.fastMemCapacityBytes) {
     StringRef str = "Total size of all copy buffers' for this block "
                     "exceeds fast memory capacity\n";
     block->getParentOp()->emitError(str);
   }
 
   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,
                    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>());

mlir/test/lib/Transforms/TestAffineDataCopy.cpp

void TestAffineDataCopy::runOnFunction() {
   }
 
   AffineCopyOptions copyOptions = {/*generateDma=*/false,
                                    /*slowMemorySpace=*/0,
                                    /*fastMemorySpace=*/0,
                                    /*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 {
 void registerTestAffineDataCopyPass() {
   PassRegistration<TestAffineDataCopy>(
       PASS_NAME, "Tests affine data copy utility functions.");
 }
 } // namespace mlir