[mlir][scf] Enhance sizes computation in tileUsingSCFForOp.

[hanchung]

Summary

The boundary is always 1 if the tile size is 1.

[mravishankar]

Just a minor readability comment.

mlir/lib/Dialect/SCF/Transforms/TileUsingInterface.cpp

171

The if-then-else-if logic is a bit hard to follow. Maybe move that into a helper method like

OpFoldResult getBoundedTileSize(OpBuilder &b, Value tileSize) {
  Optional<int64_t> ts = getConstantIntValue(tileSize)
  if (ts && ts.value() == 1)
    return getAsOpFoldResult(tileSize);
  if (tileDividesIterationDomain(...))
    return tileSize;
  return b.create<AffineMinOp>(...);
}

[hanchung]

create a helper function

[hanchung]

clang-format

mlir/lib/Dialect/SCF/Transforms/TileUsingInterface.cpp

static bool tileDividesIterationDomain(Range loopRange) {
   if (!sizeAsInt)
     return false;
   Optional<int64_t> strideAsInt = getConstantIntValue(loopRange.stride);
   if (!strideAsInt)
     return false;
   return ((sizeAsInt.value() - offsetAsInt.value()) % strideAsInt.value() == 0);
 }
 
+/// Returns the bounded tile size given the current `iv`, `loopRange` and
+/// `tileSize`, i.e., `min(tileSize, range.end() - iv)`.
+static OpFoldResult getBoundedTileSize(OpBuilder &b, Location loc,
+                                       Range loopRange, Value iv,
+                                       Value tileSize) {
+  Optional<int64_t> ts = getConstantIntValue(tileSize);
+  if (ts && ts.value() == 1)
+    return getAsOpFoldResult(tileSize);
+
+  if (tileDividesIterationDomain(
+          Range{loopRange.offset, loopRange.size, tileSize}))
+    return tileSize;
+
+  // The tile size to use (to avoid out of bounds access) is  minimum of
+  // `tileSize` and `ub - iv`, where `iv` is the induction variable of the tiled
+  // loop.
+  AffineExpr s0, s1, d0;
+  bindDims(b.getContext(), d0);
+  bindSymbols(b.getContext(), s0, s1);
+  AffineMap minMap = AffineMap::get(1, 2, {s0, s1 - d0}, b.getContext());
+  Value size = getValueOrCreateConstantIndexOp(b, loc, loopRange.size);
+  return b.create<AffineMinOp>(loc, minMap, ValueRange{iv, tileSize, size})
+      .getResult();
+}
+
 /// Generate an empty loop nest that represents the tiled loop nest shell.
 /// - `loopRanges` specifies the lb, ub and step of the untiled iteration space.
 /// - `tileSizeVals` is the tile sizes to use. Zero represent untiled loops.
 /// - In `offsets` and `sizes` return the multi-dimensional offset and size of
 /// the
 ///   tile processed within the inner most loop.
 static SmallVector<scf::ForOp>
 generateTileLoopNest(OpBuilder &builder, Location loc,
                      ArrayRef<Range> loopRanges, ArrayRef<Value> tileSizeVals,
                      SmallVector<OpFoldResult> &offsets,
                      SmallVector<OpFoldResult> &sizes) {
   assert(!loopRanges.empty() && "expected at least one loop range");
   assert(loopRanges.size() == tileSizeVals.size() &&
          "expected as many tile sizes as loop ranges");
   OpBuilder::InsertionGuard guard(builder);
   SmallVector<scf::ForOp> loops;
   offsets.resize(loopRanges.size());
   sizes.resize(loopRanges.size());
 
-  // The tile size to use (to avoid out of bounds access) is  minimum of
-  // `tileSize` and `ub - iv`, where `iv` is the induction variable
-  // of the tiled loop.
-  AffineExpr s0, s1, d0;
-  bindDims(builder.getContext(), d0);
-  bindSymbols(builder.getContext(), s0, s1);
-  AffineMap minMap = AffineMap::get(1, 2, {s0, s1 - d0}, builder.getContext());
 
   for (auto loopRange : llvm::enumerate(loopRanges)) {
     Value offset =
         getValueOrCreateConstantIndexOp(builder, loc, loopRange.value().offset);
     Value size =
         getValueOrCreateConstantIndexOp(builder, loc, loopRange.value().size);
+    Value tileSize = tileSizeVals[loopRange.index()];
     // No loops if tile size is zero. Set offset and size to the loop
     // offset and size.
-    if (matchPattern(tileSizeVals[loopRange.index()], m_Zero())) {
+    if (matchPattern(tileSize, m_Zero())) {
       offsets[loopRange.index()] = offset;
       sizes[loopRange.index()] = size;
       continue;
     }
 
     auto loop = builder.create<scf::ForOp>(
-        loc, offset, size, tileSizeVals[loopRange.index()], ValueRange{},
+        loc, offset, size, tileSize, ValueRange{},
         [&](OpBuilder &bodyBuilder, Location bodyLoc, Value iv,
             ValueRange /*iterArgs*/) {
-          bool canAvoidMap = tileDividesIterationDomain(
-              Range{loopRange.value().offset, loopRange.value().size,
-                    tileSizeVals[loopRange.index()]});
-          Value boundedTileSize =
-              (canAvoidMap)
-                  ? tileSizeVals[loopRange.index()]
-                  : builder.create<AffineMinOp>(
-                        bodyLoc, minMap,
-                        ValueRange{iv, tileSizeVals[loopRange.index()], size});
-          sizes[loopRange.index()] = boundedTileSize;
+          sizes[loopRange.index()] = getBoundedTileSize(
+              bodyBuilder, bodyLoc, loopRange.value(), iv, tileSize);
           builder.create<scf::YieldOp>(loc);
         });
     offsets[loopRange.index()] = loop.getInductionVar();
     loops.push_back(loop);
     builder.setInsertionPoint(loop.getBody()->getTerminator());
   }
   return loops;
 }