MLIR  21.0.0git
Transforms.h
Go to the documentation of this file.
1 //===- Transforms.h - Tensor Transformation Patterns ------------*- C++ -*-===//
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 
9 #ifndef MLIR_DIALECT_TENSOR_TRANSFORMS_TRANSFORMS_H
10 #define MLIR_DIALECT_TENSOR_TRANSFORMS_TRANSFORMS_H
11 
13 #include "mlir/IR/PatternMatch.h"
15 
16 namespace mlir {
17 
18 struct TilingResult;
19 
20 namespace tensor {
21 
22 //===----------------------------------------------------------------------===//
23 // Patterns
24 //===----------------------------------------------------------------------===//
25 
26 /// Method to swap an `tensor.extract_slice` with its producer when the
27 /// producer implements the `TilingInterface`. The pattern itself does not
28 /// provide a mechanism to control where the application happens. With use of
29 /// transform dialect that control is done within the transform dialect. Other
30 /// use cases can inherit from this pattern and add necessary controls.
31 FailureOr<TilingResult> replaceExtractSliceWithTiledProducer(
32  OpBuilder &builder, tensor::ExtractSliceOp sliceOp, OpResult producerOp);
33 
34 /// Method to swap `tensor.insert_slice`s with their consumers when the
35 /// consumer implements the `TilingInterface`. The size of `sliceOps` and
36 /// `consumerOperands` is expected to be the same. Every entry in
37 /// `consumerOperands` represents a use of the the corresponding
38 /// entry in `sliceOps` in the consumer. All entries of `consumerOperands` is
39 /// expected to be uses in the same consumer.
40 FailureOr<TilingResult>
41 replaceInsertSlicesWithTiledConsumer(OpBuilder &builder,
42  ArrayRef<tensor::InsertSliceOp> sliceOps,
43  ArrayRef<OpOperand *> consumerOperands);
44 
45 //===----------------------------------------------------------------------===//
46 // Populate functions.
47 //===----------------------------------------------------------------------===//
48 
49 /// Appends patterns for folding tensor subset ops into consumer load/store
50 /// ops into `patterns`. (This includes patterns for folding tensor subset ops
51 /// into vector transfer ops.)
52 void populateFoldTensorSubsetOpPatterns(RewritePatternSet &patterns);
53 
54 /// Appends patterns for folding tensor subset ops into vector transfer ops.
56  RewritePatternSet &patterns);
57 
58 /// Collects patterns to merge consecutive tensor.insert_slice/extract_slice
59 /// into one. These patterns are in this separate entry point because the
60 /// bufferization is sensitive to IR structure, particularly those
61 /// tensor.extract_slice and tensor.insert_slice ops for creating the slices.
63  RewritePatternSet &patterns);
64 
65 /// Appends patterns that are used to bubble up tensor.extract slice op above
66 /// its producer. When used as cleanup patterns of tile and fuse, enables fusing
67 /// the producer with the consumer even if the producer does not implement the
68 /// tiling interface.
69 void populateBubbleUpExtractSliceOpPatterns(RewritePatternSet &patterns);
70 
71 /// Populates `patterns` with patterns that drop redundant tensor.insert_slice
72 /// rank expansions.
74  RewritePatternSet &patterns);
75 
76 /// Populates `patterns` with patterns that fold `tensor.expand_shape` and
77 /// `tensor.collapse_shape` into other ops.
79 
80 /// Populates `patterns` with patterns that bubble up `tensor.expand_shape`
81 /// through `tensor.collapse_shape` ops.
82 void populateBubbleUpExpandShapePatterns(RewritePatternSet &patterns);
83 
84 /// Populates `patterns` with patterns that fold tensor.empty with its
85 /// consumers.
86 ///
87 /// If `singleUseOnly` is set to "true", only tensor.empty ops with a single
88 /// use are folded.
89 void populateFoldTensorEmptyPatterns(RewritePatternSet &patterns,
90  bool foldSingleUseOnly = false);
91 
92 /// Populates `patterns` with patterns that decompose `tensor.concat` into
93 /// `tensor.empty` of a tensor of the concatenated size, followed by a chain
94 /// of `tensor.insert_slice` operations on the inputs. This is intended to be
95 /// used as a fallback tensor -> tensor lowering that decomposes concat such
96 /// that it can be bufferized into a sequence of copies.
97 void populateDecomposeTensorConcatPatterns(RewritePatternSet &patterns);
98 
99 using ControlFoldFn = std::function<bool(OpOperand *)>;
100 
101 /// Populates `patterns` with patterns that replace tensor ops (such as
102 /// tensor.generate) with constants when possible.
104  const ControlFoldFn &controlFn);
105 
106 //===----------------------------------------------------------------------===//
107 // Transform helpers
108 //===----------------------------------------------------------------------===//
109 
110 /// Build a new tensor::PadOp with low/high padding that is independent of all
111 /// given independencies. If the op is already independent of all
112 /// independencies, the same PadOp result is returned.
113 ///
114 /// Failure indicates the no suitable upper bound for low/high padding could be
115 /// found.
116 ///
117 /// Example:
118 /// scf.for %iv = %lb to %ub step %step {
119 /// %high = affine.apply affine_map<(d0)[s0] -> (s0 - d0)> (%i)[%ub]
120 /// %p = tensor.pad %t low[5] high[%high] ...
121 /// ...
122 /// }
123 ///
124 /// The function builds IR such as:
125 /// %high_new = affine.apply affine_map<()[s0, s1] -> (-s0 + s1)> ()[%lb, %ub]
126 /// %p_hoistable = tensor.pad %t low[5] high[%high_new]
127 /// %dim = tensor.dim %t, %c0
128 /// %size = affine.apply affine_map<(d0)[s0, s1] -> (-d0 + s0 + s1 + 5)>
129 /// (%iv)[%ub, %dim]
130 /// %slice = tensor.extract_slice %p_hoistable [0] [%size] [1]
131 ///
132 /// The slice is returned.
133 FailureOr<Value> buildIndependentOp(OpBuilder &b, tensor::PadOp padOp,
134  ValueRange independencies);
135 
136 /// Build a new tensor::EmptyOp who's dynamic sizes are independent of all
137 /// given independencies. If the op is already independent of all
138 /// independencies, the same EmptyOp result is returned.
139 ///
140 /// Failure indicates the no suitable upper bound for the dynamic sizes could be
141 /// found.
142 FailureOr<Value> buildIndependentOp(OpBuilder &b, tensor::EmptyOp emptyOp,
143  ValueRange independencies);
144 
145 } // namespace tensor
146 } // namespace mlir
147 
148 #endif // MLIR_DIALECT_TENSOR_TRANSFORMS_TRANSFORMS_H
This class helps build Operations.
Definition: Builders.h:205
This class represents an operand of an operation.
Definition: Value.h:257
This class provides an abstraction over the different types of ranges over Values.
Definition: ValueRange.h:387
void populateFoldTensorEmptyPatterns(RewritePatternSet &patterns, bool foldSingleUseOnly=false)
Populates patterns with patterns that fold tensor.empty with its consumers.
FailureOr< TilingResult > replaceExtractSliceWithTiledProducer(OpBuilder &builder, tensor::ExtractSliceOp sliceOp, OpResult producerOp)
Method to swap an tensor.extract_slice with its producer when the producer implements the TilingInter...
FailureOr< TilingResult > replaceInsertSlicesWithTiledConsumer(OpBuilder &builder, ArrayRef< tensor::InsertSliceOp > sliceOps, ArrayRef< OpOperand * > consumerOperands)
Method to swap tensor.insert_slices with their consumers when the consumer implements the TilingInter...
void populateMergeConsecutiveInsertExtractSlicePatterns(RewritePatternSet &patterns)
Collects patterns to merge consecutive tensor.insert_slice/extract_slice into one.
void populateDecomposeTensorConcatPatterns(RewritePatternSet &patterns)
Populates patterns with patterns that decompose tensor.concat into tensor.empty of a tensor of the co...
void populateFoldTensorSubsetOpPatterns(RewritePatternSet &patterns)
Appends patterns for folding tensor subset ops into consumer load/store ops into patterns.
void populateReassociativeReshapeFoldingPatterns(RewritePatternSet &patterns)
Populates patterns with patterns that fold tensor.expand_shape and tensor.collapse_shape into other o...
void populateBubbleUpExtractSliceOpPatterns(RewritePatternSet &patterns)
Appends patterns that are used to bubble up tensor.extract slice op above its producer.
void populateBubbleUpExpandShapePatterns(RewritePatternSet &patterns)
Populates patterns with patterns that bubble up tensor.expand_shape through tensor....
void populateDropRedundantInsertSliceRankExpansionPatterns(RewritePatternSet &patterns)
Populates patterns with patterns that drop redundant tensor.insert_slice rank expansions.
FailureOr< Value > buildIndependentOp(OpBuilder &b, tensor::PadOp padOp, ValueRange independencies)
Build a new tensor::PadOp with low/high padding that is independent of all given independencies.
std::function< bool(OpOperand *)> ControlFoldFn
Definition: Transforms.h:99
void populateRewriteAsConstantPatterns(RewritePatternSet &patterns, const ControlFoldFn &controlFn)
Populates patterns with patterns that replace tensor ops (such as tensor.generate) with constants whe...
void populateFoldTensorSubsetIntoVectorTransferPatterns(RewritePatternSet &patterns)
Appends patterns for folding tensor subset ops into vector transfer ops.
Include the generated interface declarations.
const FrozenRewritePatternSet & patterns