[mlir][Linalg] Add a pattern to fold `linalg.fill`s with `tensor.expand/collapse_shape` ops.

[mravishankar]

Summary

Combine these patterns with the SwapExtractSliceOfFill patterns and
expose them through a single entry point
populateFoldLinalgFillPatterns. Rename files and tests accordingly.

[hanchung]

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

48–50

MLIR style nit: do not add braces

[springerm]

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

61

The name is a bit misleading. These patterns do not fold anything.

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

192

This still has the old name.

[mravishankar]

I can change the names to be consistent. What makes sense 'swapFillWithTensorOps` ?

[springerm]

In D159048#4625732, @mravishankar wrote:

I can change the names to be consistent. What makes sense 'swapFillWithTensorOps` ?

that sounds good to me!

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

 void populateMoveInitOperandsToInputPattern(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);
+/// Adds patterns that fold `linalg.fill` operations with its consumers.
+/// Currently folds it with consumers that are `tensor.extract_slice` (if only
+/// use) and `tensor.collapse/expand_shape` operations.
+void populateFoldLinalgFillPatterns(RewritePatternSet &patterns);
 
 /// Patterns to apply `splitReduction` below.
 void populateSplitReductionPattern(
     RewritePatternSet &patterns,
     const ControlSplitReductionFn &controlSplitReductionFn,
     bool useAlloc = false);
 
 } // namespace linalg
 } // namespace mlir
 
 #endif // MLIR_DIALECT_LINALG_TRANSFORMS_TRANSFORMS_H

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

 add_mlir_dialect_library(MLIRLinalgTransforms
   BubbleUpExtractSlice.cpp
   BufferizableOpInterfaceImpl.cpp
   Bufferize.cpp
   ConstantFold.cpp
   ConvertToDestinationStyle.cpp
   ConvertConv2DToImg2Col.cpp
   DataLayoutPropagation.cpp
   DecomposeLinalgOps.cpp
   Detensorize.cpp
   DropUnitDims.cpp
   ElementwiseOpFusion.cpp
   ElementwiseToLinalg.cpp
   EliminateEmptyTensors.cpp
   EraseUnusedOperandsAndResults.cpp
+  FoldLinalgFillPatterns.cpp
   FusePadOpWithLinalgProducer.cpp
   Fusion.cpp
   Generalization.cpp
   Hoisting.cpp
   HoistPadding.cpp
   InlineScalarOperands.cpp
   Interchange.cpp
   Loops.cpp
   NamedOpConversions.cpp
   Padding.cpp
   Promotion.cpp
   Split.cpp
   SplitReduction.cpp
   SubsetHoisting.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/FoldLinalgFillPatterns.cpp

auto newExtractOp = rewriter.create<tensor::ExtractSliceOp>(
         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());
+/// Fold `tensor.expand/collapse_shape(linalg.fill(%val, %init))` to
+/// `linalg.fill(%val, tensor.expand/collapse_shape(%init))`.
+template <typename T>
+struct SwapTensorShapeChangePatternsWithFillOp final
+    : public OpRewritePattern<T> {
+  using OpRewritePattern<T>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(T reshapeOp,
+                                PatternRewriter &rewriter) const override {
+    auto fillOp = reshapeOp->getOperand(0).template getDefiningOp<FillOp>();
+    if (!fillOp) {
+      return failure();
+    }
+    auto reshapeInitOp =
+        rewriter.create<T>(reshapeOp.getLoc(), reshapeOp.getResult().getType(),
+                           fillOp.getDpsInitOperand(0)->get(),
+                           reshapeOp.getReassociationIndices());
+    rewriter.replaceOpWithNewOp<linalg::FillOp>(reshapeOp, fillOp.getInputs(),
+                                                reshapeInitOp.getResult());
+    return success();
+  }
+};
+
+void mlir::linalg::populateFoldLinalgFillPatterns(RewritePatternSet &patterns) {
+  patterns.add<SwapExtractSliceOfFill,
+               SwapTensorShapeChangePatternsWithFillOp<tensor::CollapseShapeOp>,
+               SwapTensorShapeChangePatternsWithFillOp<tensor::ExpandShapeOp>>(
+      patterns.getContext());
 }

mlir/test/Dialect/Linalg/fold-fill-patterns.mlir

-//RUN: mlir-opt -split-input-file -test-linalg-transform-patterns=test-swap-extract-slice-with-fill-pattern %s | FileCheck %s
+//RUN: mlir-opt -split-input-file -test-linalg-transform-patterns=test-fold-linalg-fill-patterns %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> {

 //       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>
 }
