[MLIR][Affine] Add parametric tile size support for affine.for tiling

[navdeepkk]

Summary

[MLIR][Affine] Add parametric tile size support for affine.for tiling

Add support to tile affine.for ops with parametric sizes (i.e., SSA
values). Currently supports hyper-rectangular loop nests with constant
lower bounds only. Move methods

• moveLoopBody(*)
• getTileableBands(*)
• checkTilingLegality(*)
• tilePerfectlyNested(*)
• constructTiledIndexSetHyperRect(*)

to allow reuse with constant tile size API. Add a test pass -test-affine
-parametric-tile to test parametric tiling.

[bondhugula]

Thanks for contributing this! Please state at the end of the commit summary as to which methods were just moved as a result of this revision (since there appear to be methods that were purely moved to allow reuse with constant tile size API).

mlir/include/mlir/Transforms/LoopUtils.h

91

Nit: Drop "This function ..". Start with "Checks whether ..>". Please also mention at the commit summary which methods were just moved as a result of this revision.

mlir/test/lib/Transforms/TestAffineLoopParametricTiling.cpp

10

pass -> test pass

or

a pass to test parametric tiling of ..

25

Leave a space below.

66–69

Combine please.

299

Avoid auto.

321

as parameters -> as tile sizes

[bondhugula]

mlir/test/Dialect/Affine/loop-tiling-parametric.mlir

258

You can do

// CHECK-LABEL: func @tile_with_...
// CHECK-SAME: ([[ARG0:..

[bondhugula]

mlir/include/mlir/Transforms/LoopUtils.h

102–104

