Create private exclusive / single use affine computation slice for an op stmt.

- add util to create a private / exclusive / single use affine
  computation slice for an op stmt (see method doc comment); a single
  multi-result affine_apply op is prepended to the op stmt to provide all
  results needed for its operands as a function of loop iterators and symbols.
- use it for DMA pipelining (to create private slices for DMA start stmt's);
  resolve TODOs/feature request (b/117159533)
- move createComposedAffineApplyOp to Transforms/Utils; free it from taking a
  memref as input / generalize it.

PiperOrigin-RevId: 216926818

1 files changed, 7 insertions, 46 deletions

diff --git a/mlir/lib/Transforms/ComposeAffineMaps.cpp b/mlir/lib/Transforms/ComposeAffineMaps.cpp
index 0aa593202e1..9e44dab1bff 100644
--- a/mlir/lib/Transforms/ComposeAffineMaps.cpp
+++ b/mlir/lib/Transforms/ComposeAffineMaps.cpp
@@ -66,6 +66,7 @@
 #include "mlir/StandardOps/StandardOps.h"
 #include "mlir/Transforms/Pass.h"
 #include "mlir/Transforms/Passes.h"
+#include "mlir/Transforms/Utils.h"
 #include "llvm/Support/CommandLine.h"
 
 using namespace mlir;
@@ -90,46 +91,6 @@ MLFunctionPass *mlir::createComposeAffineMapsPass() {
   return new ComposeAffineMaps();
 }
 
-// Creates and inserts into 'builder' a new AffineApplyOp with the number of
-// results equal to the rank of 'memrefType'. The AffineApplyOp is composed
-// with all other AffineApplyOps reachable from input paramter 'operands'.
-// The final results of the composed AffineApplyOp are returned in output
-// paramter 'results'.
-static void createComposedAffineApplyOp(
-    MLFuncBuilder *builder, Location *loc, MemRefType *memrefType,
-    const SmallVector<MLValue *, 4> &indices,
-    const SmallVector<OperationStmt *, 4> &affineApplyOps,
-    SmallVector<SSAValue *, 4> *results) {
-  // Get rank of memref type.
-  unsigned rank = memrefType->getRank();
-  assert(indices.size() == rank);
-  // Create identity map with same number of dimensions as 'memrefType'.
-  auto map = builder->getMultiDimIdentityMap(rank);
-  // Initialize AffineValueMap with identity map.
-  AffineValueMap valueMap(map, indices);
-
-  for (auto *opStmt : affineApplyOps) {
-    assert(opStmt->is<AffineApplyOp>());
-    auto affineApplyOp = opStmt->getAs<AffineApplyOp>();
-    // Forward substitute 'affineApplyOp' into 'valueMap'.
-    valueMap.forwardSubstitute(*affineApplyOp);
-  }
-  // Compose affine maps from all ancestor AffineApplyOps.
-  // Create new AffineApplyOp from 'valueMap'.
-  unsigned numOperands = valueMap.getNumOperands();
-  SmallVector<SSAValue *, 4> operands(numOperands);
-  for (unsigned i = 0; i < numOperands; ++i) {
-    operands[i] = valueMap.getOperand(i);
-  }
-  // Create new AffineApplyOp based on 'valueMap'.
-  auto affineApplyOp =
-      builder->create<AffineApplyOp>(loc, valueMap.getAffineMap(), operands);
-  results->resize(rank);
-  for (unsigned i = 0; i < rank; ++i) {
-    (*results)[i] = affineApplyOp->getResult(i);
-  }
-}
-
 PassResult ComposeAffineMaps::runOnMLFunction(MLFunction *f) {
   // Gather all loads, stores and affine apply ops.
   struct OpGatherer : public StmtWalker<OpGatherer> {
@@ -170,14 +131,14 @@ PassResult ComposeAffineMaps::runOnMLFunction(MLFunction *f) {
 
     // Gather sequnce of AffineApplyOps reachable from 'indices'.
     SmallVector<OperationStmt *, 4> affineApplyOps;
-    getReachableAffineApplyOps(indices, &affineApplyOps);
+    getReachableAffineApplyOps(indices, affineApplyOps);
     // Skip transforming 'loadOp' if there are no affine maps to compose.
     if (affineApplyOps.size() <= 1)
       continue;
 
     SmallVector<SSAValue *, 4> results;
-    createComposedAffineApplyOp(&builder, opStmt->getLoc(), memrefType, indices,
-                                affineApplyOps, &results);
+    createComposedAffineApplyOp(&builder, opStmt->getLoc(), indices,
+                                affineApplyOps, results);
     // Create new LoadOp with new affine apply op.
     auto *newLoadResult =
         builder.create<LoadOp>(opStmt->getLoc(), loadOp->getMemRef(), results)
@@ -203,14 +164,14 @@ PassResult ComposeAffineMaps::runOnMLFunction(MLFunction *f) {
     }
     // Gather sequnce of AffineApplyOps reachable from 'indices'.
     SmallVector<OperationStmt *, 4> affineApplyOps;
-    getReachableAffineApplyOps(indices, &affineApplyOps);
+    getReachableAffineApplyOps(indices, affineApplyOps);
     // Skip transforming 'storeOp' if there are no affine maps to compose.
     if (affineApplyOps.size() <= 1)
       continue;
 
     SmallVector<SSAValue *, 4> results;
-    createComposedAffineApplyOp(&builder, opStmt->getLoc(), memrefType, indices,
-                                affineApplyOps, &results);
+    createComposedAffineApplyOp(&builder, opStmt->getLoc(), indices,
+                                affineApplyOps, results);
     // Create new StoreOp with new affine apply op.
     builder.create<StoreOp>(opStmt->getLoc(), storeOp->getValueToStore(),
                             storeOp->getMemRef(), results);