[MLIR] NFC affine for op tiling cleanup / utility rename

[bondhugula]

Summary

Rename mlir::tileCodeGen -> mlir::tilePerfectlyNested to be consistent.
NFC clean up tiling utility code, drop dead code, better comments.
Expose isPerfectlyNested and reuse.

[bondhugula]

Rebase - expose and use isPerfectlyNested utility

[rriddle]

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

96–97

Can you just use OperandRange instead?

[bondhugula]

Use OperandRange

[bondhugula]

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

96–97

Thanks for catching this.

mlir/include/mlir/Transforms/LoopUtils.h

 /// if the loop cannot be unrolled either due to restrictions or due to invalid
 /// unroll factors.
 LogicalResult loopUnrollByFactor(AffineForOp forOp, uint64_t unrollFactor);
 
 /// Unrolls this loop by the specified unroll factor or its trip count,
 /// whichever is lower.
 LogicalResult loopUnrollUpToFactor(AffineForOp forOp, uint64_t unrollFactor);
 
+/// Returns true if `loops` is a perfectly nested loop nest, where loops appear
+/// in it from outermost to innermost.
+bool LLVM_ATTRIBUTE_UNUSED isPerfectlyNested(ArrayRef<AffineForOp> loops);
+
 /// Get perfectly nested sequence of loops starting at root of loop nest
 /// (the first op being another AffineFor, and the second op - a terminator).
 /// A loop is perfectly nested iff: the first op in the loop's body is another
 /// AffineForOp, and the second op is a terminator).
 void getPerfectlyNestedLoops(SmallVectorImpl<AffineForOp> &nestedLoops,
                              AffineForOp root);
 void getPerfectlyNestedLoops(SmallVectorImpl<loop::ForOp> &nestedLoops,
                              loop::ForOp root);

 /// loops when possible.
 LLVM_NODISCARD
 LogicalResult affineForOpBodySkew(AffineForOp forOp, ArrayRef<uint64_t> shifts,
                                   bool unrollPrologueEpilogue = false);
 
 /// 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.
 LLVM_NODISCARD
-LogicalResult tileCodeGen(MutableArrayRef<AffineForOp> band,
-                          ArrayRef<unsigned> tileSizes,
-                          SmallVectorImpl<AffineForOp> *tiledNest = nullptr);
+LogicalResult
+tilePerfectlyNested(MutableArrayRef<AffineForOp> input,
 
 /// 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).

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

constructTiledIndexSetHyperRect(MutableArrayRef<AffineForOp> origLoops,
   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++) {
-    auto lbOperands = origLoops[i].getLowerBoundOperands();
-    auto ubOperands = origLoops[i].getUpperBoundOperands();
-    SmallVector<Value, 4> newLbOperands(lbOperands);
-    SmallVector<Value, 4> newUbOperands(ubOperands);
+    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.hasValue() &&
-        mayBeConstantCount.getValue() < tileSizes[i]) {
+    if (mayBeConstantCount && mayBeConstantCount.getValue() < tileSizes[i]) {
       // Trip count is less than tile size; upper bound is the trip count.
       auto ubMap = b.getConstantAffineMap(mayBeConstantCount.getValue());
       newLoops[width + i].setUpperBoundMap(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.
-      auto ub = origLoops[i].getUpperBound();
+
+      // 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();
-      // Add dim operands from original upper bound.
-      for (unsigned j = 0, e = origUbMap.getNumDims(); j < e; ++j) {
+      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) {
+      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);
     }
   }
 }
 
 /// 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(bondhugula): handle non hyper-rectangular spaces.
-LogicalResult mlir::tileCodeGen(MutableArrayRef<AffineForOp> band,
-                                ArrayRef<unsigned> tileSizes,
-                                SmallVectorImpl<AffineForOp> *tiledNest) {
+LogicalResult
+mlir::tilePerfectlyNested(MutableArrayRef<AffineForOp> input,
   // Check if the supplied for op's are all successively nested.
-  assert(!band.empty() && "no loops in band");
-  assert(band.size() == tileSizes.size() && "Too few/many tile sizes");
+  assert(!input.empty() && "no loops in input band");
+  assert(input.size() == tileSizes.size() && "Too few/many tile sizes");
 
-  for (unsigned i = 1, e = band.size(); i < e; i++)
-    assert(band[i].getParentOp() == band[i - 1] && "not a perfect nest / band");
+  assert(isPerfectlyNested(input) && "input loops not perfectly nested");
 
-  auto origLoops = band;
+  auto origLoops = input;
 
   AffineForOp rootAffineForOp = origLoops[0];
   auto loc = rootAffineForOp.getLoc();
   // Note that width is at least one since band isn't empty.
-  unsigned width = band.size();
+  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 = width; i < 2 * width; i++) {
     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[origLoops.size() - 1], innermostPointLoop);
+  moveLoopBody(origLoops.back(), innermostPointLoop);
 
   SmallVector<Value, 8> origLoopIVs;
-  extractForInductionVars(band, &origLoopIVs);
-  SmallVector<Optional<Value>, 6> ids(origLoopIVs.begin(), origLoopIVs.end());
-  FlatAffineConstraints cst;
-  getIndexSet(band, &cst);
+  extractForInductionVars(input, &origLoopIVs);
 
+  FlatAffineConstraints cst;
+  getIndexSet(input, &cst);
   if (!cst.isHyperRectangular(0, width)) {
     llvm::dbgs() << "tiled code generation unimplemented for the "
                     "non-hyperrectangular case, op:"
                  << *rootAffineForOp << "\n";
     return failure();
   }
 
   constructTiledIndexSetHyperRect(origLoops, tiledLoops, tileSizes);
static void getTileableBands(FuncOp f,
   };
 
   for (auto &block : f)
     for (auto &op : block)
       if (auto forOp = dyn_cast<AffineForOp>(op))
         getMaximalPerfectLoopNest(forOp);
 }
 
-// Reduce each tile size to the largest divisor of the corresponding trip count
-// (if the trip count is known).
+/// 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]);
-    if (!mayConst.hasValue())
+    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)
       tSizeAdjusted--;
   }
 }
 
 // Returns tile sizes to use. Checks CL options; if none are specified, sets it
 // based on a simple model that looks at the memory footprint and determines
 // tile sizes assuming identity accesses / 1:1 tile size proportional footprint
 // along each of the dimensions being tiled.
 // TODO(mlir-team): evolve this model. Tile size determination is a large area
 // to play with in general.
 void LoopTiling::getTileSizes(ArrayRef<AffineForOp> band,
                               SmallVectorImpl<unsigned> *tileSizes) {
   if (band.empty())
     return;
 
-  // Use tileSize for all loops if specified.
-  if (tileSize.hasValue()) {
+  // Use command-line tileSize for all loops if specified.
+  if (tileSize) {
     tileSizes->assign(band.size(), tileSize);
     return;
   }
 
   // Use tileSizes and fill them with default tile size if it's short.
   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];
   (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);
-  if (!fp.hasValue()) {
+  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 "
                               "sizes adjusted to trip count divisors"));
