[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 ↗

(On Diff #356683)

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?

371

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.

398

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 :) .

527

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

371

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

398

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/lib/Dialect/Linalg/Transforms/Tiling.cpp

for (IndexOp indexOp : block.getOps<linalg::IndexOp>()) {
     bindDims(b.getContext(), index, iv);
     AffineApplyOp applyOp = b.create<AffineApplyOp>(
         indexOp.getLoc(), index + iv,
         ValueRange{indexOp.getResult(), ivs[rangeIndex->second]});
     indexOp.getResult().replaceAllUsesExcept(applyOp, applyOp);
   }
 }
 
+// Insert a tile `source` into the destination tensor `dest`. The position at
+// which the tile is inserted (as well as size of tile) is taken from a given
+// ExtractSliceOp `sliceOp`.
+static Value insertSliceIntoTensor(OpBuilder &b, Location loc,
+                                   tensor::ExtractSliceOp sliceOp, Value source,
+                                   Value dest) {
+  return b.create<tensor::InsertSliceOp>(
+      loc, sliceOp.source().getType(), source, dest, sliceOp.offsets(),
+      sliceOp.sizes(), sliceOp.strides(), sliceOp.static_offsets(),
+      sliceOp.static_sizes(), sliceOp.static_strides());
+}
+
 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);
 
auto tiledLoopBodyBuilder = [&](OpBuilder &b, Location loc,
 
     // Insert a insert_slice for each output tensor.
     unsigned resultIdx = 0;
     for (OpOperand *opOperand : op.getOutputTensorOperands()) {
       // TODO: use an interface/adaptor to avoid leaking position in
       // `tiledOperands`.
       Value outputTensor = tiledOperands[opOperand->getOperandNumber()];
       if (auto sliceOp = outputTensor.getDefiningOp<tensor::ExtractSliceOp>()) {
-        tensorResults.push_back(b.create<tensor::InsertSliceOp>(
-            loc, sliceOp.source().getType(), res->getResult(resultIdx),
-            sliceOp.source(), sliceOp.offsets(), sliceOp.sizes(),
-            sliceOp.strides(), sliceOp.static_offsets(), sliceOp.static_sizes(),
-            sliceOp.static_strides()));
+        tensorResults.push_back(insertSliceIntoTensor(
+            b, loc, sliceOp, res->getResult(resultIdx), sliceOp.source()));
       } else {
         tensorResults.push_back(res->getResult(resultIdx));
       }
       ++resultIdx;
     }
     return scf::ValueVector(tensorResults.begin(), tensorResults.end());
   };
   GenerateLoopNest<LoopTy>::doit(b, op.getLoc(), loopRanges, op, iteratorTypes,
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 =
+              insertSliceIntoTensor(b, loc, sliceOp, sliceOp, iterArgs[0]);
+          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.
 static void insertTilingPatterns(RewritePatternSet &patterns,
                                  const LinalgTilingOptions &options) {
   RewritePatternList<GenericOp,
 #define GET_OP_LIST
 #include "mlir/Dialect/Linalg/IR/LinalgStructuredOps.cpp.inc"
                      >::insert(patterns, options);
+  patterns.add<PadTensorOpTilingPattern>(patterns.getContext(), options);
+  patterns.add<ExtractSliceOfPadTensorSwapPattern>(patterns.getContext());
 }
 
 static void
 applyTilingToLoopPatterns(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);
   insertTilingPatterns(patterns, options);
   (void)applyPatternsAndFoldGreedily(funcOp, std::move(patterns));
   (void)applyPatternsAndFoldGreedily(
       funcOp, getLinalgTilingCanonicalizationPatterns(ctx));
   // Drop the marker.
   funcOp.walk([](LinalgOp 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);
   }

mlir/lib/Dialect/Linalg/Utils/Utils.cpp

 #include "mlir/Dialect/Affine/IR/AffineOps.h"
 #include "mlir/Dialect/Linalg/IR/LinalgOps.h"
 #include "mlir/Dialect/Linalg/IR/LinalgTypes.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/Dialect/SCF/SCF.h"
 #include "mlir/Dialect/StandardOps/IR/Ops.h"
 #include "mlir/Dialect/StandardOps/Utils/Utils.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
+#include "mlir/Dialect/Utils/StaticValueUtils.h"
 #include "mlir/IR/AffineExpr.h"
 #include "mlir/IR/AffineExprVisitor.h"
 #include "mlir/IR/AffineMap.h"
 #include "mlir/IR/Matchers.h"
 #include "mlir/IR/OpImplementation.h"
 #include "mlir/Pass/Pass.h"
 #include "mlir/Transforms/LoopUtils.h"
 #include "llvm/Support/Debug.h"
Value makeTiledShape(OpBuilder &builder, Location loc, Value valueToTile,
   offsets.reserve(rank);
   sizes.reserve(rank);
   strides.reserve(rank);
   for (unsigned r = 0; r < rank; ++r) {
     LLVM_DEBUG(llvm::dbgs() << "makeTiledShape: for dim#" << r);
     if (!isTiled(map.getSubMap({r}), tileSizes)) {
       offsets.push_back(builder.getIndexAttr(0));
       Value dim = createOrFoldDimOp(builder, loc, valueToTile, r);
-      sizes.push_back(dim);
+      sizes.push_back(getAsOpFoldResult(dim));
       strides.push_back(builder.getIndexAttr(1));
       LLVM_DEBUG(llvm::dbgs() << ": not tiled: use size: " << dim << "\n");
       continue;
     }
     LLVM_DEBUG(llvm::dbgs() << ": tiled: figure out subsize...\n");
 
     // Tiling creates a new slice at the proper index, the slice step is 1
     // (i.e. the op does not subsample, stepping occurs in the loop).

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

+// RUN: mlir-opt %s -linalg-tile="linalg-tile-sizes=2,3" -cse -split-input-file | \
+// RUN: FileCheck %s -check-prefix=TILE2
+// RUN: mlir-opt %s -linalg-tile="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
+    low[3, 4] high[5, 3] into %output_tensor{
+    ^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:   %[[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, {{.*}}] [15, {{.*}}] [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
+    low[3, 4] high[5, 3] into %output_tensor {
+    ^bb0(%arg1: index, %arg2: index):
+      linalg.yield %pad_value : f32
+    } : tensor<7x9xf32> to tensor<15x16xf32>
+  return %0 : tensor<15x16xf32>
+}