+
+// -----
+
+func.func @fold_fill_with_collapse_shape(%cst : f32, %init : tensor<?x?x?x?x?xf32>) -> tensor<?x?xf32> {
+  %0 = linalg.fill ins(%cst : f32) outs(%init : tensor<?x?x?x?x?xf32>) -> tensor<?x?x?x?x?xf32>
+  %1 = tensor.collapse_shape %0 [[0, 1, 2], [3, 4]] : tensor<?x?x?x?x?xf32> into tensor<?x?xf32>
+  return %1 : tensor<?x?xf32>
+}
+// CHECK-LABEL: func @fold_fill_with_collapse_shape
+//  CHECK-SAME:     %[[INIT:[a-zA-Z0-9]+]]: tensor<?x?x?x?x?xf32>
+//       CHECK:   %[[INIT_COLLAPSE:.+]] = tensor.collapse_shape %[[INIT]]
+//       CHECK:   %[[FILL:.+]] = linalg.fill
+//  CHECK-SAME:       outs(%[[INIT_COLLAPSE]] :
+
+// -----
+
+func.func @fold_fill_with_expand_shape(%cst : f32, %init : tensor<?x?xf32>) -> tensor<?x2x3x?x4xf32> {
+  %0 = linalg.fill ins(%cst : f32) outs(%init : tensor<?x?xf32>) -> tensor<?x?xf32>
+  %1 = tensor.expand_shape %0 [[0, 1], [2, 3, 4]] : tensor<?x?xf32> into tensor<?x2x3x?x4xf32>
+  return %1 : tensor<?x2x3x?x4xf32>
+}
+// CHECK-LABEL: func @fold_fill_with_expand_shape
+//  CHECK-SAME:     %[[INIT:[a-zA-Z0-9]+]]: tensor<?x?xf32>
+//       CHECK:   %[[INIT_EXPAND:.+]] = tensor.expand_shape %[[INIT]]
+//       CHECK:   %[[FILL:.+]] = linalg.fill
+//  CHECK-SAME:       outs(%[[INIT_EXPAND]] :

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",
+  Option<bool> testFoldLinalgFillPatterns{
+      *this, "test-fold-linalg-fill-patterns",
       llvm::cl::desc(
-          "Test patterns to swap tensor.extract_slice(linalg.fill())"),
+          "Test patterns to fold linalg.fill with tensor operations"),
       llvm::cl::init(false)};
   Option<bool> testEraseUnusedOperandsAndResults{
       *this, "test-erase-unused-operands-and-results",
       llvm::cl::desc("Test patterns to erase unused operands and results"),
       llvm::cl::init(false)};
   Option<bool> testEraseUnnecessaryInputs{
       *this, "test-erase-unnecessary-inputs",
       llvm::cl::desc("Test patterns to erase unnecessary inputs"),

 }
 
 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);
+  populateFoldLinalgFillPatterns(patterns);
   (void)applyPatternsAndFoldGreedily(funcOp, std::move(patterns));
 }
 
 static void applyEraseUnusedOperandsAndResultsPatterns(func::FuncOp funcOp) {
   RewritePatternSet patterns(funcOp.getContext());
   populateEraseUnusedOperandsAndResultsPatterns(patterns);
   (void)applyPatternsAndFoldGreedily(funcOp, std::move(patterns));
 }
void TestLinalgTransforms::runOnOperation() {
   if (testGeneralizeTensorPackOp)
     return applyGeneralizeTensorPackPatterns(getOperation());
   if (testGeneralizeTensorUnPackOp)
     return applyGeneralizeTensorUnPackPatterns(getOperation());
   if (testSwapSubTensorPadTensor)
     return applyExtractSliceOfPadTensorSwapPattern(getOperation());
   if (testBubbleUpExtractSliceOpPattern)
     return applyBubbleUpExtractSliceOpPattern(getOperation());
-  if (testSwapExtractSliceWithFill)
+  if (testFoldLinalgFillPatterns)
     return applySwapExtractSliceWithFillPattern(getOperation());
   if (testEraseUnusedOperandsAndResults)
     return applyEraseUnusedOperandsAndResultsPatterns(getOperation());
   if (testEraseUnnecessaryInputs)
     return applyEraseUnnecessaryInputs(getOperation());
 }
 
 namespace mlir {
 namespace test {
 void registerTestLinalgTransforms() {
   PassRegistration<TestLinalgTransforms>();
 }
 } // namespace test
 } // namespace mlir