void LoopTiling::getTileSizes(ArrayRef<AffineForOp> band,
   }
 
   // Divide all loops equally in an attempt to reduce footprint.
   // TODO(bondhugula): this is approximate. Ideally, obtain reuse factor /
   // profitability along each dimension and weight tile sizes based on that as
   // one possible approach. Or compute a polynomial in tile sizes and solve for
   // it.
 
-  // For an n-d tileable band, compute n^th root of the excess.
+  // For an n-d tileable band, compute the n^th root of the excess.
   unsigned tSize =
       static_cast<unsigned>(floorl(std::pow(excessFactor, 1.0 / band.size())));
   // We'll keep a running product to determine the last tile size better.
   unsigned cumulProductOfTileSizes = 1;
   for (unsigned i = 0, e = band.size(); i < e; i++) {
     if (i < e - 1)
       (*tileSizes)[i] = tSize;
     else
for (auto &band : bands) {
     getTileSizes(band, &tileSizes);
     if (llvm::DebugFlag) {
       auto diag = band[0].emitRemark("using tile sizes [");
       for (auto tSize : tileSizes)
         diag << tSize << ' ';
       diag << "]\n";
     }
     SmallVector<AffineForOp, 6> tiledNest;
-    if (failed(tileCodeGen(band, tileSizes, &tiledNest)))
+    if (failed(tilePerfectlyNested(band, tileSizes, &tiledNest)))
       return signalPassFailure();
 
     // Separate full and partial tiles.
     if (separate) {
       auto intraTileLoops =
           MutableArrayRef<AffineForOp>(tiledNest).drop_front(band.size());
       separateFullTiles(intraTileLoops);
     }
   }
 }
 
 constexpr unsigned LoopTiling::kDefaultTileSize;

mlir/lib/Transforms/Utils/LoopUtils.cpp

bool mlir::isValidLoopInterchangePermutation(ArrayRef<AffineForOp> loops,
   unsigned maxLoopDepth = loops.size();
   std::vector<SmallVector<DependenceComponent, 2>> depCompsVec;
   getDependenceComponents(loops[0], maxLoopDepth, &depCompsVec);
   return checkLoopInterchangeDependences(depCompsVec, loops, loopPermMap);
 }
 
 /// Returns true if `loops` is a perfectly nested loop nest, where loops appear
 /// in it from outermost to innermost.
-static bool LLVM_ATTRIBUTE_UNUSED
-isPerfectlyNested(ArrayRef<AffineForOp> loops) {
+bool LLVM_ATTRIBUTE_UNUSED
+mlir::isPerfectlyNested(ArrayRef<AffineForOp> loops) {
   assert(!loops.empty() && "no loops provided");
 
   // We already know that the block can't be empty.
   auto hasTwoElements = [](Block *block) {
     auto secondOpIt = std::next(block->begin());
     return secondOpIt != block->end() && &*secondOpIt == &block->back();
   };