[MLIR][Linalg] Use GenericLoopNestRangeBuilder in tiling code.

[pifon2a]

Summary

Preparation for adding support for tiling to parallel loops.

[merge_guards_bot]

Unit tests: pass. 62414 tests passed, 0 failed and 839 were skipped.

clang-tidy: pass.

clang-format: pass.

Build artifacts: diff.json, clang-tidy.txt, clang-format.patch, CMakeCache.txt, console-log.txt, test-results.xml

Pre-merge checks is in beta. Report issue. Please join beta or enable it for your project.

[herhut]

Looks good to me. I wonder about users of TiledLinAlgOp. Are there none? This should be a breaking change.

@nicolasvasilache should take a look, too.

[nicolasvasilache]

Glanced over quickly, looks innocuous enough to me an a good step towards allow generalization to affine, various ploops (affine or not).
If all tests pass (i.e. the check-mlir target), then go for it, we'll adapt if any issues comes up.

Thanks for doing this generalization effort.

mlir/include/mlir/Dialect/Linalg/Utils/Utils.h

 /// `createAndFold` builder method. If `folder` is null, the regular `create`
 /// method is called.
 SmallVector<Value, 4> applyMapToValues(OpBuilder &b, Location loc,
                                        AffineMap map, ArrayRef<Value> values,
                                        OperationFolder *folder = nullptr);
 
 struct TiledLinalgOp {
   LinalgOp op;
-  SmallVector<loop::ForOp, 8> loops;
+  SmallVector<Operation *, 8> loops;
 };
 
 /// Performs standalone tiling of a single LinalgOp by `tileSizes`.
 /// and permute the loop nest according to `permutation`
 /// The permutation is expressed as a list of integers that specify
 /// the new ordering of the loop nest. The length of `permutation`
 /// must be equal to the length of `tileSizes`.
 /// E.g. the permutation `(i,j,k) -> (j,k,i)` will be expressed with

mlir/lib/Dialect/Linalg/Transforms/Tiling.cpp

makeTiledViews(OpBuilder &b, Location loc, LinalgOp linalgOp,
   if (folder)
     for (auto v : llvm::concat<Value>(lbs, subViewSizes))
       if (v.use_empty())
         v.getDefiningOp()->erase();
 
   return res;
 }
 
+template <typename LoopTy>
 Optional<TiledLinalgOp>
