[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/lib/Dialect/Linalg/IR/LinalgOps.cpp

LogicalResult matchAndRewrite(tensor::InsertSliceOp insertOp,
 
     rewriter.replaceOpWithNewOp<tensor::InsertSliceOp>(
         insertOp, srcPadOp.getSource(), insertOp.getDest(), newOffsets,
         newSizes, insertOp.getMixedStrides());
     return success();
   }
 };
 
+/// 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();
+  }
+};
+
 } // namespace
 
 void FillOp::getCanonicalizationPatterns(RewritePatternSet &results,
                                          MLIRContext *context) {
   results
       .add<FoldFillWithPad, FoldFillWithTensorReshape<tensor::CollapseShapeOp>,
            FoldFillWithTensorReshape<tensor::ExpandShapeOp>,
-           FoldInsertPadIntoFill>(context);
+           FoldInsertPadIntoFill, SwapExtractSliceOfFill>(context);
 }
 
 //===----------------------------------------------------------------------===//
 // GenericOps
 //===----------------------------------------------------------------------===//
 void GenericOp::build(
     OpBuilder &builder, OperationState &result, TypeRange resultTensorTypes,
     ValueRange inputs, ValueRange outputs, ArrayAttr indexingMaps,

mlir/test/Dialect/Linalg/canonicalize.mlir

 //   CHECK-DAG:   %[[INIT1:.+]] = linalg.init_tensor [2, 3, 4] : tensor<2x3x4xf32>
 //   CHECK-DAG:   %[[CAST:.+]] = tensor.cast %[[ARG0]] : tensor<?x?x?xf32> to tensor<4x3x2xf32>
 //   CHECK-DAG:   %[[INIT2:.+]] = linalg.init_tensor [3, 2, 4] : tensor<3x2x4xf32>
 //       CHECK:   %[[GENERIC:.+]]:2 = linalg.generic
 //  CHECK-SAME:       ins(%[[CAST]] :
 //  CHECK-SAME:       outs(%[[INIT2]], %[[INIT1]] :
 //       CHECK:   %[[RETURN_CAST:.+]] = tensor.cast %[[GENERIC]]#0 : tensor<3x2x4xf32> to tensor<?x?x?xf32>
 //       CHECK:   return %[[RETURN_CAST]], %[[GENERIC]]#1
+
+// -----
+
+// 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>
+}