[mlir][linalg] Swap tensor.extract_slice(linalg.fill)

[antiagainst]

Summary

This commit adds a pattern to swap

tensor.extract_slice(linalg.fill(%cst, %init))

into

linalg.fill(%cst, tensor.extract_slice(%init))

when the linalg.fill op have no other users.
This helps to reduce the fill footprint.

[mravishankar]

I am apprehensive of adding this as a canonicalization. Canonicalization that involves multiple operations has always been a footgun

[nicolasvasilache]

+1, please just create a new file with a proper name (and suffixed with Patterns.cpp) and let's trigger this on-demand.
It is quite close to a canonicalization but it is not really a canonicalization and can be surprising..

[mravishankar]

Friday evening my brain doesnt work as well. This is basically a specialized implementation of https://github.com/llvm/llvm-project/blob/main/mlir/include/mlir/Dialect/Tensor/Transforms/Transforms.h#L29 which works on any operation that implements the TilingInterface. If you can just use that method, we dont need to diverge (but that does introduce a more heavy weight treatment of this, and the implementation here is simpler). Either way, this shouldnt be a canonicalization.

[antiagainst]

Address comments

[antiagainst]

Address comments

[antiagainst]

Fix comments

[antiagainst]

@mravishankar @nicolasvasilache : moved to a separate file and entry point now!

In D134102#3800844, @mravishankar wrote:

Friday evening my brain doesnt work as well. This is basically a specialized implementation of https://github.com/llvm/llvm-project/blob/main/mlir/include/mlir/Dialect/Tensor/Transforms/Transforms.h#L29 which works on any operation that implements the TilingInterface. If you can just use that method, we dont need to diverge (but that does introduce a more heavy weight treatment of this, and the implementation here is simpler). Either way, this shouldnt be a canonicalization.

Thanks for the suggestion! I tried to use that but found the above 1) does not support non-1 strides, & 2) cannot handle rank-reducing extract_slice ops. I'm not sure how easy to extend the above for both; esp. the above is used for tiling, which later bufferization relies on the tiled structure. OTOH, this pattern itself is already very simple; just a few lines. I'd prefer to keep the current impl.

mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h

 void populateFoldUnitExtentDimsPatterns(RewritePatternSet &patterns);
 
 /// Patterns that are used to inline constant operands into linalg generic ops.
 void populateInlineConstantOperandsPatterns(RewritePatternSet &patterns);
 
 /// Patterns that are used to bubble up extract slice op above linalg op.
 void populateBubbleUpExtractSliceOpPatterns(RewritePatternSet &patterns);
 
+/// Adds patterns that waps tensor.extract_slice(linalg.fill(%cst, %init)) into
+/// linalg.fill(%cst, tensor.extract_slice(%init)).
+void populateSwapExtractSliceWithFillPatterns(RewritePatternSet &patterns);
+
 /// Return true if two `linalg.generic` operations with producer/consumer
 /// relationship through `fusedOperand` can be fused using elementwise op
 /// fusion.
 bool areElementwiseOpsFusable(OpOperand *fusedOperand);
 
 /// Fuse two `linalg.generic` operations that have a producer-consumer
 /// relationship captured through `fusedOperand`. The method expects
 /// that `areElementwiseOpsFusable` returns true for the given `fusedOperand`.

mlir/lib/Dialect/Linalg/Transforms/CMakeLists.txt

