[mlir][linalg] Tile PadTensorOp

[springerm]

Summary

Tiling can be enabled with linalg-tile-pad-tensor-ops. Only scf::ForOp can be generated at the moment.

Depends On D105602

[springerm]

small update

[springerm]

more static types

[nicolasvasilache]

mlir/include/mlir/Dialect/Linalg/Passes.td

174

Can this be folded into the existing tiling pass ?

These are pure mechanics and don't have any policy / heuristics logic anyway.
They are not meant to drive anything other than tests.
Moving them to purely test patterns would also be appropriate.

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

47 ↗

(On Diff #356683)

s/asOpFoldResult/getAsOpFoldResult ?

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

158

s/writeTileToTensor/insertSliceIntoTensor?

362

The part that performs the tiling rewrite logic should be separated in a standalone transformation.

The pattern should call that transformation, similarly to other similar pattterns.
The part about using filters should be completely separated out and put in the same places as the other patterns that use filters.

The main reason is that these filters are meant to go away in the future when we have something better to control pattern composition.

389

This is a good observation that we have also discussed in other generalization contexts.

Viewing tiling as: given a subset of data, give me the subset of the op that produces the data, generalizes better as rewrite extract_slice(op(x, y, z)) into op(f(x), g(y), h(z)).
Bonus points if in a followup revision you are able to update linalg.generic tiling to operate similarly :) .

519

Can these just be folded in the pass above?

[springerm]

rebase

[springerm]

rebase

[springerm]

no change

[springerm]

address comments

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

362

I removed the filter and instead set an attribute manually. This attribute is needed, so that the pattern does not apply infinitely. Basically, this pattern just generates the loop structure. So once the structure is generated, it should not be generated a second time.

What do you mean by "tiling rewrite logic"? Do you mean how to rewrite the PadTensorOp inside the loop? That part is already as minimal as possible. Basically, I just move the PadTensorOp inside the loop nest (by cloning it and replacing the old one with the result of the loop nest). The actual rewrite logic specific to the PadTensorOp is handled by ExtractSliceOfPadTensorSwapPattern.

[springerm]

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

389

Probably requires a larger refactoring, but then it should be possible to use the exact same code path for generating the for-loop for each operation (PadTensorOp and LinalgOp). Looking into this afterwards.

[springerm]

address comments

[springerm]

Moved the tiling functionality into a separate function, as discussed.

[nicolasvasilache]

Cool thanks for this!

[springerm]

fix test

mlir/include/mlir/Dialect/Linalg/Passes.h

 
 std::unique_ptr<Pass> createLinalgFusionOfTensorOpsPass();
 std::unique_ptr<Pass> createFoldReshapeOpsByLinearizationPass();
 
 std::unique_ptr<OperationPass<FuncOp>>
 createLinalgTilingPass(ArrayRef<int64_t> tileSizes = {});
 
 std::unique_ptr<OperationPass<FuncOp>>
+createLinalgPadTensorTilingPass(ArrayRef<int64_t> tileSizes = {});
+
+std::unique_ptr<OperationPass<FuncOp>>
 createLinalgTilingToParallelLoopsPass(ArrayRef<int64_t> tileSizes = {});
 
 std::unique_ptr<OperationPass<FuncOp>>
 createLinalgTilingToTiledLoopPass(ArrayRef<int64_t> tileSizes = {},
                                   ArrayRef<StringRef> distributionTypes = {});
 
 std::unique_ptr<OperationPass<FuncOp>>
 createLinalgPromotionPass(bool dynamicBuffers, bool useAlloca);

mlir/include/mlir/Dialect/Linalg/Passes.td

let dependentDialects = [
     "scf::SCFDialect"
   ];
   let options = [
     ListOption<"tileSizes", "linalg-tile-sizes", "int64_t", "Tile sizes",
                "llvm::cl::ZeroOrMore, llvm::cl::MiscFlags::CommaSeparated">
   ];
 }
 
