[mlir][linalg] Swap extract_slice(fill(x)) ops

[springerm]

Summary

This pattern is similar to FoldFillWithTensorReshape, which performs the same swapping with reshapes.

Fill the smaller extracted tensor slice instead of x. This allows for additional simplifications in case x is the result of another extract_slice.

mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp

auto newInit = rewriter.create<TensorReshapeOp>(
         reshapeOp.getReassociation());
     rewriter.replaceOpWithNewOp<FillOp>(reshapeOp, ValueRange{oldFill.value()},
                                         ValueRange{newInit});
 
     return success();
   }
 };
 
+/// Swap extract_slice(fill) to fill(extract_slice).
+///
+/// Only swap the two ops if the extract_slice is the only user of the fill.
+struct SwapExtractSliceOfFill : OpRewritePattern<tensor::ExtractSliceOp> {
+  using OpRewritePattern<tensor::ExtractSliceOp>::OpRewritePattern;
+  LogicalResult matchAndRewrite(tensor::ExtractSliceOp extractSliceOp,
+                                PatternRewriter &rewriter) const override {
+    auto oldFill = extractSliceOp.getSource().getDefiningOp<FillOp>();
+    if (!oldFill)
+      return failure();
+    // Only swap the ops if there is no other user of the fill.
+    if (!extractSliceOp.getSource().hasOneUse())
+      return failure();
+    // Extract from the old fill's source.
+    rewriter.updateRootInPlace(extractSliceOp, [&]() {
+      extractSliceOp.getSourceMutable().assign(oldFill.output());
+    });
+    // Create a new fill and remove the old one.
+    rewriter.setInsertionPointAfter(extractSliceOp);
+    auto newFill =
+        rewriter.create<FillOp>(oldFill.getLoc(), ValueRange{oldFill.value()},
+                                ValueRange{extractSliceOp.getResult()});
+    rewriter.eraseOp(oldFill);
+    // Use the new fill instead of the extract_slice.
+    rewriter.replaceAllUsesExcept(extractSliceOp.getResult(),
+                                  newFill.getResult(0), newFill);
+    return success();
+  }
+};
+
 /// Fold tensor.pad(linalg.fill) into linalg.fill if the padding value and the
 /// filling value are the same.
 struct FoldFillWithPad final : public OpRewritePattern<tensor::PadOp> {
   using OpRewritePattern::OpRewritePattern;
 
   LogicalResult matchAndRewrite(tensor::PadOp padOp,
                                 PatternRewriter &rewriter) const override {
     auto fillOp = padOp.getSource().getDefiningOp<linalg::FillOp>();

 
 } // namespace
 
 void FillOp::getCanonicalizationPatterns(RewritePatternSet &results,
                                          MLIRContext *context) {
   results
       .add<FoldFillWithPad, FoldFillWithTensorReshape<tensor::CollapseShapeOp>,
            FoldFillWithTensorReshape<tensor::ExpandShapeOp>,
-           FoldInsertPadIntoFill>(context);
+           FoldInsertPadIntoFill, SwapExtractSliceOfFill>(context);
 }
 
 //===----------------------------------------------------------------------===//
 // GenericOp
 //===----------------------------------------------------------------------===//
 
 static void buildGenericRegion(
     OpBuilder &builder, Location loc, Region &region, ValueRange inputs,

mlir/test/Dialect/Linalg/canonicalize.mlir

func.func @fold_fill_reshape() -> tensor<6x4xf32> {
   %reshape = tensor.collapse_shape %fill [[0, 1, 2], [3]]
       : tensor<1x2x3x4xf32> into tensor<6x4xf32>
   // CHECK: return %[[FILL]] : tensor<6x4xf32>
   return %reshape : tensor<6x4xf32>
 }
 
 // -----
 
+// CHECK-LABEL: func @fold_fill_extract_slice(
+//  CHECK-SAME:     %[[t:.*]]: tensor<1x1xf32>
+func.func @fold_fill_extract_slice(%t: tensor<1x1xf32>) -> (tensor<f32>) {
+  %cst = arith.constant 0.000000e+00 : f32
+  // CHECK: %[[e:.*]] = tensor.extract_slice %[[t]]
+  // CHECK: %[[f:.*]] = linalg.fill {{.*}} outs(%[[e]] : tensor<f32>)
+  %0 = linalg.fill ins(%cst : f32) outs(%t : tensor<1x1xf32>) -> tensor<1x1xf32>
+  %1 = tensor.extract_slice %0[0, 0] [1, 1] [1, 1] : tensor<1x1xf32> to tensor<f32>
+  // CHECK: return %[[f]]
+  return %1 : tensor<f32>
+}
+
+// -----
+
 //       CHECK: func @fold_fill_reshape_dynamic
 //  CHECK-SAME:   %[[ARG0:.+]]: tensor<?x?x?x?x?xf32>
 func.func @fold_fill_reshape_dynamic(%arg0 : tensor<?x?x?x?x?xf32>) -> tensor<?x?xf32> {
   %zero = arith.constant 0.0 : f32
   // CHECK: %[[RESHAPE:.+]] = tensor.collapse_shape %[[ARG0]]
   %0 = linalg.fill ins(%zero : f32) outs(%arg0 : tensor<?x?x?x?x?xf32>) -> tensor<?x?x?x?x?xf32>
   // CHECK: %[[RESULT:.+]] = linalg.fill ins(%{{.+}}{{.*}}outs(%[[RESHAPE]]
   %1 = tensor.collapse_shape %0 [[0, 1, 2], [3, 4]]