MLIR  22.0.0git
SwapExtractSliceWithFillPatterns.cpp
Go to the documentation of this file.
1 //===- SwapExtractSliceWithFillPatterns.cpp -------------------------------===//
2 //
3 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4 // See https://llvm.org/LICENSE.txt for license information.
5 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6 //
7 //===----------------------------------------------------------------------===//
8 
10 #include "mlir/IR/PatternMatch.h"
11 
12 using namespace mlir;
13 using namespace mlir::linalg;
14 
15 /// swaps:
16 /// `tensor.extract_slice(linalg.fill(%cst, %init))`
17 /// with:
18 /// `linalg.fill(%cst, tensor.extract_slice(%init))`
19 ///
20 /// when the linalg.fill op have no other users.
21 /// This helps to reduce the fill footprint.
23  : public OpRewritePattern<tensor::ExtractSliceOp> {
25 
26  LogicalResult matchAndRewrite(tensor::ExtractSliceOp extractOp,
27  PatternRewriter &rewriter) const override {
28  auto fillOp = extractOp.getSource().getDefiningOp<FillOp>();
29  if (!fillOp || !fillOp->hasOneUse())
30  return failure();
31 
32  auto newExtractOp = tensor::ExtractSliceOp::create(
33  rewriter, extractOp.getLoc(), extractOp.getType(),
34  fillOp.getOutputs()[0], extractOp.getMixedOffsets(),
35  extractOp.getMixedSizes(), extractOp.getMixedStrides());
36  rewriter.replaceOpWithNewOp<FillOp>(extractOp, fillOp.getInputs(),
37  ValueRange{newExtractOp.getResult()});
38  return success();
39  }
40 };
41 
44  patterns.add<SwapExtractSliceOfFill>(patterns.getContext());
45 }
Ty getType(Args &&...args)
Get or construct an instance of the type Ty with provided arguments.
Definition: Builders.h:89
A special type of RewriterBase that coordinates the application of a rewrite pattern on the current I...
Definition: PatternMatch.h:769
OpTy replaceOpWithNewOp(Operation *op, Args &&...args)
Replace the results of the given (original) op with a new op that is created without verification (re...
Definition: PatternMatch.h:519
This class provides an abstraction over the different types of ranges over Values.
Definition: ValueRange.h:387
void populateSwapExtractSliceWithFillPatterns(RewritePatternSet &patterns)
Adds patterns that waps tensor.extract_slice(linalg.fill(cst, init)) into linalg.fill(cst,...
Include the generated interface declarations.
const FrozenRewritePatternSet & patterns
swaps: tensor.extract_slice(linalg.fill(cst, init)) with: linalg.fill(cst, tensor....
LogicalResult matchAndRewrite(tensor::ExtractSliceOp extractOp, PatternRewriter &rewriter) const override
OpRewritePattern is a wrapper around RewritePattern that allows for matching and rewriting against an...
Definition: PatternMatch.h:314
OpRewritePattern(MLIRContext *context, PatternBenefit benefit=1, ArrayRef< StringRef > generatedNames={})
Patterns must specify the root operation name they match against, and can also specify the benefit of...
Definition: PatternMatch.h:319