Why has this method been exposed instead of being local/static to LoopUtils.cpp? (I don't see it being used anywhere other that.) Similarly, for other methods. They don't have to be exposed in LoopUtils.h unless they are used from other files.

[navdeepkk]

Address comments on initial diff (290671)

[bondhugula]

mlir/include/mlir/Transforms/LoopUtils.h

103–112

Please add argument names.

105–108

An extra argument to this method has been added without any documentation. It's not even clear here what it is. Parameter names have also been deleted from the existing one.

[navdeepkk]

Address comments on diff (290734).

Add names to input arguments of functions tilePerfectlyNested and
tilePerfectlyNestedParametric. Change type of input arguments
to these functions to match the coding style.

[bondhugula]

mlir/include/mlir/Transforms/LoopUtils.h

98–99

This is again changing the API for the constant tile size case and this isn't mentioned in the commit summary. Is this intended for this revision? Also, please have the output parameters at the end for consistency. tiledNest has been moved to the front. Ideally, please leave the constant tile size API unchanged in this revision.

103–112

Likewise - fix order of arguments. You can't expose the construction fn input in the API.

104

Why is a user expected to provide a function argument here? This doesn't seem to be meaningful.

[navdeepkk]

Address comments on diff (290923).

Move tilePerfectlyNested/tilePerfectlyNestedParametric() back to the
respective passes. Expose moveLoopBody() and checkTilingLegality() in
LoopUtils.h to eliminate redundancy, were removed diff(290734).

[navdeepkk]

Remove auto from appropriate places in TestAffineLoopParametricTiling.cpp

[bondhugula]

mlir/include/mlir/Transforms/LoopUtils.h

105

I see this comment marked as done but this code is still there.

[bondhugula]

mlir/include/mlir/Transforms/LoopUtils.h

105

These should not be exposed.

mlir/lib/Transforms/Utils/LoopUtils.cpp

1012–1013

You can instead use getOps compactly here.

[bondhugula]

mlir/lib/Transforms/Utils/LoopUtils.cpp

424–425

This shouldn't be actually returning LogicalResult - but just a bool. Also this should be marked static.

500

This should be marked static as well.

[bondhugula]

mlir/lib/Transforms/Utils/LoopUtils.cpp

506

dest -> dest

mlir/test/lib/Transforms/TestAffineLoopParametricTiling.cpp

276

This is the utility you are adding and this should be in LoopUtils.cpp! and so should all the other local functions this depends on. TestAffineLoopParametricTiling.cpp is just expected to have the test pass logic.

[navdeepkk]

Address comments on diff (291253).

Split tilePerfecltyNested(*) into performPreTilingChecks(*), constructTiledLoopNest(*),
checkIfHyperRectangular(*) and move all of them to utils to promote code reuse.

[bondhugula]

constructTiledSetHyper... appears to have been moved but not mentioned in the commit summary as an NFC move. Has it been updated in some way?

[bondhugula]

mlir/include/mlir/Transforms/LoopUtils.h

108

You don't need LLVM_NODISCARD are you are using it from the test pass here?

mlir/lib/Transforms/Utils/LoopUtils.cpp

569

Type: wether

570

If this function was simply moved, please mention so. If it was refactored, one might as well using the filling ctor for ops.

... ops(input.begin(), input.end());

843

auto -> AffineMap

869

auto -> AffineDimExpr

mlir/test/Dialect/Affine/loop-tiling-parametric.mlir

4

Please add a summary comment on what this test case file is testing and from where the tile sizes are being read.

[navdeepkk]

In D87353#2275746, @bondhugula wrote:

constructTiledSetHyper... appears to have been moved but not mentioned in the commit summary as an NFC move. Has it been updated in some way?

Sorry, forgot to mention it. It has not been modified in any way just moved to utils because it is being called from tilePerfectlyNested(), which was in turn moved to promote code reuse.

mlir/include/mlir/Transforms/LoopUtils.h

108

Yes, it is being called from line 87 in the test pass.

[bondhugula]

mlir/lib/Transforms/Utils/LoopUtils.cpp

604–606

Please combine the two.

688

Nit:
bound -> bounds

as the lower -> as the corresponding lower

751–753

Is this comment accurate? The new bound isn't going to be s0 - 5 ceildiv s1 - there should be a + 5 to it?

759–760

Likewise. Anyway, both these cases have the same handling, so I'm not sure why the distinction in the comments.

856

auto -> AffineMap

942

Comment here and a comment below for constructParametr.. - otherwise, it's not immediately clear what this one and what the other does.

mlir/test/lib/Transforms/TestAffineLoopParametricTiling.cpp

15–21

You don't several of these headers I think. Please prune. (AffineValueMap, LoopAnalysis, AffineAnalysis, Analysis/Utils.h etc. may not be needed).

[navdeepkk]

Address comments on diff(291964).

[bondhugula]

Looks good!

mlir/test/lib/Transforms/TestAffineLoopParametricTiling.cpp

17

You won't need Block.h - instead, something else lower is needed perhaps.

[navdeepkk]

Remove include directive "block.h" from "TestAffineLoopParametricTiling.cpp"

mlir/include/mlir/Transforms/LoopUtils.h

 /// operation-wise shifts. The shifts are with respect to the original execution
 /// order, and are multiplied by the loop 'step' before being applied. If
 /// `unrollPrologueEpilogue` is set, fully unroll the prologue and epilogue
 /// loops when possible.
 LLVM_NODISCARD
 LogicalResult affineForOpBodySkew(AffineForOp forOp, ArrayRef<uint64_t> shifts,
                                   bool unrollPrologueEpilogue = false);
 
+/// Identify valid and profitable bands of loops to tile. This is currently just
+/// a temporary placeholder to test the mechanics of tiled code generation.
+/// Returns all maximal outermost perfect loop nests to tile.
+void getTileableBands(FuncOp f,
+                      std::vector<SmallVector<AffineForOp, 6>> *bands);
+
 /// Tiles the specified band of perfectly nested loops creating tile-space loops
-/// and intra-tile loops. A band is a contiguous set of loops. `tiledNest` when
-/// non-null is set to the loops of the tiled nest from outermost to innermost.
-/// Loops in `input` are erased when the tiling is successful.
+/// and intra-tile loops. A band is a contiguous set of loops.
 LLVM_NODISCARD
 LogicalResult
 tilePerfectlyNested(MutableArrayRef<AffineForOp> input,
                     ArrayRef<unsigned> tileSizes,
                     SmallVectorImpl<AffineForOp> *tiledNest = nullptr);
 
+/// Tiles the specified band of perfectly nested loops creating tile-space
+/// loops and intra-tile loops, using SSA values as tiling parameters. A band
+/// is a contiguous set of loops.
+LLVM_NODISCARD
+LogicalResult tilePerfectlyNestedParametric(
+    MutableArrayRef<AffineForOp> input, ArrayRef<Value> tileSizes,
+    SmallVectorImpl<AffineForOp> *tiledNest = nullptr);
+
 /// Performs loop interchange on 'forOpA' and 'forOpB'. Requires that 'forOpA'
 /// and 'forOpB' are part of a perfectly nested sequence of loops.
 void interchangeLoops(AffineForOp forOpA, AffineForOp forOpB);
 
 /// Checks if the loop interchange permutation 'loopPermMap', of the perfectly
 /// nested sequence of loops in 'loops', would violate dependences (loop 'i' in
 /// 'loops' is mapped to location 'j = 'loopPermMap[i]' in the interchange).
 bool isValidLoopInterchangePermutation(ArrayRef<AffineForOp> loops,

mlir/lib/Dialect/Affine/Transforms/LoopTiling.cpp

 std::unique_ptr<OperationPass<FuncOp>>
 mlir::createLoopTilingPass(uint64_t cacheSizeBytes) {
   return std::make_unique<LoopTiling>(cacheSizeBytes);
 }
 std::unique_ptr<OperationPass<FuncOp>> mlir::createLoopTilingPass() {
   return std::make_unique<LoopTiling>();
 }
 
-// Move the loop body of AffineForOp 'src' from 'src' into the specified
-// location in destination's body, ignoring the terminator.
-static inline void moveLoopBody(AffineForOp src, AffineForOp dest,
-                                Block::iterator loc) {
-  auto &insts = src.getBody()->getOperations();
-  dest.getBody()->getOperations().splice(loc, insts, insts.begin(),
-                                         std::prev(insts.end()));
-}
-
-// Move the loop body of AffineForOp 'src' from 'src' to the start of dest's
-// body.
-static inline void moveLoopBody(AffineForOp src, AffineForOp dest) {
-  moveLoopBody(src, dest, dest.getBody()->begin());
-}
-
-/// Constructs and sets new loop bounds after tiling for the case of
-/// hyper-rectangular index sets, where the bounds of one dimension do not
-/// depend on other dimensions. Bounds of each dimension can thus be treated
-/// independently, and deriving the new bounds is much simpler and faster
-/// than for the case of tiling arbitrary polyhedral shapes.
-static void
-constructTiledIndexSetHyperRect(MutableArrayRef<AffineForOp> origLoops,
-                                MutableArrayRef<AffineForOp> newLoops,
-                                ArrayRef<unsigned> tileSizes) {
-  assert(!origLoops.empty());
-  assert(origLoops.size() == tileSizes.size());
-
-  OpBuilder b(origLoops[0].getOperation());
-  unsigned width = origLoops.size();
-
-  // Bounds for tile space loops.
-  for (unsigned i = 0; i < width; i++) {
-    OperandRange newLbOperands = origLoops[i].getLowerBoundOperands();
-    OperandRange newUbOperands = origLoops[i].getUpperBoundOperands();
-    newLoops[i].setLowerBound(newLbOperands, origLoops[i].getLowerBoundMap());
-    newLoops[i].setUpperBound(newUbOperands, origLoops[i].getUpperBoundMap());
-    newLoops[i].setStep(tileSizes[i]);
-  }
-  // Bounds for intra-tile loops.
-  for (unsigned i = 0; i < width; i++) {
-    int64_t largestDiv = getLargestDivisorOfTripCount(origLoops[i]);
-    auto mayBeConstantCount = getConstantTripCount(origLoops[i]);
-    // The lower bound is just the tile-space loop.
-    AffineMap lbMap = b.getDimIdentityMap();
-    newLoops[width + i].setLowerBound(
-        /*operands=*/newLoops[i].getInductionVar(), lbMap);
-
-    // Set the upper bound.
-    if (mayBeConstantCount && mayBeConstantCount.getValue() < tileSizes[i]) {
-      // Trip count is less than the tile size: upper bound is lower bound +
-      // trip count.
-      auto ubMap = b.getSingleDimShiftAffineMap(mayBeConstantCount.getValue());
-      newLoops[width + i].setUpperBound(
-          /*operands=*/newLoops[i].getInductionVar(), ubMap);
-    } else if (largestDiv % tileSizes[i] != 0) {
-      // Intra-tile loop ii goes from i to min(i + tileSize, ub_i).
-      // Construct the upper bound map; the operands are the original operands
-      // with 'i' (tile-space loop) appended to it. The new upper bound map is
-      // the original one with an additional expression i + tileSize appended.
-
-      // Add dim operands from original upper bound.
-      SmallVector<Value, 4> ubOperands;
-      auto ub = origLoops[i].getUpperBound();
-      ubOperands.reserve(ub.getNumOperands() + 1);
-      auto origUbMap = ub.getMap();
-      for (unsigned j = 0, e = origUbMap.getNumDims(); j < e; ++j)
-        ubOperands.push_back(ub.getOperand(j));
-
-      // Add dim operand for new loop upper bound.
-      ubOperands.push_back(newLoops[i].getInductionVar());
-
-      // Add symbol operands from original upper bound.
-      for (unsigned j = 0, e = origUbMap.getNumSymbols(); j < e; ++j)
-        ubOperands.push_back(ub.getOperand(origUbMap.getNumDims() + j));
-
-      SmallVector<AffineExpr, 4> boundExprs;
-      boundExprs.reserve(1 + origUbMap.getNumResults());
-      auto dim = b.getAffineDimExpr(origUbMap.getNumDims());
-      // The new upper bound map is the original one with an additional
-      // expression i + tileSize appended.
-      boundExprs.push_back(dim + tileSizes[i]);
-      boundExprs.append(origUbMap.getResults().begin(),
-                        origUbMap.getResults().end());
-      auto ubMap =
-          AffineMap::get(origUbMap.getNumDims() + 1, origUbMap.getNumSymbols(),
-                         boundExprs, b.getContext());
-      newLoops[width + i].setUpperBound(/*operands=*/ubOperands, ubMap);
-    } else {
-      // No need of the min expression.
-      auto dim = b.getAffineDimExpr(0);
-      auto ubMap = AffineMap::get(1, 0, dim + tileSizes[i]);
-      newLoops[width + i].setUpperBound(newLoops[i].getInductionVar(), ubMap);
-    }
-  }
-}
-
-/// This function checks whether hyper-rectangular loop tiling of the nest
-/// represented by `origLoops` is valid. The validity condition is from Irigoin
-/// and Triolet, which states that two tiles cannot depend on each other. We
-/// simplify such condition to just checking whether there is any negative
-/// dependence direction, since we have the prior knowledge that the tiling
-/// results will be hyper-rectangles, which are scheduled in the
-/// lexicographically increasing order on the vector of loop indices. This
-/// function will return failure when any dependence component is negative along
-/// any of `origLoops`.
-static LogicalResult
-checkTilingLegality(MutableArrayRef<mlir::AffineForOp> origLoops) {
-  assert(!origLoops.empty() && "no original loops provided");
-
-  // We first find out all dependences we intend to check.
-  SmallVector<Operation *, 8> loadAndStoreOps;
-  origLoops[0].getOperation()->walk([&](Operation *op) {
-    if (isa<AffineReadOpInterface, AffineWriteOpInterface>(op))
-      loadAndStoreOps.push_back(op);
-  });
-
-  unsigned numOps = loadAndStoreOps.size();
-  unsigned numLoops = origLoops.size();
-  FlatAffineConstraints dependenceConstraints;
-  for (unsigned d = 1; d <= numLoops + 1; ++d) {
-    for (unsigned i = 0; i < numOps; ++i) {
-      Operation *srcOp = loadAndStoreOps[i];
-      MemRefAccess srcAccess(srcOp);
-      for (unsigned j = 0; j < numOps; ++j) {
-        Operation *dstOp = loadAndStoreOps[j];
-        MemRefAccess dstAccess(dstOp);
-
-        SmallVector<DependenceComponent, 2> depComps;
-        dependenceConstraints.reset();
-        DependenceResult result = checkMemrefAccessDependence(
-            srcAccess, dstAccess, d, &dependenceConstraints, &depComps);
-
-        // Skip if there is no dependence in this case.
-        if (!hasDependence(result))
-          continue;
-
-        // Check whether there is any negative direction vector in the
-        // dependence components found above, which means that dependence is
-        // violated by the default hyper-rect tiling method.
-        LLVM_DEBUG(llvm::dbgs() << "Checking whether tiling legality violated "
-                                   "for dependence at depth: "
-                                << Twine(d) << " between:\n";);
-        LLVM_DEBUG(srcAccess.opInst->dump(););
-        LLVM_DEBUG(dstAccess.opInst->dump(););
-        for (unsigned k = 0, e = depComps.size(); k < e; k++) {
-          DependenceComponent depComp = depComps[k];
-          if (depComp.lb.hasValue() && depComp.ub.hasValue() &&
-              depComp.lb.getValue() < depComp.ub.getValue() &&
-              depComp.ub.getValue() < 0) {
-            LLVM_DEBUG(llvm::dbgs()
-                       << "Dependence component lb = "
-                       << Twine(depComp.lb.getValue())
-                       << " ub = " << Twine(depComp.ub.getValue())
-                       << " is negative  at depth: " << Twine(d)
-                       << " and thus violates the legality rule.\n");
-            return failure();
-          }
-        }
-      }
-    }
-  }
-
-  return success();
-}
-/// Tiles the specified band of perfectly nested loops creating tile-space loops
-/// and intra-tile loops. A band is a contiguous set of loops.
-//  TODO: handle non hyper-rectangular spaces.
-LogicalResult
-mlir::tilePerfectlyNested(MutableArrayRef<AffineForOp> input,
-                          ArrayRef<unsigned> tileSizes,
-                          SmallVectorImpl<AffineForOp> *tiledNest) {
-  // Check if the supplied for op's are all successively nested.
-  assert(!input.empty() && "no loops in input band");
-  assert(input.size() == tileSizes.size() && "Too few/many tile sizes");
-
-  assert(isPerfectlyNested(input) && "input loops not perfectly nested");
-
-  auto origLoops = input;
-
-  // Perform tiling legality test.
-  if (failed(checkTilingLegality(origLoops)))
-    origLoops[0].emitRemark("tiled code is illegal due to dependences");
-
-  AffineForOp rootAffineForOp = origLoops[0];
-  auto loc = rootAffineForOp.getLoc();
-  // Note that width is at least one since band isn't empty.
-  unsigned width = input.size();
-
-  SmallVector<AffineForOp, 6> tiledLoops(2 * width);
-
-  // The outermost among the loops as we add more..
-  auto *topLoop = rootAffineForOp.getOperation();
-  AffineForOp innermostPointLoop;
-
-  // Add intra-tile (or point) loops.
-  for (unsigned i = 0; i < width; i++) {
-    OpBuilder b(topLoop);
-    // Loop bounds will be set later.
-    auto pointLoop = b.create<AffineForOp>(loc, 0, 0);
-    pointLoop.getBody()->getOperations().splice(
-        pointLoop.getBody()->begin(), topLoop->getBlock()->getOperations(),
-        topLoop);
-    tiledLoops[2 * width - 1 - i] = pointLoop;
-    topLoop = pointLoop.getOperation();
-    if (i == 0)
-      innermostPointLoop = pointLoop;
-  }
-
-  // Add tile space loops;
-  for (unsigned i = width; i < 2 * width; i++) {
-    OpBuilder b(topLoop);
-    // Loop bounds will be set later.
-    auto tileSpaceLoop = b.create<AffineForOp>(loc, 0, 0);
-    tileSpaceLoop.getBody()->getOperations().splice(
-        tileSpaceLoop.getBody()->begin(), topLoop->getBlock()->getOperations(),
-        topLoop);
-    tiledLoops[2 * width - i - 1] = tileSpaceLoop;
-    topLoop = tileSpaceLoop.getOperation();
-  }
-
-  // Move the loop body of the original nest to the new one.
-  moveLoopBody(origLoops.back(), innermostPointLoop);
-
-  SmallVector<Value, 8> origLoopIVs;
-  extractForInductionVars(input, &origLoopIVs);
-
-  FlatAffineConstraints cst;
-  SmallVector<Operation *, 8> ops;
-  ops.reserve(input.size());
-  for (AffineForOp forOp : input)
-    ops.push_back(forOp);
-  getIndexSet(ops, &cst);
-  if (!cst.isHyperRectangular(0, width)) {
-    rootAffineForOp.emitError("tiled code generation unimplemented for the "
-                              "non-hyperrectangular case");
-    return failure();
-  }
-
-  constructTiledIndexSetHyperRect(origLoops, tiledLoops, tileSizes);
-
-  // Replace original IVs with intra-tile loop IVs.
-  for (unsigned i = 0; i < width; i++)
-    origLoopIVs[i].replaceAllUsesWith(tiledLoops[i + width].getInductionVar());
-
-  // Erase the old loop nest.
-  rootAffineForOp.erase();
-
-  if (tiledNest)
-    *tiledNest = std::move(tiledLoops);
-
-  return success();
-}
-
-// Identify valid and profitable bands of loops to tile. This is currently just
-// a temporary placeholder to test the mechanics of tiled code generation.
-// Returns all maximal outermost perfect loop nests to tile.
-static void getTileableBands(FuncOp f,
-                             std::vector<SmallVector<AffineForOp, 6>> *bands) {
-  // Get maximal perfect nest of 'affine.for' insts starting from root
-  // (inclusive).
-  auto getMaximalPerfectLoopNest = [&](AffineForOp root) {
-    SmallVector<AffineForOp, 6> band;
-    getPerfectlyNestedLoops(band, root);
-    bands->push_back(band);
-  };
-
-  for (auto &block : f)
-    for (auto &op : block)
-      if (auto forOp = dyn_cast<AffineForOp>(op))
-        getMaximalPerfectLoopNest(forOp);
-}
-
 /// Reduces each tile size to the largest divisor of the corresponding trip
 /// count (if the trip count is known).
 static void adjustToDivisorsOfTripCounts(ArrayRef<AffineForOp> band,
                                          SmallVectorImpl<unsigned> *tileSizes) {
   assert(band.size() == tileSizes->size() && "invalid tile size count");
   for (unsigned i = 0, e = band.size(); i < e; i++) {
     unsigned &tSizeAdjusted = (*tileSizes)[i];
-    auto mayConst = getConstantTripCount(band[i]);
+    Optional<uint64_t> mayConst = getConstantTripCount(band[i]);
     if (!mayConst)
       continue;
     // Adjust the tile size to largest factor of the trip count less than
     // tSize.
     uint64_t constTripCount = mayConst.getValue();
     if (constTripCount > 1 && tSizeAdjusted > constTripCount / 2)
       tSizeAdjusted = constTripCount / 2;
     while (constTripCount % tSizeAdjusted != 0)
void LoopTiling::getTileSizes(ArrayRef<AffineForOp> band,
   if (!this->tileSizes.empty()) {
     tileSizes->assign(this->tileSizes.begin(), this->tileSizes.end());
     tileSizes->resize(band.size(), kDefaultTileSize);
     return;
   }
   tileSizes->resize(band.size());
 
   // The first loop in the band.
-  auto rootForOp = band[0];
+  AffineForOp rootForOp = band[0];
   (void)rootForOp;
 
   // Obtain memory footprint and set tile sizes so that a tile fits in
   // the cache size. This is an approximation with the assumption that the
   // footprint increases with the tile size linearly in that dimension (i.e.,
   // assumes one-to-one access function).
-  auto fp = getMemoryFootprintBytes(band[0], 0);
+  Optional<int64_t> fp = getMemoryFootprintBytes(band[0], 0);
   if (!fp) {
     // Fill with default tile sizes if footprint is unknown.
     std::fill(tileSizes->begin(), tileSizes->end(),
               LoopTiling::kDefaultTileSize);
     if (avoidMaxMinBounds)
       adjustToDivisorsOfTripCounts(band, tileSizes);
     LLVM_DEBUG(
         rootForOp.emitWarning("memory footprint unknown: using default tile "
void LoopTiling::runOnFunction() {
   // Tile each band.
   for (auto &band : bands) {
     // Set up tile sizes; fill missing tile sizes at the end with default tile
     // size or tileSize if one was provided.
     SmallVector<unsigned, 6> tileSizes;
     getTileSizes(band, &tileSizes);
     if (llvm::DebugFlag) {
       auto diag = band[0].emitRemark("using tile sizes [");
-      for (auto tSize : tileSizes)
+      for (unsigned tSize : tileSizes)
         diag << tSize << ' ';
       diag << "]\n";
     }
     SmallVector<AffineForOp, 6> tiledNest;
     if (failed(tilePerfectlyNested(band, tileSizes, &tiledNest)))
       return signalPassFailure();
 
     // Separate full and partial tiles.

mlir/lib/Transforms/Utils/LoopUtils.cpp

LogicalResult mlir::affineForOpBodySkew(AffineForOp forOp,
   if (unrollPrologueEpilogue && prologue)
     loopUnrollFull(prologue);
   if (unrollPrologueEpilogue && !epilogue && epilogue != prologue)
     loopUnrollFull(epilogue);
 
   return success();
 }
 
-// Collect perfectly nested loops starting from `rootForOps`.  Loops are
-// perfectly nested if each loop is the first and only non-terminator operation
-// in the parent loop.  Collect at most `maxLoops` loops and append them to
-// `forOps`.
+/// Checks the legality of tiling of a hyper-rectangular loop nest by simply
+/// checking if there is a 'negative' dependence in the memrefs present in
+/// the loop nest. If yes then tiling is invalid.
+static bool
+checkTilingLegalityImpl(MutableArrayRef<mlir::AffineForOp> origLoops) {
+  assert(!origLoops.empty() && "no original loops provided");
+
+  // We first find out all dependences we intend to check.
+  SmallVector<Operation *, 8> loadAndStoreOps;
+  origLoops[0].getOperation()->walk([&](Operation *op) {
+    if (isa<AffineReadOpInterface, AffineWriteOpInterface>(op))
+      loadAndStoreOps.push_back(op);
+  });
+
+  unsigned numOps = loadAndStoreOps.size();
+  unsigned numLoops = origLoops.size();
+  FlatAffineConstraints dependenceConstraints;
+  for (unsigned d = 1; d <= numLoops + 1; ++d) {
+    for (unsigned i = 0; i < numOps; ++i) {
+      Operation *srcOp = loadAndStoreOps[i];
+      MemRefAccess srcAccess(srcOp);
+      for (unsigned j = 0; j < numOps; ++j) {
+        Operation *dstOp = loadAndStoreOps[j];
+        MemRefAccess dstAccess(dstOp);
+
+        SmallVector<DependenceComponent, 2> depComps;
+        dependenceConstraints.reset();
+        DependenceResult result = checkMemrefAccessDependence(
+            srcAccess, dstAccess, d, &dependenceConstraints, &depComps);
+
+        // Skip if there is no dependence in this case.
+        if (!hasDependence(result))
+          continue;
+
+        // Check whether there is any negative direction vector in the
+        // dependence components found above, which means that dependence is
+        // violated by the default hyper-rect tiling method.
+        LLVM_DEBUG(llvm::dbgs() << "Checking whether tiling legality violated "
+                                   "for dependence at depth: "
+                                << Twine(d) << " between:\n";);
+        LLVM_DEBUG(srcAccess.opInst->dump(););
+        LLVM_DEBUG(dstAccess.opInst->dump(););
+        for (unsigned k = 0, e = depComps.size(); k < e; k++) {
+          DependenceComponent depComp = depComps[k];
+          if (depComp.lb.hasValue() && depComp.ub.hasValue() &&
+              depComp.lb.getValue() < depComp.ub.getValue() &&
+              depComp.ub.getValue() < 0) {
+            LLVM_DEBUG(llvm::dbgs()
+                       << "Dependence component lb = "
+                       << Twine(depComp.lb.getValue())
+                       << " ub = " << Twine(depComp.ub.getValue())
+                       << " is negative  at depth: " << Twine(d)
+                       << " and thus violates the legality rule.\n");
+            return false;
+          }
+        }
+      }
+    }
+  }
+
+  return true;
+}
+
+/// Checks whether hyper-rectangular loop tiling of the nest
+/// represented by `origLoops` is valid. The validity condition is from Irigoin
+/// and Triolet, which states that two tiles cannot depend on each other. We
+/// simplify such condition to just checking whether there is any negative
+/// dependence direction, since we have the prior knowledge that the tiling
+/// results will be hyper-rectangles, which are scheduled in the
+/// lexicographically increasing order on the vector of loop indices. This
+/// function will return failure when any dependence component is negative along
+/// any of `origLoops`.
+LogicalResult
+checkTilingLegality(MutableArrayRef<mlir::AffineForOp> origLoops) {
+  return success(checkTilingLegalityImpl(origLoops));
+}
+
+/// Check if the input data is valid and wheter tiled code will be legal or not.
+template <typename t>
+void performPreTilingChecks(MutableArrayRef<AffineForOp> input,
+                            ArrayRef<t> tileSizes) {
+  // Check if the supplied for op's are all successively nested.
+  assert(!input.empty() && "no loops in input band");
+  assert(input.size() == tileSizes.size() && "Too few/many tile sizes");
+
+  assert(isPerfectlyNested(input) && "input loops not perfectly nested");
+
+  // Perform tiling legality test.
+  if (failed(checkTilingLegality(input)))
+    input[0].emitRemark("tiled code is illegal due to dependences");
+}
+
+/// Move the loop body of AffineForOp 'src' from 'src' into the specified
+/// location in destination's body, ignoring the terminator.
+static void moveLoopBodyImpl(AffineForOp src, AffineForOp dest,
+                             Block::iterator loc) {
+  auto &ops = src.getBody()->getOperations();
+  dest.getBody()->getOperations().splice(loc, ops, ops.begin(),
+                                         std::prev(ops.end()));
+}
+
+/// Move the loop body of AffineForOp 'src' from 'src' to the start of dest
+/// body.
+void moveLoopBody(AffineForOp src, AffineForOp dest) {
+  moveLoopBodyImpl(src, dest, dest.getBody()->begin());
+}
+
+/// Constructs tiled loop nest, without setting the loop bounds and move the
+/// body of the original loop nest to the tiled loop nest.
+void constructTiledLoopNest(MutableArrayRef<AffineForOp> origLoops,
+                            AffineForOp rootAffineForOp, unsigned width,
+                            MutableArrayRef<AffineForOp> tiledLoops) {
+  Location loc = rootAffineForOp.getLoc();
+
+  // The outermost among the loops as we add more..
+  Operation *topLoop = rootAffineForOp.getOperation();
+  AffineForOp innermostPointLoop;
+
+  // Add intra-tile (or point) loops.
+  for (unsigned i = 0; i < width; i++) {
+    OpBuilder b(topLoop);
+    // Loop bounds will be set later.
+    AffineForOp pointLoop = b.create<AffineForOp>(loc, 0, 0);
+    pointLoop.getBody()->getOperations().splice(
+        pointLoop.getBody()->begin(), topLoop->getBlock()->getOperations(),
+        topLoop);
+    tiledLoops[2 * width - 1 - i] = pointLoop;
+    topLoop = pointLoop.getOperation();
+    if (i == 0)
+      innermostPointLoop = pointLoop;
+  }
+
+  // Add tile space loops;
+  for (unsigned i = width; i < 2 * width; i++) {
+    OpBuilder b(topLoop);
+    // Loop bounds will be set later.
+    AffineForOp tileSpaceLoop = b.create<AffineForOp>(loc, 0, 0);
+    tileSpaceLoop.getBody()->getOperations().splice(
+        tileSpaceLoop.getBody()->begin(), topLoop->getBlock()->getOperations(),
+        topLoop);
+    tiledLoops[2 * width - i - 1] = tileSpaceLoop;
+    topLoop = tileSpaceLoop.getOperation();
+  }
+
+  // Move the loop body of the original nest to the new one.
+  moveLoopBody(origLoops.back(), innermostPointLoop);
+}
+
+/// Checks whether a loop nest is hyper-rectangular or not.
+LogicalResult checkIfHyperRectangular(MutableArrayRef<AffineForOp> input,
+                                      AffineForOp rootAffineForOp,
+                                      unsigned width) {
+  FlatAffineConstraints cst;
+  SmallVector<Operation *, 8> ops(input.begin(), input.end());
+  getIndexSet(ops, &cst);
+  if (!cst.isHyperRectangular(0, width)) {
+    rootAffineForOp.emitError("tiled code generation unimplemented for the "
+                              "non-hyperrectangular case");
+    return failure();
+  }
+  return success();
+}
+
+/// Set lower and upper bounds of intra-tile loops for parametric tiling.
+//  TODO: Handle non-constant lower bounds.
+static void setIntraTileBoundsParametric(OpBuilder &b, AffineForOp origLoop,
+                                         AffineForOp newInterTileLoop,
+                                         AffineForOp newIntraTileLoop,
+                                         Value tileSize) {
+  // The lower bound for the intra-tile loop is represented by an affine map
+  // as (%i, %t0)->((%i - %origlb) * %t0 + %origlb). Similarly, the upper bound
+  // for the intra-tile loop is represented by an affine map as (%i, %t0)->((%i
+  // - %origlb) * %t0) + (%t0 * %origLoopStep) + %origlb), where %i is loop IV
+  // of the corresponding inter-tile loop, %t0 is the corresponding tiling
+  // parameter, %origlb is lower bound and %origLoopStep is the loop step of the
+  // corresponding inter-tile loop.
+
+  assert(origLoop.hasConstantLowerBound() &&
+         "expected input loops to have constant lower bound.");
+
+  // Get lower bound of original loop as an affine expression.
+  AffineExpr origLowerBoundExpr;
+  origLowerBoundExpr =
+      b.getAffineConstantExpr(origLoop.getConstantLowerBound());
+
+  // Add dim operands from original lower/upper bound.
+  SmallVector<Value, 4> lbOperands, ubOperands;
+  AffineBound lb = origLoop.getLowerBound();
+  AffineBound ub = origLoop.getUpperBound();
+  lbOperands.reserve(lb.getNumOperands() + 2);
+  ubOperands.reserve(ub.getNumOperands() + 2);
+  AffineMap origLbMap = lb.getMap();
+  AffineMap origUbMap = ub.getMap();
+  for (unsigned j = 0, e = origLbMap.getNumDims(); j < e; ++j)
+    lbOperands.push_back(lb.getOperand(j));
+  for (unsigned j = 0, e = origUbMap.getNumDims(); j < e; ++j)
+    ubOperands.push_back(ub.getOperand(j));
+
+  // Add a new dim operand in lb/ubOperands corresponding to the origLoop
+  // IV.
+  lbOperands.push_back(newInterTileLoop.getInductionVar());
+  ubOperands.push_back(newInterTileLoop.getInductionVar());
+
+  // Get loop IV as an affine expression for lower/upper bound. Size of
+  // lb/ubOperands is guaranteed to be atleast one.
+  AffineExpr lbLoopIvExpr = b.getAffineDimExpr(lbOperands.size() - 1);
+  AffineExpr ubLoopIvExpr = b.getAffineDimExpr(ubOperands.size() - 1);
+
+  // Add symbol operands from original lower/upper bound.
+  for (unsigned j = 0, e = origLbMap.getNumSymbols(); j < e; ++j)
+    lbOperands.push_back(lb.getOperand(origLbMap.getNumDims() + j));
+  for (unsigned j = 0, e = origUbMap.getNumSymbols(); j < e; ++j)
+    ubOperands.push_back(ub.getOperand(origUbMap.getNumDims() + j));
+
+  // Add a new symbol operand which is the tile size for this loop.
+  lbOperands.push_back(tileSize);
+  ubOperands.push_back(tileSize);
+
+  SmallVector<AffineExpr, 4> lbBoundExprs;
+  SmallVector<AffineExpr, 4> ubBoundExprs;
+  lbBoundExprs.reserve(origLbMap.getNumResults());
+  ubBoundExprs.reserve(origUbMap.getNumResults());
+
+  // Get tiling parameter as an affine expression for lb/ub.
+  AffineExpr lbTileParameter = b.getAffineSymbolExpr(origLbMap.getNumSymbols());
+  AffineExpr ubTileParameter = b.getAffineSymbolExpr(origUbMap.getNumSymbols());
+
+  // Insert lb as inter-tile ((loop IV - origlb) * tilingParameter) + origlb.
+  lbBoundExprs.push_back(
+      ((lbLoopIvExpr - origLowerBoundExpr) * lbTileParameter) +
+      origLowerBoundExpr);
+
+  // Get the origLoopStep as an affine expression.
+  AffineExpr origLoopStep = b.getAffineConstantExpr(origLoop.getStep());
+
+  // Insert ub as inter-tile ((loop IV - origlb) * tilingParameter) +
+  // (tilingParameter * origLoopStep) + origlb.
+  ubBoundExprs.push_back(
+      ((ubLoopIvExpr - origLowerBoundExpr) * ubTileParameter) +
+      (ubTileParameter * origLoopStep) + origLowerBoundExpr);
+
+  ubBoundExprs.append(origUbMap.getResults().begin(),
+                      origUbMap.getResults().end());
+
+  AffineMap lbMap =
+      AffineMap::get(origLbMap.getNumDims() + 1, origLbMap.getNumSymbols() + 1,
+                     lbBoundExprs, b.getContext());
+  newIntraTileLoop.setLowerBound(lbOperands, lbMap);
+
+  AffineMap ubMap =
+      AffineMap::get(origUbMap.getNumDims() + 1, origUbMap.getNumSymbols() + 1,
+                     ubBoundExprs, b.getContext());
+  newIntraTileLoop.setUpperBound(ubOperands, ubMap);
+
+  // Original loop step must be preserved.
+  newIntraTileLoop.setStep(origLoop.getStep());
+}
+
+/// Set lower and upper bounds of inter-tile loops for parametric tiling.
+//  TODO: Handle non-constant lower bounds.
+static void setInterTileBoundsParametric(OpBuilder &b, AffineForOp origLoop,
+                                         AffineForOp newLoop, Value tileSize) {
+  OperandRange newLbOperands = origLoop.getLowerBoundOperands();
+
+  // The lower bounds for inter-tile loops are same as the correspondig lower
+  // bounds of original loops.
+  newLoop.setLowerBound(newLbOperands, origLoop.getLowerBoundMap());
+
+  // The new upper bound map for inter-tile loops, assuming constant lower
+  // bounds, are now originalLowerBound + ceildiv((orignalUpperBound -
+  // originalLowerBound), tiling paramter); where tiling parameter is the
+  // respective tile size for that loop. For e.g. if the original ubmap was
+  // ()->(1024), the new map will be
+  // ()[s0]->(ceildiv((1024 -lb) % s0)), where s0 is the tiling parameter.
+  // Therefore a new symbol operand is inserted in the map and the result
+  // expression is overwritten.
+
+  assert(origLoop.hasConstantLowerBound() &&
+         "expected input loops to have constant lower bound.");
+
+  // Get lower bound of original loop as an affine expression.
+  AffineExpr origLowerBoundExpr;
+  origLowerBoundExpr =
+      b.getAffineConstantExpr(origLoop.getConstantLowerBound());
+
+  // Add dim operands from original upper bound.
+  SmallVector<Value, 4> ubOperands;
+  AffineBound ub = origLoop.getUpperBound();
+  ubOperands.reserve(ub.getNumOperands() + 1);
+  AffineMap origUbMap = ub.getMap();
+  for (unsigned j = 0, e = origUbMap.getNumDims(); j < e; ++j)
+    ubOperands.push_back(ub.getOperand(j));
+
+  // Add symbol operands from original upper bound.
+  for (unsigned j = 0, e = origUbMap.getNumSymbols(); j < e; ++j)
+    ubOperands.push_back(ub.getOperand(origUbMap.getNumDims() + j));
+
+  // Add a new symbol operand which is the tile size for this loop.
+  ubOperands.push_back(tileSize);
+
+  // Get tiling parameter as an affine expression.
+  AffineExpr tileParameter = b.getAffineSymbolExpr(origUbMap.getNumSymbols());
+
+  SmallVector<AffineExpr, 4> boundExprs;
+  boundExprs.reserve(origUbMap.getNumResults());
+  int64_t origUpperBound;
+  AffineExpr origUpperBoundExpr;
+
+  // If upper bound for the original loop is constant, then the constant can
+  // be obtained as an affine expression straight away.
+  if (origLoop.hasConstantUpperBound()) {
+    origUpperBound = origLoop.getConstantUpperBound();
+
+    // Get original constant upper bound as an affine expression.
+    origUpperBoundExpr = b.getAffineConstantExpr(origUpperBound);
+
+    // Insert the bound as originalLowerBoundceildiv((originalUpperBound -
+    // originalLowerBound), tilingParameter).
+    boundExprs.push_back(
+        origLowerBoundExpr +
+        (origUpperBoundExpr - origLowerBoundExpr).ceilDiv(tileParameter));
+  } else {
+    // If upper bound for the original loop is not constant then two cases
+    // are possible, although there handeling is the same, 1.) The result of
+    // ubmap has only one result expression. For e.g.
+    //    affine.for %i = 5 to %ub
+    //
+    // A symbol operand is added which represents the tiling paramater. The
+    // new loop bounds here will be like ()[s0, s1] -> ((s0 - 5) ceildiv s1 + 5)
+    // where 's0' is the original upper bound and 's1' is the tiling
+    // parameter. 2.) When ubMap has more than one result expression. For e.g.
+    //    #map0 = affine_map<()[s0, s1] -> (s0, s1)
+    //    affine.for %i = 5 to min #map0()[%s0, %s1]
+    //
+    // A symbol operand is added which represents the tiling parameter. The
+    // new loop bounds will be like ()[s0, s1, s2] -> ((s0 - 5) ceildiv s2 + 5,
+    // (s1 -5) ceildiv s2 + 5), where s2 is the tiling parameter.
+
+    // Insert the bounds as originalLowerBound + ceildiv((originalUpperBound -
+    // originalLowerBound), tilingParameter).
+    for (AffineExpr origUpperBoundExpr : origUbMap.getResults())
+      boundExprs.push_back(
+          origLowerBoundExpr +
+          (origUpperBoundExpr - origLowerBoundExpr).ceilDiv(tileParameter));
+  }
+
+  AffineMap ubMap =
+      AffineMap::get(origUbMap.getNumDims(), origUbMap.getNumSymbols() + 1,
+                     boundExprs, b.getContext());
+  newLoop.setUpperBound(ubOperands, ubMap);
+
+  // Original loop step must be preserved.
+  newLoop.setStep(origLoop.getStep());
+}
+
+/// Constructs and sets new loop bounds after tiling for the case of
+/// hyper-rectangular index sets, where the bounds of one dimension do not
+/// depend on other dimensions and tiling parameters are captured from SSA
+/// values. Bounds of each dimension can thus be treated independently,
+/// and deriving the new bounds is much simpler and faster than for the case of
+/// tiling arbitrary polyhedral shapes.
+static void constructParametricallyTiledIndexSetHyperRect(
+    MutableArrayRef<AffineForOp> origLoops,
+    MutableArrayRef<AffineForOp> newLoops, ArrayRef<Value> tileSizes) {
+  assert(!origLoops.empty() && "expected atleast one loop in band");
+  assert(origLoops.size() == tileSizes.size() &&
+         "expected tiling parameter for each loop in band.");
+
+  OpBuilder b(origLoops[0].getOperation());
+  unsigned width = origLoops.size();
+
+  // Set bounds for tile space loops.
+  for (unsigned i = 0; i < width; ++i) {
+    setInterTileBoundsParametric(b, origLoops[i], newLoops[i], tileSizes[i]);
+  }
+
+  // Set bounds for intra-tile loops.
+  for (unsigned i = 0; i < width; ++i) {
+    setIntraTileBoundsParametric(b, origLoops[i], newLoops[i],
+                                 newLoops[i + width], tileSizes[i]);
+  }
+}
+
+/// Constructs and sets new loop bounds after tiling for the case of
+/// hyper-rectangular index sets, where the bounds of one dimension do not
+/// depend on other dimensions. Bounds of each dimension can thus be treated
+/// independently, and deriving the new bounds is much simpler and faster
+/// than for the case of tiling arbitrary polyhedral shapes.
+static void
+constructTiledIndexSetHyperRect(MutableArrayRef<AffineForOp> origLoops,
+                                MutableArrayRef<AffineForOp> newLoops,
+                                ArrayRef<unsigned> tileSizes) {
+  assert(!origLoops.empty());
+  assert(origLoops.size() == tileSizes.size());
+
+  OpBuilder b(origLoops[0].getOperation());
+  unsigned width = origLoops.size();
+
+  // Bounds for tile space loops.
+  for (unsigned i = 0; i < width; i++) {
+    OperandRange newLbOperands = origLoops[i].getLowerBoundOperands();
+    OperandRange newUbOperands = origLoops[i].getUpperBoundOperands();
+    newLoops[i].setLowerBound(newLbOperands, origLoops[i].getLowerBoundMap());
+    newLoops[i].setUpperBound(newUbOperands, origLoops[i].getUpperBoundMap());
+    newLoops[i].setStep(tileSizes[i]);
+  }
+  // Bounds for intra-tile loops.
+  for (unsigned i = 0; i < width; i++) {
+    int64_t largestDiv = getLargestDivisorOfTripCount(origLoops[i]);
+    Optional<uint64_t> mayBeConstantCount = getConstantTripCount(origLoops[i]);
+    // The lower bound is just the tile-space loop.
+    AffineMap lbMap = b.getDimIdentityMap();
+    newLoops[width + i].setLowerBound(
+        /*operands=*/newLoops[i].getInductionVar(), lbMap);
+
+    // Set the upper bound.
+    if (mayBeConstantCount && mayBeConstantCount.getValue() < tileSizes[i]) {
+      // Trip count is less than the tile size: upper bound is lower bound +
+      // trip count.
+      AffineMap ubMap =
+          b.getSingleDimShiftAffineMap(mayBeConstantCount.getValue());
+      newLoops[width + i].setUpperBound(
+          /*operands=*/newLoops[i].getInductionVar(), ubMap);
+    } else if (largestDiv % tileSizes[i] != 0) {
+      // Intra-tile loop ii goes from i to min(i + tileSize, ub_i).
+      // Construct the upper bound map; the operands are the original operands
+      // with 'i' (tile-space loop) appended to it. The new upper bound map is
+      // the original one with an additional expression i + tileSize appended.
+
+      // Add dim operands from original upper bound.
+      SmallVector<Value, 4> ubOperands;
+      AffineBound ub = origLoops[i].getUpperBound();
+      ubOperands.reserve(ub.getNumOperands() + 1);
+      AffineMap origUbMap = ub.getMap();
+      for (unsigned j = 0, e = origUbMap.getNumDims(); j < e; ++j)
+        ubOperands.push_back(ub.getOperand(j));
+
+      // Add dim operand for new loop upper bound.
+      ubOperands.push_back(newLoops[i].getInductionVar());
+
+      // Add symbol operands from original upper bound.
+      for (unsigned j = 0, e = origUbMap.getNumSymbols(); j < e; ++j)
+        ubOperands.push_back(ub.getOperand(origUbMap.getNumDims() + j));
+
+      SmallVector<AffineExpr, 4> boundExprs;
+      boundExprs.reserve(1 + origUbMap.getNumResults());
+      AffineExpr dim = b.getAffineDimExpr(origUbMap.getNumDims());
+      // The new upper bound map is the original one with an additional
+      // expression i + tileSize appended.
+      boundExprs.push_back(dim + tileSizes[i]);
+      boundExprs.append(origUbMap.getResults().begin(),
+                        origUbMap.getResults().end());
+      AffineMap ubMap =
+          AffineMap::get(origUbMap.getNumDims() + 1, origUbMap.getNumSymbols(),
+                         boundExprs, b.getContext());
+      newLoops[width + i].setUpperBound(/*operands=*/ubOperands, ubMap);
+    } else {
+      // No need of the min expression.
+      AffineExpr dim = b.getAffineDimExpr(0);
+      AffineMap ubMap = AffineMap::get(1, 0, dim + tileSizes[i]);
+      newLoops[width + i].setUpperBound(newLoops[i].getInductionVar(), ubMap);
+    }
+  }
+}
+
+/// Tiles the specified band of perfectly nested loops creating tile-space loops
+/// and intra-tile loops. A band is a contiguous set of loops.
+//  TODO: handle non hyper-rectangular spaces.
+LogicalResult
+mlir::tilePerfectlyNested(MutableArrayRef<AffineForOp> input,
+                          ArrayRef<unsigned> tileSizes,
+                          SmallVectorImpl<AffineForOp> *tiledNest) {
+  performPreTilingChecks(input, tileSizes);
+
+  MutableArrayRef<AffineForOp> origLoops = input;
+  AffineForOp rootAffineForOp = origLoops[0];
+  // Note that width is at least one since band isn't empty.
+  unsigned width = input.size();
+  SmallVector<AffineForOp, 6> tiledLoops(2 * width);
+
+  // Construct a tiled loop nest without setting their bounds. Bounds are
+  // set later.
+  constructTiledLoopNest(origLoops, rootAffineForOp, width, tiledLoops);
+
+  SmallVector<Value, 8> origLoopIVs;
+  extractForInductionVars(input, &origLoopIVs);
+
+  if (failed(checkIfHyperRectangular(input, rootAffineForOp, width)))
+    return failure();
+
+  // Set loop bounds for the tiled loop nest.
+  constructTiledIndexSetHyperRect(origLoops, tiledLoops, tileSizes);
+
+  // Replace original IVs with intra-tile loop IVs.
+  for (unsigned i = 0; i < width; i++)
+    origLoopIVs[i].replaceAllUsesWith(tiledLoops[i + width].getInductionVar());
+
+  // Erase the old loop nest.
+  rootAffineForOp.erase();
+
+  if (tiledNest)
+    *tiledNest = std::move(tiledLoops);
+
+  return success();
+}
+
+/// Tiles the specified band of perfectly nested loops creating tile-space
+/// loops and intra-tile loops, using SSA values as tiling parameters. A band
+/// is a contiguous set of loops.
+//  TODO: handle non hyper-rectangular spaces.
+LogicalResult
+mlir::tilePerfectlyNestedParametric(MutableArrayRef<AffineForOp> input,
+                                    ArrayRef<Value> tileSizes,
+                                    SmallVectorImpl<AffineForOp> *tiledNest) {
+  performPreTilingChecks(input, tileSizes);
+
+  MutableArrayRef<AffineForOp> origLoops = input;
+  AffineForOp rootAffineForOp = origLoops[0];
+  // Note that width is at least one since band isn't empty.
+  unsigned width = input.size();
+  SmallVector<AffineForOp, 6> tiledLoops(2 * width);
+
+  // Construct a tiled loop nest without setting their bounds. Bounds are
+  // set later.
+  constructTiledLoopNest(origLoops, rootAffineForOp, width, tiledLoops);
+
+  SmallVector<Value, 8> origLoopIVs;
+  extractForInductionVars(input, &origLoopIVs);
+
+  if (failed(checkIfHyperRectangular(input, rootAffineForOp, width)))
+    return failure();
+
+  // Set loop bounds for the tiled loop nest.
+  constructParametricallyTiledIndexSetHyperRect(origLoops, tiledLoops,
+                                                tileSizes);
+
+  // Replace original IVs with intra-tile loop IVs.
+  for (unsigned i = 0; i < width; i++)
+    origLoopIVs[i].replaceAllUsesWith(tiledLoops[i + width].getInductionVar());
+
+  // Erase the old loop nest.
+  rootAffineForOp.erase();
+
+  if (tiledNest)
+    *tiledNest = std::move(tiledLoops);
+
+  return success();
+}
+
+/// Collect perfectly nested loops starting from `rootForOps`.  Loops are
+/// perfectly nested if each loop is the first and only non-terminator operation
+/// in the parent loop.  Collect at most `maxLoops` loops and append them to
+/// `forOps`.
 template <typename T>
 static void getPerfectlyNestedLoopsImpl(
     SmallVectorImpl<T> &forOps, T rootForOp,
     unsigned maxLoops = std::numeric_limits<unsigned>::max()) {
   for (unsigned i = 0; i < maxLoops; ++i) {
     forOps.push_back(rootForOp);
     Block &body = rootForOp.region().front();
     if (body.begin() != std::prev(body.end(), 2))
void mlir::getPerfectlyNestedLoops(SmallVectorImpl<AffineForOp> &nestedLoops,
   getPerfectlyNestedLoopsImpl(nestedLoops, root);
 }
 
 void mlir::getPerfectlyNestedLoops(SmallVectorImpl<scf::ForOp> &nestedLoops,
                                    scf::ForOp root) {
   getPerfectlyNestedLoopsImpl(nestedLoops, root);
 }
 
+/// Identify valid and profitable bands of loops to tile. This is currently just
+/// a temporary placeholder to test the mechanics of tiled code generation.
+/// Returns all maximal outermost perfect loop nests to tile.
+void mlir::getTileableBands(FuncOp f,
+                            std::vector<SmallVector<AffineForOp, 6>> *bands) {
+  // Get maximal perfect nest of 'affine.for' insts starting from root
+  // (inclusive).
+  for (AffineForOp forOp : f.getOps<AffineForOp>()) {
+    SmallVector<AffineForOp, 6> band;
+    getPerfectlyNestedLoops(band, forOp);
+    bands->push_back(band);
+  }
+}
+
 /// Unrolls this loop completely.
 LogicalResult mlir::loopUnrollFull(AffineForOp forOp) {
   Optional<uint64_t> mayBeConstantTripCount = getConstantTripCount(forOp);
   if (mayBeConstantTripCount.hasValue()) {
     uint64_t tripCount = mayBeConstantTripCount.getValue();
     if (tripCount == 1)
       return promoteIfSingleIteration(forOp);
     return loopUnrollByFactor(forOp, tripCount);

mlir/test/Dialect/Affine/loop-tiling-parametric.mlir

+// RUN: mlir-opt %s -split-input-file -test-affine-parametric-tile | FileCheck %s
+// Test cases to test the utility introduced to tile affine for loops using
+// SSA values as tiling parameters(tile sizes). The tile sizes are expected
+// to be passed as input arguments(before any other argument) to the function
+// enclosing the loop nest. Currently hyper-rectangular loop nests with constant
+// lower bounds are supported.
+
+// -----
+
+// CHECK-DAG: [[LBI:#map[0-9]+]] = affine_map<(d0)[s0] -> (d0 * s0)>
+// CHECK-DAG: [[UBI0:#map[0-9]+]] = affine_map<(d0)[s0] -> (d0 * s0 + s0, 256)>
+// CHECK-DAG: [[UBI1:#map[0-9]+]] = affine_map<(d0)[s0] -> (d0 * s0 + s0, 512)>
+// CHECK-DAG: [[UBI2:#map[0-9]+]] = affine_map<(d0)[s0] -> (d0 * s0 + s0, 1024)>
+// CHECK-DAG: [[UBO0:#map[0-9]+]] = affine_map<()[s0] -> (256 ceildiv s0)>
+// CHECK-DAG: [[UBO1:#map[0-9]+]] = affine_map<()[s0] -> (512 ceildiv s0)>
+// CHECK-DAG: [[UBO2:#map[0-9]+]] = affine_map<()[s0] -> (1024 ceildiv s0)>
+
+// CHECK: func @loop_tiling_3d([[ARG0:%arg[0-9]+]]: index, [[ARG1:%arg[0-9]+]]: index, [[ARG2:%arg[0-9]+]]: index)
+// CHECK-NEXT:   affine.for [[ARG3:%arg[0-9]+]] = 0 to [[UBO0]](){{.*}}[[ARG0]]
+// CHECK-NEXT:     affine.for [[ARG4:%arg[0-9]+]] = 0 to [[UBO1]](){{.*}}[[ARG1]]
+// CHECK-NEXT:       affine.for [[ARG5:%arg[0-9]+]] = 0 to [[UBO2]](){{.*}}[[ARG2]]
+// CHECK-NEXT:         affine.for %[[I:.*]] = [[LBI]]{{.*}}[[ARG3]]{{.*}}[[ARG0]]{{.*}} to min [[UBI0]]{{.*}}[[ARG3]]{{.*}}[[ARG0]]
+// CHECK-NEXT:          affine.for %[[J:.*]] = [[LBI]]{{.*}}[[ARG4]]{{.*}}[[ARG1]]{{.*}} to min [[UBI1]]{{.*}}[[ARG4]]{{.*}}[[ARG1]]
+// CHECK-NEXT:            affine.for %[[K:.*]] = [[LBI]]{{.*}}[[ARG5]]{{.*}}[[ARG2]]{{.*}} to min [[UBI2]]{{.*}}[[ARG5]]{{.*}}[[ARG2]]
+// CHECK-NEXT:              "test.foo"(%[[I]], %[[J]], %[[K]])
+func @loop_tiling_3d(%t0 : index, %t1 : index, %t2 : index) {
+  affine.for %i = 0 to 256 {
+    affine.for %j = 0 to 512 {
+      affine.for %k = 0 to 1024 {
+        "test.foo"(%i, %j, %k) : (index, index, index) -> ()
+      }
+    }
+  }
+  return
+}
+
+// -----
+
+// CHECK-DAG: [[LBI:#map[0-9]+]] = affine_map<(d0)[s0] -> (d0 * s0)>
+// CHECK-DAG: [[UBI0:#map[0-9]+]] = affine_map<(d0)[s0] -> (d0 * s0 + s0 * 4, 256)>
+// CHECK-DAG: [[UBI1:#map[0-9]+]] = affine_map<(d0)[s0] -> (d0 * s0 + s0 * 3, 512)>
+// CHECK-DAG: [[UBI2:#map[0-9]+]] = affine_map<(d0)[s0] -> (d0 * s0 + s0 * 2, 1024)>
+// CHECK-DAG: [[UBO0:#map[0-9]+]] = affine_map<()[s0] -> (256 ceildiv s0)>
+// CHECK-DAG: [[UBO1:#map[0-9]+]] = affine_map<()[s0] -> (512 ceildiv s0)>
+// CHECK-DAG: [[UBO2:#map[0-9]+]] = affine_map<()[s0] -> (1024 ceildiv s0)>
+
+// CHECK: func @loop_tiling_non_unit_step([[ARG0:%arg[0-9]+]]: index, [[ARG1:%arg[0-9]+]]: index, [[ARG2:%arg[0-9]+]]: index)
+// CHECK-NEXT:  affine.for [[ARG3:%arg[0-9]+]] = 0 to [[UBO0]](){{.*}}[[ARG0]]{{.*}}step 4
+// CHECK-NEXT:    affine.for [[ARG4:%arg[0-9]+]] = 0 to [[UBO1]](){{.*}}[[ARG1]]{{.*}} step 3
+// CHECK-NEXT:       affine.for [[ARG5:%arg[0-9]+]] = 0 to [[UBO2]](){{.*}}[[ARG2]]{{.*}} step 2
+// CHECK-NEXT:         affine.for %[[I:.*]] = [[LBI]]{{.*}}[[ARG3]]{{.*}}[[ARG0]]{{.*}} to min [[UBI0]]{{.*}}[[ARG3]]{{.*}}[[ARG0]]{{.*}} step 4
+// CHECK-NEXT:          affine.for %[[J:.*]] = [[LBI]]{{.*}}[[ARG4]]{{.*}}[[ARG1]]{{.*}} to min [[UBI1]]{{.*}}[[ARG4]]{{.*}}[[ARG1]]{{.*}} step 3
+// CHECK-NEXT:            affine.for %[[K:.*]] = [[LBI]]{{.*}}[[ARG5]]{{.*}}[[ARG2]]{{.*}} to min [[UBI2]]{{.*}}[[ARG5]]{{.*}}[[ARG2]]{{.*}} step 2
+// CHECK-NEXT:              "test.foo"(%[[I]], %[[J]], %[[K]])
+func @loop_tiling_non_unit_step(%t0: index, %t1: index, %t2: index){
+  affine.for %i = 0 to 256 step 4 {
+    affine.for %j = 0 to 512  step 3 {
+      affine.for %k = 0 to 1024 step 2 {
+        "test.foo"(%i, %j, %k) : (index, index, index) -> ()
+      }
+    }
+  }
+  return
+}
+
+// -----
+
+// CHECK-DAG: [[LBI0:#map[0-9]+]] = affine_map<(d0)[s0] -> (d0 * s0)>
+// CHECK-DAG: [[UBI0:#map[0-9]+]] = affine_map<(d0)[s0, s1, s2] -> (d0 * s2 + s2, s0, 4096 floordiv s1)>
+// CHECK-DAG: [[UBO0:#map[0-9]+]] = affine_map<()[s0, s1, s2] -> (s0 ceildiv s2, (4096 floordiv s1) ceildiv s2)>
+
+// CHECK: func @tile_loop_with_div_in_upper_bound([[ARG0:%arg[0-9]+]]: index, %{{.*}}: memref<?xi32>, %{{.*}}: index, %{{.*}}: index)
+#ub = affine_map<()[s0, s1] -> (s0, 4096 floordiv s1)>
+func @tile_loop_with_div_in_upper_bound(%t5 : index, %A : memref<? x i32>, %L : index, %U : index) {
+  %c0 = constant 0 : index
+  %M = dim %A, %c0 : memref<? x i32>
+  affine.for %i = 0 to min #ub()[%M, %U] {
+    addi %i, %i : index
+  }
+  // CHECK:  affine.for [[ARG1:%arg[0-9]+]] = 0 to min [[UBO0]]()[%{{.*}}, %{{.*}}, [[ARG0]]]
+  // CHECK-NEXT:    affine.for %[[I:.*]] = [[LBI0]]([[ARG1]]){{.*}}[[ARG0]]{{.*}} to min [[UBI0]]({{.*}})[{{.*}}, {{.*}}, [[ARG0]]]
+  // CHECK-NEXT:      addi %[[I]], %[[I]]
+  return
+}
+
+// -----
+
+// CHECK-DAG: [[LBI0:#map[0-9]+]] = affine_map<(d0)[s0] -> (d0 * s0)>
+// CHECK-DAG: [[UBI0:#map[0-9]+]] = affine_map<(d0)[s0, s1, s2] -> (d0 * s2 + s2 * 4, s0, 4096 floordiv s1)>
+// CHECK-DAG: [[UBO0:#map[0-9]+]] = affine_map<()[s0, s1, s2] -> (s0 ceildiv s2, (4096 floordiv s1) ceildiv s2)>
+
+// CHECK: func @tile_loop_with_div_in_upper_bound_non_unit_step([[ARG0:%arg[0-9]+]]: index, %{{.*}}: memref<?xi32>, %{{.*}}: index, %{{.*}}: index)
+#ub = affine_map<()[s0, s1] -> (s0, 4096 floordiv s1)>
+func @tile_loop_with_div_in_upper_bound_non_unit_step(%t5 : index, %A : memref<? x i32>, %L : index, %U : index) {
+  %c0 = constant 0 : index
+  %M = dim %A, %c0 : memref<? x i32>
+  affine.for %i = 0 to min #ub()[%M, %U] step 4 {
+    addi %i, %i : index
+  }
+  // CHECK: affine.for [[ARG1:%arg[0-9]+]] = 0 to min [[UBO0]]()[%{{.*}}, %{{.*}}, [[ARG0]]]{{.*}} step 4{{.*}}
+  // CHECK-NEXT:    affine.for %[[I:.*]] = [[LBI0]]([[ARG1]]){{.*}}[[ARG0]]{{.*}} to min [[UBI0]]({{.*}})[{{.*}}, {{.*}}, [[ARG0]]]{{.*}} step 4{{.*}}
+  // CHECK-NEXT:      addi %[[I]], %[[I]]
+  return
+}
+
+// -----
+
+// CHECK-DAG: [[LBI0:#map[0-9]+]] = affine_map<(d0)[s0] -> ((d0 - 8) * s0 + 8)>
+// CHECK-DAG: [[UBI2:#map[0-9]+]] = affine_map<(d0)[s0, s1] -> ((d0 - 8) * s1 + s1 * 4 + 8, s0 + 16)>
+// CHECK-DAG: [[UBI1:#map[0-9]+]] = affine_map<(d0)[s0, s1] -> ((d0 - 8) * s1 + s1 + 8, s0 + 16)>
+// CHECK-DAG: [[UBI0:#map[0-9]+]] = affine_map<(d0)[s0] -> ((d0 - 8) * s0 + s0 + 8, 256)>
+// CHECK-DAG: [[UBO1:#map[0-9]+]] = affine_map<()[s0, s1] -> ((s0 + 8) ceildiv s1 + 8)>
+// CHECK-DAG: [[UBO0:#map[0-9]+]] = affine_map<()[s0] -> (248 ceildiv s0 + 8)>
+
+// CHECK: func @tile_loop_with_non_zero_lb([[ARG0:%arg[0-9]+]]: index, [[ARG1:%arg[0-9]+]]: index, [[ARG2:%arg[0-9]+]]: index, %{{.*}}: index)
+// CHECK-NEXT:  affine.for [[ARG3:%arg[0-9+]]] = 8 to [[UBO0]]{{.*}}[[ARG0]]{{.*}}
+// CHECK-NEXT:    affine.for [[ARG4:%arg[0-9+]]] = 8 to [[UBO1]]{{.*}}[[ARG1]]{{.*}}
+// CHECK-NEXT:      affine.for [[ARG5:%arg[0-9+]]] = 8 to [[UBO1]]{{.*}}[[ARG2]]{{.*}} step 4
+// CHECK-NEXT:        affine.for %[[I:.*]] = [[LBI0]]([[ARG3]]){{.*}}[[ARG0]]{{.*}} to min [[UBI0]]([[ARG3]]){{.*}}[[ARG0]]{{.*}}
+// CHECK-NEXT:          affine.for %[[J:.*]] = [[LBI0]]([[ARG4]]){{.*}}[[ARG1]]{{.*}} to min [[UBI1]]([[ARG4]]){{.*}}[[ARG1]]{{.*}}
+// CHECK-NEXT:            affine.for %[[K:.*]] = [[LBI0]]([[ARG5]]){{.*}}[[ARG2]]{{.*}} to min [[UBI2]]([[ARG5]]){{.*}}[[ARG2]]{{.*}}step 4{{.*}}
+// CHECK-NEXT:              "test.foo"(%[[I]], %[[J]], %[[K]]) : (index, index, index) -> ()
+#ubi = affine_map<()[s0] -> (s0 + 16)>
+func @tile_loop_with_non_zero_lb(%t0: index, %t1: index, %t2: index, %U: index){
+  affine.for %i = 8 to 256 {
+    affine.for %j = 8 to #ubi()[%U] {
+      affine.for %k = 8 to #ubi()[%U] step 4 {
+        "test.foo"(%i, %j, %k) : (index, index, index) -> ()
+      }
+    }
+  }
+  return
+}
+
+// -----
+
+// CHECK-DAG: [[LBI:#map[0-9]+]] = affine_map<(d0)[s0] -> (d0 * s0)>
+// CHECK-DAG: [[UBI0:#map[0-9]+]] = affine_map<(d0)[s0] -> (d0 * s0 + s0, 256)>
+// CHECK-DAG: [[UBI1:#map[0-9]+]] = affine_map<(d0)[s0] -> (d0 * s0 + s0, 250)>
+// CHECK-DAG: [[UBO0:#map[0-9]+]] = affine_map<()[s0] -> (256 ceildiv s0)>
+// CHECK-DAG: [[UBO1:#map[0-9]+]] = affine_map<()[s0] -> (250 ceildiv s0)>
+
+// CHECK: func @simple_matmul([[ARG0:%arg[0-9]+]]: index, [[ARG1:%arg[0-9]+]]: index, [[ARG2:%arg[0-9]+]]: index{{.*}})
+// CHECK-NEXT:   affine.for [[ARG3:%arg[0-9]+]] = 0 to [[UBO0]](){{.*}}[[ARG0]]{{.*}}
+// CHECK-NEXT:     affine.for [[ARG4:%arg[0-9]+]] = 0 to [[UBO0]](){{.*}}[[ARG1]]{{.*}}
+// CHECK-NEXT:       affine.for [[ARG5:%arg[0-9]+]] = 0 to [[UBO1]](){{.*}}[[ARG2]]{{.*}}
+// CHECK-NEXT:         affine.for %[[I:.*]] = [[LBI]]{{.*}}[[ARG3]]{{.*}}[[ARG0]]{{.*}} to min [[UBI0]]{{.*}}[[ARG3]]{{.*}}[[ARG0]]{{.*}}
+// CHECK-NEXT:          affine.for %[[J:.*]] = [[LBI]]{{.*}}[[ARG4]]{{.*}}[[ARG1]]{{.*}} to min [[UBI0]]{{.*}}[[ARG4]]{{.*}}[[ARG1]]{{.*}}
+// CHECK-NEXT:            affine.for %[[K:.*]] = [[LBI]]{{.*}}[[ARG5]]{{.*}}[[ARG2]]{{.*}} to min [[UBI1]]{{.*}}[[ARG5]]{{.*}}[[ARG2]]{{.*}}
+// CHECK-NEXT:                 affine.load %{{.*}}[%[[I]], %[[K]]]
+// CHECK-NEXT:                 affine.load %{{.*}}[%[[K]], %[[J]]]
+// CHECK-NEXT:                 affine.load %{{.*}}[%[[I]], %[[J]]]
+// CHECK-NEXT:                 mulf %{{.*}}
+// CHECK-NEXT:                 addf %{{.*}}
+// CHECK-NEXT:                 affine.store %{{.*}}[%[[I]], %[[J]]]
+func @simple_matmul(%t6 : index, %t7 : index, %t8 : index, %arg0: memref<256x256xvector<64xf32>>, %arg1: memref<256x256xvector<64xf32>>, %arg2: memref<256x256xvector<64xf32>>) -> memref<256x256xvector<64xf32>> {
+  affine.for %i = 0 to 256 {
+    affine.for %j = 0 to 256 {
+      affine.for %k = 0 to 250 {
+        %l = affine.load %arg0[%i, %k] : memref<256x256xvector<64xf32>>
+        %r = affine.load %arg1[%k, %j] : memref<256x256xvector<64xf32>>
+        %o = affine.load %arg2[%i, %j] : memref<256x256xvector<64xf32>>
+        %m = mulf %l, %r : vector<64xf32>
+        %a = addf %o, %m : vector<64xf32>
+        affine.store %a, %arg2[%i, %j] : memref<256x256xvector<64xf32>>
+      }
+    }
+  }
+  return %arg2 : memref<256x256xvector<64xf32>>
+}
+
+// -----
+
+// CHECK-DAG: [[LBI0:#map[0-9]+]] = affine_map<(d0)[s0] -> (d0 * s0)>
+// CHECK-DAG: [[UBI0:#map[0-9]+]] = affine_map<(d0)[s0, s1] -> (d0 * s1 + s1, s0)>
+// CHECK-DAG: [[UBO0:#map[0-9]+]] = affine_map<()[s0, s1] -> (s0 ceildiv s1)>
+
+// CHECK: func @tile_with_symbolic_loop_upper_bounds([[ARG0:%arg[0-9]+]]: index, [[ARG1:%arg[0-9]+]]: index{{.*}}){{.*}}
+// CHECK:        affine.for [[ARG2:%arg[0-9]+]] = 0 to [[UBO0]](){{.*}}[[ARG0]]{{.*}}
+// CHECK-NEXT:     affine.for [[ARG3:%arg[0-9]+]] = 0 to [[UBO0]](){{.*}}[[ARG1]]{{.*}}
+// CHECK-NEXT:       affine.for %[[I0:.*]] = [[LBI0]]{{.*}}[[ARG2]]{{.*}}[[ARG0]]{{.*}} to min [[UBI0]]{{.*}}[[ARG2]]{{.*}}[[ARG0]]{{.*}}
+// CHECK-NEXT:         affine.for %[[I1:.*]] = [[LBI0]]{{.*}}[[ARG3]]{{.*}}[[ARG1]]{{.*}} to min [[UBI0]]{{.*}}[[ARG3]]{{.*}}[[ARG1]]{{.*}}
+// CHECK-NEXT:           affine.store %{{.*}}, %{{.*}}[%[[I0]], %[[I1]]] : memref<?x?xf32>
+// CHECK-NEXT:           affine.for %[[I2:.*]] = 0 to %{{.*}} {
+// CHECK-NEXT:             affine.load %{{.*}}%[[I0]], %[[I2]]
+// CHECK-NEXT:             affine.load %{{.*}}%[[I2]], %[[I1]]
+// CHECK-NEXT:             mulf
+// CHECK-NEXT:             affine.load %{{.*}}%[[I0]], %[[I1]]
+// CHECK-NEXT:             addf
+// CHECK-NEXT:             affine.store %{{.*}}%[[I0]], %[[I1]]
+func @tile_with_symbolic_loop_upper_bounds(%t9 : index, %t10: index, %arg0: memref<?x?xf32>, %arg1: memref<?x?xf32>, %arg2: memref<?x?xf32>) {
+  %cst = constant 0.000000e+00 : f32
+  %c0 = constant 0 : index
+  %0 = dim %arg0, %c0 : memref<?x?xf32>
+  affine.for %i0 = 0 to %0 {
+    affine.for %i1 = 0 to %0 {
+      affine.store %cst, %arg2[%i0, %i1] : memref<?x?xf32>
+      affine.for %i2 = 0 to %0 {
+        %1 = affine.load %arg0[%i0, %i2] : memref<?x?xf32>
+        %2 = affine.load %arg1[%i2, %i1] : memref<?x?xf32>
+        %3 = mulf %1, %2 : f32
+        %4 = affine.load %arg2[%i0, %i1] : memref<?x?xf32>
+        %5 = addf %4, %3 : f32
+        affine.store %5, %arg2[%i0, %i1] : memref<?x?xf32>
+      }
+    }
+  }
+  return
+}
+
+// -----
+
+// CHECK-DAG: [[LBI0:#map[0-9]+]] = affine_map<(d0)[s0] -> (d0 * s0)>
+// CHECK-DAG: [[UBI0:#map[0-9]+]] = affine_map<(d0)[s0, s1, s2] -> (d0 * s2 + s2, s0 + s1)>
+// CHECK-DAG: [[UBO0:#map[0-9]+]] = affine_map<()[s0, s1, s2] -> ((s0 + s1) ceildiv s2)>
+
+// CHECK: func @tile_with_loop_upper_bounds_in_two_symbols([[ARG0:%arg[0-9]+]]: index{{.*}}){{.*}}
+func @tile_with_loop_upper_bounds_in_two_symbols(%t11 : index, %arg0: memref<?xf32>, %limit: index) {
+  %c0 = constant 0 : index
+  %dim0 = dim %arg0, %c0 : memref<?xf32>
+  affine.for %i0 = 0 to affine_map<()[s0, s1] -> (s0 + s1)> ()[%dim0, %limit] {
+    %v0 = affine.load %arg0[%i0] : memref<?xf32>
+  }
+  // CHECK:  affine.for [[ARG1:%arg[0-9]+]] = 0 to [[UBO0]]()[%{{.*}}, %{{.*}}, [[ARG0]]]
+  // CHECK-NEXT:    affine.for %[[I:.*]] = [[LBI0]]([[ARG1]]){{.*}}[[ARG0]]{{.*}} to min [[UBI0]]([[ARG1]])[{{.*}}, {{.*}}, [[ARG0]]]
+  // CHECK-NEXT:      affine.load %{{.*}}[%[[I]]]
+  return
+}
+
+// -----
+
+// CHECK-DAG: [[LBI0:#map[0-9]+]] = affine_map<(d0)[s0] -> (d0 * s0)>
+// CHECK-DAG: [[UBI1:#map[0-9]+]] = affine_map<(d0, d1)[s0, s1] -> (d1 * s1 + s1, d0 + s0 + 4)>
+// CHECK-DAG: [[UBI0:#map[0-9]+]] = affine_map<(d0, d1)[s0, s1] -> (d1 * s1 + s1, d0 + s0 + 2)>
+// CHECK-DAG: [[LBO0:#map[0-9]+]] = affine_map<() -> (0)>
+// CHECK-DAG: [[UBO1:#map[0-9]+]] = affine_map<(d0)[s0, s1] -> ((d0 + s0 + 4) ceildiv s1)>
+// CHECK-DAG: [[UBO0:#map[0-9]+]] = affine_map<(d0)[s0, s1] -> ((d0 + s0 + 2) ceildiv s1)>
+
+// CHECK: func @tile_with_upper_bounds_in_dimensions_and_symbols([[ARG0:%arg[0-9]+]]: index, [[ARG1:%arg[0-9]+]]: index, [[ARG2:%arg[0-9]+]]: index, [[ARG3:%arg[0-9]+]]: index{{.*}}){{.*}}
+// CHECK-NEXT: affine.for [[ARG4:%arg[0-9]+]] = 0 to [[UBO0]]({{.*}}){{.*}}[[ARG0]]
+// CHECK-NEXT:   affine.for [[ARG5:%arg[0-9]+]] = 0 to [[UBO1]]({{.*}}){{.*}}[[ARG1]]
+// CHECK-NEXT:     affine.for {{.*}} = [[LBI0]]([[ARG4]]){{.*}}[[ARG0]]{{.*}} to min [[UBI0]]({{.*}}, [[ARG4]]){{.*}}[[ARG0]]{{.*}}
+// CHECK-NEXT:       affine.for {{.*}} = [[LBI0]]([[ARG5]]){{.*}}[[ARG1]]{{.*}} to min [[UBI1]]({{.*}}, [[ARG5]]){{.*}}[[ARG1]]{{.*}}
+func @tile_with_upper_bounds_in_dimensions_and_symbols(%t12 : index, %t13 :index, %M: index, %N:  index, %K: index) {
+  affine.for %i = 0 to affine_map<(d0)[s0] -> (d0 + s0 + 2)>(%M)[%K] {
+    affine.for %j = 0 to affine_map<(d0)[s0] -> (d0 + s0 + 4)>(%N)[%K] {
+      "test.foo" () : () -> ()
+    }
+  }
+  return
+}
+
+// -----
+
+// CHECK-DAG: [[LBI0:#map[0-9]+]] = affine_map<(d0)[s0] -> (d0 * s0)>
+// CHECK-DAG: [[UBI1:#map[0-9]+]] = affine_map<(d0, d1)[s0, s1] -> (d1 * s1 + s1 * 4, d0 + s0 + 4)>
+// CHECK-DAG: [[UBI0:#map[0-9]+]] = affine_map<(d0, d1)[s0, s1] -> (d1 * s1 + s1 * 2, d0 + s0 + 2)>
+// CHECK-DAG: [[LBO0:#map[0-9]+]] = affine_map<() -> (0)>
+// CHECK-DAG: [[UBO1:#map[0-9]+]] = affine_map<(d0)[s0, s1] -> ((d0 + s0 + 4) ceildiv s1)>
+// CHECK-DAG: [[UBO0:#map[0-9]+]] = affine_map<(d0)[s0, s1] -> ((d0 + s0 + 2) ceildiv s1)>
+
+// CHECK: func @tile_with_upper_bounds_in_dimensions_and_symbols_non_unit_steps
+// CHECK-SAME: ([[ARG0:%arg[0-9]+]]: index, [[ARG1:%arg[0-9]+]]: index, [[ARG2:%arg[0-9]+]]: index, [[ARG3:%arg[0-9]+]]: index{{.*}}){{.*}}
+// CHECK-NEXT: affine.for [[ARG4:%arg[0-9]+]] = 0 to [[UBO0]]({{.*}}){{.*}}[[ARG0]]{{.*}} step 2{{.*}}
+// CHECK-NEXT:   affine.for [[ARG5:%arg[0-9]+]] = 0 to [[UBO1]]({{.*}}){{.*}}[[ARG1]]{{.*}} step 4{{.*}}
+// CHECK-NEXT:     affine.for {{.*}} = [[LBI0]]([[ARG4]]){{.*}}[[ARG0]]{{.*}} to min [[UBI0]]({{.*}}, [[ARG4]]){{.*}}[[ARG0]]{{.*}} step 2{{.*}}
+// CHECK-NEXT:       affine.for {{.*}} = [[LBI0]]([[ARG5]]){{.*}}[[ARG1]]{{.*}} to min [[UBI1]]({{.*}}, [[ARG5]]){{.*}}[[ARG1]]{{.*}} step 4{{.*}}
+func @tile_with_upper_bounds_in_dimensions_and_symbols_non_unit_steps(%t12 : index, %t13 :index, %M: index, %N :  index, %K: index) {
+  affine.for %i = 0 to affine_map<(d0)[s0] -> (d0 + s0 + 2)>(%M)[%K] step 2 {
+    affine.for %j = 0 to affine_map<(d0)[s0] -> (d0 + s0 + 4)>(%N)[%K] step 4 {
+      "test.foo" () : () -> ()
+    }
+  }
+  return
+}

mlir/test/lib/Transforms/CMakeLists.txt

 # Exclude tests from libMLIR.so
 add_mlir_library(MLIRTestTransforms
   TestAllReduceLowering.cpp
+  TestAffineLoopParametricTiling.cpp
   TestBufferPlacement.cpp
   TestExpandTanh.cpp
   TestCallGraph.cpp
   TestConstantFold.cpp
   TestConvVectorization.cpp
   TestConvertCallOp.cpp
   TestConvertGPUKernelToCubin.cpp
   TestConvertGPUKernelToHsaco.cpp

mlir/test/lib/Transforms/TestAffineLoopParametricTiling.cpp

+//= TestAffineLoopParametricTiling.cpp -- Parametric Affine loop tiling pass =//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// This file implements a test pass to test parametric tiling of perfectly
+// nested affine for loops.
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Dialect/Affine/IR/AffineOps.h"
+#include "mlir/Dialect/Affine/Passes.h"
+#include "mlir/Transforms/LoopUtils.h"
+
+using namespace mlir;
+
+#define DEBUG_TYPE "test-affine-parametric-tile"
+
+namespace {
+
+struct TestAffineLoopParametricTiling
+    : public PassWrapper<TestAffineLoopParametricTiling, FunctionPass> {
+  void runOnFunction() override;
+};
+} // end anonymous namespace
+
+/// Checks if the function enclosing the loop nest has any arguments passed to
+/// it, which can be used as tiling parameters. Assumes that atleast 'n'
+/// arguments are passed, where 'n' is the number of loops in the loop nest.
+static void checkIfTilingParametersExist(ArrayRef<AffineForOp> band) {
+  assert(!band.empty() && "no loops in input band");
+  AffineForOp topLoop = band[0];
+
+  if (FuncOp funcOp = dyn_cast<FuncOp>(topLoop.getParentOp()))
+    assert(funcOp.getNumArguments() >= band.size() && "Too few tile sizes");
+}
+
+/// Captures tiling parameters, which are expected to be passed as arguments
+/// to the function enclosing the loop nest. Also checks if the required
+/// parameters are of index type. This approach is temporary for testing
+/// purposes.
+static void getTilingParameters(ArrayRef<AffineForOp> band,
+                                SmallVectorImpl<Value> &tilingParameters) {
+  AffineForOp topLoop = band[0];
+  Region *funcOpRegion = topLoop.getParentRegion();
+  unsigned nestDepth = band.size();
+
+  for (BlockArgument blockArgument :
+       funcOpRegion->getArguments().take_front(nestDepth)) {
+    if (blockArgument.getArgNumber() < nestDepth) {
+      assert(blockArgument.getType().isIndex() &&
+             "expected tiling parameters to be of index type.");
+      tilingParameters.push_back(blockArgument);
+    }
+  }
+}
+
+void TestAffineLoopParametricTiling::runOnFunction() {
+  // Bands of loops to tile.
+  std::vector<SmallVector<AffineForOp, 6>> bands;
+  getTileableBands(getFunction(), &bands);
+
+  // Tile each band.
+  for (SmallVectorImpl<AffineForOp> &band : bands) {
+    // Capture the tiling parameters from the arguments to the function
+    // enclosing this loop nest.
+    SmallVector<AffineForOp, 6> tiledNest;
+    SmallVector<Value, 6> tilingParameters;
+    // Check if tiling parameters are present.
+    checkIfTilingParametersExist(band);
+
+    // Get function arguments as tiling parameters.
+    getTilingParameters(band, tilingParameters);
+
+    if (failed(
+            tilePerfectlyNestedParametric(band, tilingParameters, &tiledNest)))
+      return signalPassFailure();
+  }
+}
+
+namespace mlir {
+void registerTestAffineLoopParametricTilingPass() {
+  PassRegistration<TestAffineLoopParametricTiling>(
+      "test-affine-parametric-tile",
+      "Tile affine loops using SSA values as tile sizes");
+}
+} // namespace mlir

mlir/tools/mlir-opt/mlir-opt.cpp

 void registerMemRefBoundCheck();
 void registerPassManagerTestPass();
 void registerPatternsTestPass();
 void registerPrintOpAvailabilityPass();
 void registerSideEffectTestPasses();
 void registerSimpleParametricTilingPass();
 void registerSymbolTestPasses();
 void registerTestAffineDataCopyPass();
+void registerTestAffineLoopParametricTilingPass();
 void registerTestAffineLoopUnswitchingPass();
 void registerTestAllReduceLoweringPass();
 void registerTestBufferPlacementPreparationPass();
 void registerTestCallGraphPass();
 void registerTestConstantFold();
 void registerTestConvVectorization();
 void registerTestConvertGPUKernelToCubinPass();
 void registerTestConvertGPUKernelToHsacoPass();
void registerTestPasses() {
   registerTestConvVectorization();
   registerTestConstantFold();
 #if MLIR_CUDA_CONVERSIONS_ENABLED
   registerTestConvertGPUKernelToCubinPass();
 #endif
 #if MLIR_ROCM_CONVERSIONS_ENABLED
   registerTestConvertGPUKernelToHsacoPass();
 #endif
+  registerTestAffineLoopParametricTilingPass();
   registerTestBufferPlacementPreparationPass();
   registerTestDominancePass();
   registerTestFunc();
   registerTestExpandTanhPass();
   registerTestGpuMemoryPromotionPass();
   registerTestInterfaces();
   registerTestLinalgHoisting();
   registerTestLinalgTransforms();