+def LinalgPadTensorOpTiling : FunctionPass<"linalg-tile-pad-tensor-ops"> {
+  let summary = "Tile linalg.pad_tensor operations";
+  let constructor = "mlir::createLinalgPadTensorTilingPass()";
+  let dependentDialects = [
+    "AffineDialect",
+    "linalg::LinalgDialect",
+    "memref::MemRefDialect",
+    "scf::SCFDialect"
+  ];
+  let options = [
+    ListOption<"tileSizes", "linalg-tile-sizes", "int64_t", "Tile sizes",
+               "llvm::cl::ZeroOrMore, llvm::cl::MiscFlags::CommaSeparated">
+  ];
+}
+
 def LinalgTilingToParallelLoops
     : FunctionPass<"linalg-tile-to-parallel-loops"> {
   let summary = "Tile operations in the linalg dialect to parallel loops";
   let constructor = "mlir::createLinalgTilingToParallelLoopsPass()";
   let options = [
     ListOption<"tileSizes", "linalg-tile-sizes", "int64_t", "Tile sizes",
                "llvm::cl::ZeroOrMore, llvm::cl::MiscFlags::CommaSeparated">
   ];

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

AffineApplyOp applyOp = b.create<AffineApplyOp>(
         ValueRange{indexOp.getResult(), ivs[rangeIndex->second]});
     indexOp.getResult().replaceAllUsesExcept(applyOp, applyOp);
   }
 }
 
 template <typename LoopTy>
 static Optional<TiledLinalgOp>
 tileLinalgOpImpl(OpBuilder &b, LinalgOp op, ValueRange tileSizes,
                  const LinalgTilingOptions &options) {
   auto nLoops = op.getNumLoops();
   // Initial tile sizes may be too big, only take the first nLoops.
   tileSizes = tileSizes.take_front(nLoops);
 
   if (llvm::all_of(tileSizes, isZero))
     return llvm::None;
 
   if (auto convOp = dyn_cast<linalg::ConvOp>(op.getOperation())) {
mlir::linalg::tileLinalgOp(OpBuilder &b, LinalgOp op,
   case LinalgTilingLoopType::TiledLoops:
     return tileLinalgOpImpl<linalg::TiledLoopOp>(b, op, options);
   default:;
   }
   return llvm::None;
 }
 
 namespace {
+struct PadTensorOpTilingPattern : public OpRewritePattern<PadTensorOp> {
+  PadTensorOpTilingPattern(MLIRContext *ctx, LinalgTilingOptions opt)
+      : OpRewritePattern<PadTensorOp>(ctx), options(opt),
+        filter(LinalgTransformationFilter(ArrayRef<Identifier>{},
+                                          Identifier::get("tiled", ctx))) {}
+
+  LogicalResult matchAndRewrite(PadTensorOp op,
+                                PatternRewriter &rewriter) const override {
+    // Can tile only PadTensorOp that have an output operand.
+    if (!op.output())
+      return failure();
+    if (failed(filter.checkAndNotify(rewriter, op)))
+      return failure();
+    Location loc = op.getLoc();
+    OpBuilder::InsertionGuard g(rewriter);
+    rewriter.setInsertionPoint(op);
+
+    // Clone PadTensorOp so that the existing op can be replaced more easily.
+    Operation *newPadOp = rewriter.clone(*op.getOperation());
+    // Get rank and tile sizes.
+    int64_t rank = op.getResultType().getRank();
+    SmallVector<Value> tileSizes =
+        options.tileSizeComputationFunction(rewriter, op);
+    assert(tileSizes.size() == rank);
+    // Compute lower and upper bounds of the loop nest.
+    SmallVector<Value> lbs, dims, steps;
+    for (int64_t i = 0; i < rank; ++i) {
+      if (!isZero(tileSizes[i])) {
+        lbs.push_back(rewriter.create<ConstantIndexOp>(loc, 0));
+        dims.push_back(rewriter.create<tensor::DimOp>(loc, op.output(), i));
+        steps.push_back(tileSizes[i]);
+      }
+    }
+    // Generate loop nest: One loop per dimension.
+    LoopNest loopNest = mlir::scf::buildLoopNest(
+        rewriter, loc, lbs, /*ubs=*/dims, steps, ValueRange(op.output()),
+        [&](OpBuilder &b, Location loc, ValueRange localIvs,
+            ValueRange iterArgs) -> scf::ValueVector {
+          // Compute offsets and sizes of ExtractSliceOp.
+          SmallVector<Value> offsets =
+              computeTileOffsets(b, loc, localIvs, tileSizes);
+          SmallVector<Value> sizes =
+              computeTileSizes(b, loc, localIvs, tileSizes, dims);
+          // Create ExtractSliceOp: Extract a tile from the PadTensorOp.
+          // Note: The PadTensorOp is located outside of the loop nest. It is
+          // later moved inside by ExtractSliceOfPadTensorSwapPattern.
+          auto map = AffineMap::getMultiDimIdentityMap(rank, b.getContext());
+          Value tiledOutput = makeTiledShape(b, loc, newPadOp->getResult(0),
+                                             tileSizes, map, offsets, sizes);
+          auto sliceOp = tiledOutput.getDefiningOp<tensor::ExtractSliceOp>();
+          assert(sliceOp && "expected ExtractSliceOp");
+          // Insert the tile into the output tensor.
+          Value yieldValue = b.create<tensor::InsertSliceOp>(
+              loc, sliceOp.source().getType(), /*source=*/sliceOp,
+              /*dest=*/iterArgs[0], sliceOp.offsets(), sliceOp.sizes(),
+              sliceOp.strides(), sliceOp.static_offsets(),
+              sliceOp.static_sizes(), sliceOp.static_strides());
+          return scf::ValueVector({yieldValue});
+        });
+    // Replace all uses of the original PadTensorOp.
+    rewriter.replaceOp(op, loopNest.getResults()[0]);
+    filter.replaceLinalgTransformationFilter(rewriter, newPadOp);
+    return success();
+  }
+
+  LinalgTilingOptions options;
+  LinalgTransformationFilter filter;
+};
+} // namespace
+
+namespace {
 /// Helper classes for type list expansion.
 template <typename... OpTypes>
 class CanonicalizationPatternList;
 
 template <>
 class CanonicalizationPatternList<> {
 public:
   static void insert(RewritePatternSet &patterns) {}
void mlir::linalg::populateLinalgTilingCanonicalizationPatterns(
   AffineMaxOp::getCanonicalizationPatterns(patterns, ctx);
   scf::ForOp::getCanonicalizationPatterns(patterns, ctx);
   scf::ParallelOp::getCanonicalizationPatterns(patterns, ctx);
   ConstantIndexOp::getCanonicalizationPatterns(patterns, ctx);
   tensor::ExtractSliceOp::getCanonicalizationPatterns(patterns, ctx);
   memref::SubViewOp::getCanonicalizationPatterns(patterns, ctx);
   tensor::CastOp::getCanonicalizationPatterns(patterns, ctx);
   memref::ViewOp::getCanonicalizationPatterns(patterns, ctx);
+  PadTensorOp::getCanonicalizationPatterns(patterns, ctx);
   ctx->getLoadedDialect<LinalgDialect>()->getCanonicalizationPatterns(patterns);
   CanonicalizationPatternList<
 #define GET_OP_LIST
 #include "mlir/Dialect/Linalg/IR/LinalgStructuredOps.cpp.inc"
       >::insert(patterns);
 }
 
 /// Populate the given list with patterns that apply Linalg tiling.
applyTilingToLoopPatterns(LinalgTilingLoopType loopType, FuncOp funcOp,
   (void)applyPatternsAndFoldGreedily(
       funcOp, getLinalgTilingCanonicalizationPatterns(ctx));
   // Drop the marker.
   funcOp.walk([](LinalgOp op) {
     op->removeAttr(LinalgTransforms::kLinalgTransformMarker);
   });
 }
 
+static void applyPadTensorOpTilingToLoopPatterns(
+    LinalgTilingLoopType loopType, FuncOp funcOp, ArrayRef<int64_t> tileSizes,
+    ArrayRef<StringRef> distributionTypes = {}) {
+  auto options = LinalgTilingOptions()
+                     .setTileSizes(tileSizes)
+                     .setLoopType(loopType)
+                     .setDistributionTypes(distributionTypes);
+  MLIRContext *ctx = funcOp.getContext();
+  RewritePatternSet patterns(ctx);
+  patterns.add<PadTensorOpTilingPattern>(patterns.getContext(), options);
+  patterns.add<ExtractSliceOfPadTensorSwapPattern>(patterns.getContext());
+  (void)applyPatternsAndFoldGreedily(funcOp, std::move(patterns));
+  (void)applyPatternsAndFoldGreedily(
+      funcOp, getLinalgTilingCanonicalizationPatterns(ctx));
+  // Drop the marker.
+  funcOp.walk([](PadTensorOp op) {
+    op->removeAttr(LinalgTransforms::kLinalgTransformMarker);
+  });
+}
+
 namespace {
 struct LinalgTilingPass : public LinalgTilingBase<LinalgTilingPass> {
   LinalgTilingPass() = default;
   LinalgTilingPass(ArrayRef<int64_t> sizes) { tileSizes = sizes; }
 
   void runOnFunction() override {
     applyTilingToLoopPatterns(LinalgTilingLoopType::Loops, getFunction(),
                               tileSizes);
   }
 };
 
+struct LinalgPadTensorOpTilingPass
+    : public LinalgPadTensorOpTilingBase<LinalgPadTensorOpTilingPass> {
+  LinalgPadTensorOpTilingPass() = default;
+  LinalgPadTensorOpTilingPass(ArrayRef<int64_t> sizes) { tileSizes = sizes; }
+
+  void runOnFunction() override {
+    applyPadTensorOpTilingToLoopPatterns(LinalgTilingLoopType::Loops,
+                                         getFunction(), tileSizes);
+  }
+};
+
 struct LinalgTilingToParallelLoopsPass
     : public LinalgTilingToParallelLoopsBase<LinalgTilingToParallelLoopsPass> {
   LinalgTilingToParallelLoopsPass() = default;
   LinalgTilingToParallelLoopsPass(ArrayRef<int64_t> sizes) {
     tileSizes = sizes;
   }
 
   void runOnFunction() override {

 } // namespace
 
 std::unique_ptr<OperationPass<FuncOp>>
 mlir::createLinalgTilingPass(ArrayRef<int64_t> tileSizes) {
   return std::make_unique<LinalgTilingPass>(tileSizes);
 }
 
 std::unique_ptr<OperationPass<FuncOp>>
+mlir::createLinalgPadTensorTilingPass(ArrayRef<int64_t> tileSizes) {
+  return std::make_unique<LinalgPadTensorOpTilingPass>(tileSizes);
+}
+
+std::unique_ptr<OperationPass<FuncOp>>
 mlir::createLinalgTilingToParallelLoopsPass(ArrayRef<int64_t> tileSizes) {
   return std::make_unique<LinalgTilingToParallelLoopsPass>(tileSizes);
 }
 
 std::unique_ptr<OperationPass<FuncOp>>
 mlir::createLinalgTilingToTiledLoopPass(ArrayRef<int64_t> tileSizes,
                                         ArrayRef<StringRef> distributionTypes) {
   return std::make_unique<LinalgTilingToTiledLoopsPass>(tileSizes,
                                                         distributionTypes);
 }

mlir/test/Dialect/Linalg/tile-pad-tensor-op.mlir

+// RUN: mlir-opt %s -linalg-tile-pad-tensor-ops="linalg-tile-sizes=2,3" -cse -split-input-file | \
+// RUN: FileCheck %s -check-prefix=TILE2
+// RUN: mlir-opt %s -linalg-tile-pad-tensor-ops="linalg-tile-sizes=0,3" -cse -split-input-file | \
+// RUN: FileCheck %s -check-prefix=TILE1
+
+// TILE2-LABEL: func @dynamic_pad_tensor(
+//  TILE2-SAME:     %[[IN:.*]]: tensor<?x?xf32>, %[[OUT:.*]]: tensor<?x?xf32>
+//   TILE2-DAG:   %[[C0:.*]] = constant 0 : index
+//   TILE2-DAG:   %[[C1:.*]] = constant 1 : index
+//   TILE2-DAG:   %[[C2:.*]] = constant 2 : index
+//   TILE2-DAG:   %[[C3:.*]] = constant 3 : index
+//       TILE2:   %[[DIM0:.*]] = tensor.dim %[[OUT]], %[[C0]]
+//       TILE2:   %[[DIM1:.*]] = tensor.dim %[[OUT]], %[[C1]]
+//       TILE2:   %[[RESULT:.*]] = scf.for {{.*}} = %[[C0]] to %[[DIM0]] step %[[C2]]
+//       TILE2:     scf.for {{.*}} = %[[C0]] to %[[DIM1]] step %[[C3]] iter_args(%[[INNER_OUT:.*]] =
+//       TILE2:       %[[SWAP_RESULT:.*]] = scf.if
+//       TILE2:         tensor.generate
+//       TILE2:       else
+//       TILE2:         %[[SLICE:.*]] = tensor.extract_slice %[[IN]][{{.*}}, {{.*}}] [{{.*}}, {{.*}}] [1, 1]
+//       TILE2:         %[[PAD:.*]] = linalg.pad_tensor %[[SLICE]]
+//       TILE2:       tensor.insert_slice %[[SWAP_RESULT]] into %[[INNER_OUT]][{{.*}}, {{.*}}] [{{.*}}, {{.*}}] [1, 1]
+//       TILE2:   return %[[RESULT]]
+
+// TILE1-LABEL: func @dynamic_pad_tensor(
+//  TILE1-SAME:     %[[IN:.*]]: tensor<?x?xf32>, %[[OUT:.*]]: tensor<?x?xf32>
+//   TILE1-DAG:   %[[C0:.*]] = constant 0 : index
+//   TILE1-DAG:   %[[C1:.*]] = constant 1 : index
+//   TILE1-DAG:   %[[C3:.*]] = constant 3 : index
+//       TILE1:   %[[DIM1:.*]] = tensor.dim %[[OUT]], %[[C1]]
+//       TILE1:   %[[RESULT:.*]] = scf.for {{.*}} = %[[C0]] to %[[DIM1]] step %[[C3]] iter_args(%[[INNER_OUT:.*]] =
+//       TILE1:     %[[DIM0:.*]] = tensor.dim %[[OUT]], %[[C0]]
+//       TILE1:     %[[SWAP_RESULT:.*]] = scf.if
+//       TILE1:       tensor.generate
+//       TILE1:     else
+//       TILE1:       %[[SLICE:.*]] = tensor.extract_slice %[[IN]][{{.*}}, {{.*}}] [{{.*}}, {{.*}}] [1, 1]
+//       TILE1:       %[[PAD:.*]] = linalg.pad_tensor %[[SLICE]] low[3, %{{.*}}] high[{{.*}}, {{.*}}]
+//       TILE1:     tensor.insert_slice %[[SWAP_RESULT]] into %[[INNER_OUT]][0, {{.*}}] [%[[DIM0]], {{.*}}] [1, 1]
+//       TILE1:   return %[[RESULT]]
+
+func @dynamic_pad_tensor(%input_tensor: tensor<?x?xf32>,
+                         %output_tensor: tensor<?x?xf32>,
+                         %pad_value: f32) -> tensor<?x?xf32> {
+  %0 = linalg.pad_tensor %input_tensor, %output_tensor
+    low[3, 4] high[5, 3] {
+    ^bb0(%arg1: index, %arg2: index):
+      linalg.yield %pad_value : f32
+    } : tensor<?x?xf32> to tensor<?x?xf32>
+  return %0 : tensor<?x?xf32>
+}
+
+// -----
+
+// TILE2-LABEL: func @static_pad_tensor(
+//  TILE2-SAME:     %[[IN:.*]]: tensor<7x9xf32>, %[[OUT:.*]]: tensor<15x16xf32>
+//   TILE2-DAG:   %[[C0:.*]] = constant 0 : index
+//   TILE2-DAG:   %[[C2:.*]] = constant 2 : index
+//   TILE2-DAG:   %[[C3:.*]] = constant 3 : index
+//   TILE2-DAG:   %[[C15:.*]] = constant 15 : index
+//   TILE2-DAG:   %[[C16:.*]] = constant 16 : index
+//       TILE2:   %[[RESULT:.*]] = scf.for {{.*}} = %[[C0]] to %[[C15]] step %[[C2]]
+//       TILE2:     scf.for {{.*}} = %[[C0]] to %[[C16]] step %[[C3]] iter_args(%[[INNER_OUT:.*]] =
+//       TILE2:       %[[SWAP_RESULT:.*]] = scf.if
+//       TILE2:         tensor.generate
+//       TILE2:       else
+//       TILE2:         %[[SLICE:.*]] = tensor.extract_slice %[[IN]][{{.*}}, {{.*}}] [{{.*}}, {{.*}}] [1, 1]
+//       TILE2:         %[[PAD:.*]] = linalg.pad_tensor %[[SLICE]]
+//       TILE2:       tensor.insert_slice %[[SWAP_RESULT]] into %[[INNER_OUT]][{{.*}}, {{.*}}] [{{.*}}, {{.*}}] [1, 1]
+//       TILE2:   return %[[RESULT]]
+
+
+// TILE1-LABEL: func @static_pad_tensor(
+//  TILE1-SAME:     %[[IN:.*]]: tensor<7x9xf32>, %[[OUT:.*]]: tensor<15x16xf32>
+//   TILE1-DAG:   %[[C0:.*]] = constant 0 : index
+//   TILE1-DAG:   %[[C3:.*]] = constant 3 : index
+//   TILE1-DAG:   %[[C15:.*]] = constant 15 : index
+//   TILE1-DAG:   %[[C16:.*]] = constant 16 : index
+//       TILE1:   %[[RESULT:.*]] = scf.for {{.*}} = %[[C0]] to %[[C16]] step %[[C3]] iter_args(%[[INNER_OUT:.*]] =
+//       TILE1:     %[[SWAP_RESULT:.*]] = scf.if
+//       TILE1:       tensor.generate
+//       TILE1:     else
+//       TILE1:       %[[SLICE:.*]] = tensor.extract_slice %[[IN]][0, {{.*}}] [7, {{.*}}] [1, 1]
+//       TILE1:       %[[PAD:.*]] = linalg.pad_tensor %[[SLICE]] low[3, %{{.*}}] high[5, {{.*}}]
+//       TILE1:     tensor.insert_slice %[[SWAP_RESULT]] into %[[INNER_OUT]][0, {{.*}}] [%[[C15]], {{.*}}] [1, 1]
+//       TILE1:   return %[[RESULT]]
+
+func @static_pad_tensor(%input_tensor: tensor<7x9xf32>,
+                        %output_tensor: tensor<15x16xf32>,
+                        %pad_value: f32) -> tensor<15x16xf32> {
+  %0 = linalg.pad_tensor %input_tensor, %output_tensor
+    low[3, 4] high[5, 3] {
+    ^bb0(%arg1: index, %arg2: index):
+      linalg.yield %pad_value : f32
+    } : tensor<7x9xf32> to tensor<15x16xf32>
+  return %0 : tensor<15x16xf32>
+}