add_mlir_dialect_library(MLIRLinalgTransforms
   InlineScalarOperands.cpp
   Interchange.cpp
   Loops.cpp
   LinalgStrategyPasses.cpp
   NamedOpConversions.cpp
   Promotion.cpp
   Split.cpp
   SplitReduction.cpp
+  SwapExtractSliceWithFillPatterns.cpp
   Tiling.cpp
   TilingInterfaceImpl.cpp
   Transforms.cpp
   Vectorization.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/Linalg
 

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

+//===- SwapExtractSliceWithFillPatterns.cpp -------------------------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Dialect/Linalg/Transforms/Transforms.h"
+#include "mlir/IR/PatternMatch.h"
+
+using namespace mlir;
+using namespace mlir::linalg;
+
+/// Swaps tensor.extract_slice(linalg.fill(%cst, %init)) into linalg.fill(%cst,
+/// tensor.extract_slice(%init)) when the linalg.fill op have no other users.
+/// This helps to reduce the fill footprint.
+struct SwapExtractSliceOfFill final
+    : public OpRewritePattern<tensor::ExtractSliceOp> {
+  using OpRewritePattern::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(tensor::ExtractSliceOp extractOp,
+                                PatternRewriter &rewriter) const override {
+    auto fillOp = extractOp.getSource().getDefiningOp<FillOp>();
+    if (!fillOp || !fillOp->hasOneUse())
+      return failure();
+
+    auto newExtractOp = rewriter.create<tensor::ExtractSliceOp>(
+        extractOp.getLoc(), extractOp.getType(), fillOp.getOutputs()[0],
+        extractOp.getMixedOffsets(), extractOp.getMixedSizes(),
+        extractOp.getMixedStrides());
+    rewriter.replaceOpWithNewOp<FillOp>(extractOp, fillOp.getInputs(),
+                                        ValueRange{newExtractOp.getResult()});
+    return success();
+  }
+};
+
+void mlir::linalg::populateSwapExtractSliceWithFillPatterns(
+    RewritePatternSet &patterns) {
+  patterns.add<SwapExtractSliceOfFill>(patterns.getContext());
+}

mlir/test/Dialect/Linalg/swap-extract-slice-with-fill.mlir

+//RUN: mlir-opt -split-input-file -test-linalg-transform-patterns=test-swap-extract-slice-with-fill-pattern %s | FileCheck %s
+
+// CHECK-LABEL: func.func @swap_fill_insert_slice
+//  CHECK-SAME: (%[[INIT:.+]]: tensor<?x?x?xf32>, %[[OFFSET0:.+]]: index, %[[SIZE1:.+]]: index)
+//       CHECK:   %[[F0:.+]] = arith.constant 0.000000e+00 : f32
+//       CHECK:   %[[EXT:.+]] = tensor.extract_slice %[[INIT]][%[[OFFSET0]], 8, 4] [1, %[[SIZE1]], 6] [1, 3, 1]
+//       CHECK:   %[[FILL:.+]] = linalg.fill ins(%[[F0]] : f32) outs(%[[EXT]] : tensor<?x6xf32>) -> tensor<?x6xf32>
+//       CHECK:   return %[[FILL]]
+func.func @swap_fill_insert_slice(%init : tensor<?x?x?xf32>, %offset0: index, %size1: index) -> tensor<?x6xf32> {
+  %f0 = arith.constant 0.000000e+00 : f32
+  %0 = linalg.fill ins(%f0 : f32) outs(%init : tensor<?x?x?xf32>) -> tensor<?x?x?xf32>
+  %1 = tensor.extract_slice %0[%offset0, 8, 4] [1, %size1, 6] [1, 3, 1]
+    : tensor<?x?x?xf32> to tensor<?x6xf32>
+  return %1: tensor<?x6xf32>
+}
+
+// -----
+
+// CHECK-LABEL: func.func @dont_swap_fill_insert_slice_multi_user
+//       CHECK:   linalg.fill
+//       CHECK:   tensor.extract_slice
+func.func @dont_swap_fill_insert_slice_multi_user(%init : tensor<?x?x?xf32>, %offset0: index, %size1: index) -> (tensor<?x?x?xf32>, tensor<2x?x6xf32>) {
+  %f0 = arith.constant 0.000000e+00 : f32
+  %0 = linalg.fill ins(%f0 : f32) outs(%init : tensor<?x?x?xf32>) -> tensor<?x?x?xf32>
+  %1 = tensor.extract_slice %0[%offset0, 8, 4] [2, %size1, 6] [1, 3, 1]
+    : tensor<?x?x?xf32> to tensor<2x?x6xf32>
+  return %0, %1: tensor<?x?x?xf32>, tensor<2x?x6xf32>
+}

mlir/test/lib/Dialect/Linalg/TestLinalgTransforms.cpp

Option<std::string> loopType{
       llvm::cl::desc("Specify the type of loops to generate: for, parallel or "
                      "tiled_loop"),
       llvm::cl::init("for")};
   Option<bool> testBubbleUpExtractSliceOpPattern{
       *this, "test-bubble-up-extract-slice-op-pattern",
       llvm::cl::desc("Test rewrite of linalgOp + extract_slice into "
                      "extract_slice + linalgOp"),
       llvm::cl::init(false)};
+  Option<bool> testSwapExtractSliceWithFill{
+      *this, "test-swap-extract-slice-with-fill-pattern",
+      llvm::cl::desc(
+          "Test patterns to swap tensor.extract_slice(linalg.fill())"),
+      llvm::cl::init(false)};
 };
 } // namespace
 
 static void applyPatterns(func::FuncOp funcOp) {
   MLIRContext *ctx = funcOp.getContext();
   RewritePatternSet patterns(ctx);
 
   //===--------------------------------------------------------------------===//

 }
 
 static void applyBubbleUpExtractSliceOpPattern(func::FuncOp funcOp) {
   RewritePatternSet patterns(funcOp.getContext());
   populateBubbleUpExtractSliceOpPatterns(patterns);
   (void)applyPatternsAndFoldGreedily(funcOp, std::move(patterns));
 }
 
+static void applySwapExtractSliceWithFillPattern(func::FuncOp funcOp) {
+  RewritePatternSet patterns(funcOp.getContext());
+  populateSwapExtractSliceWithFillPatterns(patterns);
+  (void)applyPatternsAndFoldGreedily(funcOp, std::move(patterns));
+}
+
 /// Apply transformations specified as patterns.
 void TestLinalgTransforms::runOnOperation() {
   auto lambda = [&](void *) {
     getOperation().walk([](LinalgOp op) {
       op->removeAttr(LinalgTransforms::kLinalgTransformMarker);
     });
   };
   std::unique_ptr<void, decltype(lambda)> cleanupGuard{(void *)1, lambda};
return applyTilePattern(getOperation(), loopType, tileSizes, peeledLoops,
                             /*scalarizeDynamicDims=*/false);
   if (testTileScalarizeDynamicDims)
     return applyTilePattern(getOperation(), loopType, tileSizes,
                             /*peeledLoops=*/{}, /*scalarizeDynamicDims=*/true);
   if (testSplitReduction)
     return applySplitReduction(getOperation());
   if (testBubbleUpExtractSliceOpPattern)
     return applyBubbleUpExtractSliceOpPattern(getOperation());
+  if (testSwapExtractSliceWithFill)
+    return applySwapExtractSliceWithFillPattern(getOperation());
 }
 
 namespace mlir {
 namespace test {
 void registerTestLinalgTransforms() {
   PassRegistration<TestLinalgTransforms>();
 }
 } // namespace test
 } // namespace mlir