-mlir::linalg::tileLinalgOp(OpBuilder &b, LinalgOp op, ArrayRef<Value> tileSizes,
-                           ArrayRef<unsigned> permutation,
-                           OperationFolder *folder) {
+tileLinalgOpImpl(OpBuilder &b, LinalgOp op, ArrayRef<Value> tileSizes,
+                 ArrayRef<unsigned> permutation, OperationFolder *folder) {
   assert(op.hasBufferSemantics() && "expected linalg op with buffer semantics");
   // 1. Enforce the convention that "tiling by zero" skips tiling a particular
   // dimension. This convention is significantly simpler to handle instead of
   // adjusting affine maps to account for missing dimensions.
   assert(op.getNumParallelLoops() + op.getNumReductionLoops() +
                  op.getNumWindowLoops() ==
              tileSizes.size() &&
          "expected matching number of tile sizes and loops");
std::tie(loopRanges, loopIndexToRangeIndex) =
                           viewSizes, tileSizes, folder);
   if (!permutation.empty())
     applyPermutationToVector(loopRanges, permutation);
 
   // 3. Create the tiled loops.
   LinalgOp res = op;
   SmallVector<IndexHandle, 4> ivs(loopRanges.size());
   auto pivs = makeHandlePointers(MutableArrayRef<IndexHandle>(ivs));
-  LoopNestRangeBuilder(pivs, loopRanges)([&] {
+  // Convert SubViewOp::Range to linalg_range.
+  SmallVector<Value, 4> linalgRanges;
+  for (auto &range : loopRanges) {
+    linalgRanges.push_back(
+        linalg_range(range.offset, range.size, range.stride));
+  }
+  GenericLoopNestRangeBuilder<LoopTy>(pivs, linalgRanges)([&] {
     auto b = ScopedContext::getBuilder();
     auto loc = ScopedContext::getLocation();
     SmallVector<Value, 4> ivValues(ivs.begin(), ivs.end());
 
     // If we have to apply a permutation to the tiled loop nest, we have to
     // reorder the induction variables This permutation is the right one
     // assuming that loopRanges have previously been permuted by
     // (i,j,k)->(k,i,j) So this permutation should be the inversePermutation of
     // that one: (d0,d1,d2)->(d2,d0,d1)
     if (!permutation.empty())
       ivValues = applyMapToValues(b, loc, invPermutationMap, ivValues, folder);
 
     auto views =
         makeTiledViews(b, loc, op, ivValues, tileSizes, viewSizes, folder);
     auto operands = getAssumedNonViewOperands(op);
     views.append(operands.begin(), operands.end());
     res = op.clone(b, loc, views);
   });
 
   // 4. Transforms index arguments of `linalg.generic` w.r.t. to the tiling.
   transformIndexedGenericOpIndices(b, res, pivs, loopIndexToRangeIndex);
 
   // 5. Gather the newly created loops and return them with the new op.
-  SmallVector<ForOp, 8> loops;
+  SmallVector<Operation *, 8> loops;
   loops.reserve(ivs.size());
   for (auto iv : ivs)
     loops.push_back(loop::getForInductionVarOwner(iv));
 
   return TiledLinalgOp{res, loops};
 }
 
-Optional<TiledLinalgOp> mlir::linalg::tileLinalgOp(
-    OpBuilder &b, LinalgOp op, ArrayRef<int64_t> tileSizes,
+template <typename LoopTy>
+Optional<TiledLinalgOp>
+tileLinalgOpImpl(OpBuilder &b, LinalgOp op, ArrayRef<int64_t> tileSizes,
                  ArrayRef<unsigned> permutation, OperationFolder *folder) {
   assert(op.hasBufferSemantics() && "expected linalg op with buffer semantics");
   if (tileSizes.empty())
     return llvm::None;
 
   // The following uses the convention that "tiling by zero" skips tiling a
   // particular dimension. This convention is significantly simpler to handle
   // instead of adjusting affine maps to account for missing dimensions.
   auto nLoops = op.getNumParallelLoops() + op.getNumReductionLoops() +
tileLinalgOpImpl(OpBuilder &b, LinalgOp op, ArrayRef<int64_t> tileSizes,
   for (auto ts : tileSizes)
     tileSizeValues.push_back(constant_index(folder, ts));
   // Pad tile sizes with zero values to enforce our convention.
   if (tileSizeValues.size() < nLoops) {
     for (unsigned i = tileSizeValues.size(); i < nLoops; ++i)
       tileSizeValues.push_back(constant_index(folder, 0));
   }
 
-  return tileLinalgOp(b, op, tileSizeValues, permutation, folder);
+  return tileLinalgOpImpl<LoopTy>(b, op, tileSizeValues, permutation, folder);
+}
+
+Optional<TiledLinalgOp>
+mlir::linalg::tileLinalgOp(OpBuilder &b, LinalgOp op, ArrayRef<Value> tileSizes,
+                           ArrayRef<unsigned> permutation,
+                           OperationFolder *folder) {
+  return tileLinalgOpImpl<loop::ForOp>(b, op, tileSizes, permutation, folder);
 }
 
+Optional<TiledLinalgOp> mlir::linalg::tileLinalgOp(
+    OpBuilder &b, LinalgOp op, ArrayRef<int64_t> tileSizes,
+    ArrayRef<unsigned> permutation, OperationFolder *folder) {
+  return tileLinalgOpImpl<loop::ForOp>(b, op, tileSizes, permutation, folder);
+}
+
+template <typename LoopTy>
 static void tileLinalgOps(FuncOp f, ArrayRef<int64_t> tileSizes) {
   OpBuilder b(f);
   OperationFolder folder(f.getContext());
   f.walk([tileSizes, &b, &folder](LinalgOp op) {
     if (!op.hasBufferSemantics())
       return;
     auto opLoopsPair =
-        tileLinalgOp(b, op, tileSizes, /*permutation=*/{}, &folder);
+        tileLinalgOpImpl<LoopTy>(b, op, tileSizes, /*permutation=*/{}, &folder);
     // If tiling occurred successfully, erase old op.
     if (opLoopsPair)
       op.erase();
   });
   f.walk([](LinalgOp op) {
     if (!op.getOperation()->hasNoSideEffect())
       return;
     if (op.getOperation()->use_empty())
       op.erase();
   });
 }
 
 namespace {
-struct LinalgTilingPass : public FunctionPass<LinalgTilingPass> {
+
+template <typename LoopTy>
+struct LinalgTilingPass : public FunctionPass<LinalgTilingPass<LoopTy>> {
   LinalgTilingPass() = default;
-  LinalgTilingPass(ArrayRef<int64_t> sizes);
+  LinalgTilingPass(ArrayRef<int64_t> sizes) {
+    this->tileSizes.assign(sizes.begin(), sizes.end());
+  }
 
-  void runOnFunction() override { tileLinalgOps(getFunction(), tileSizes); }
+  void runOnFunction() override {
+    tileLinalgOps<LoopTy>(this->getFunction(), tileSizes);
+  }
 
   SmallVector<int64_t, 8> tileSizes;
 };
-} // namespace
 
-LinalgTilingPass::LinalgTilingPass(ArrayRef<int64_t> sizes) {
-  this->tileSizes.assign(sizes.begin(), sizes.end());
-}
+} // namespace
 
 std::unique_ptr<OpPassBase<FuncOp>>
 mlir::createLinalgTilingPass(ArrayRef<int64_t> tileSizes) {
-  return std::make_unique<LinalgTilingPass>(tileSizes);
+  return std::make_unique<LinalgTilingPass<loop::ForOp>>(tileSizes);
 }
 
-static PassRegistration<LinalgTilingPass>
+static PassRegistration<LinalgTilingPass<loop::ForOp>>
     pass("linalg-tile", "Tile operations in the linalg dialect", [] {
-      auto pass = std::make_unique<LinalgTilingPass>();
+      auto pass = std::make_unique<LinalgTilingPass<loop::ForOp>>();
       pass->tileSizes.assign(clTileSizes.begin(), clTileSizes.end());
       return pass;
     });