MLIR  19.0.0git
Transforms.h
Go to the documentation of this file.
1 //===- Transforms.h - SCF dialect transformation utilities ------*- 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 // This header file defines transformations on SCF operations.
10 //
11 //===----------------------------------------------------------------------===//
12 
13 #ifndef MLIR_DIALECT_SCF_TRANSFORMS_TRANSFORMS_H_
14 #define MLIR_DIALECT_SCF_TRANSFORMS_TRANSFORMS_H_
15 
17 #include "mlir/Support/LLVM.h"
18 #include "llvm/ADT/ArrayRef.h"
19 
20 namespace mlir {
21 
22 struct LogicalResult;
23 class Region;
24 class RewriterBase;
25 class Operation;
26 class Value;
27 
28 namespace scf {
29 
30 class IfOp;
31 class ForOp;
32 class ParallelOp;
33 class WhileOp;
34 
35 /// Fuses all adjacent scf.parallel operations with identical bounds and step
36 /// into one scf.parallel operations. Uses a naive aliasing and dependency
37 /// analysis.
38 /// User can additionally customize alias checking with `mayAlias` hook.
39 /// `mayAlias` must return false if 2 values are guaranteed to not alias.
40 void naivelyFuseParallelOps(Region &region,
41  llvm::function_ref<bool(Value, Value)> mayAlias);
42 
43 /// Rewrite a for loop with bounds/step that potentially do not divide evenly
44 /// into a for loop where the step divides the iteration space evenly, followed
45 /// by another scf.for for the last (partial) iteration (if any; returned via
46 /// `partialIteration`). This transformation is called "loop peeling".
47 ///
48 /// This transformation is beneficial for a wide range of transformations such
49 /// as vectorization or loop tiling: It enables additional canonicalizations
50 /// inside the peeled loop body such as rewriting masked loads into unmaked
51 /// loads.
52 ///
53 /// E.g., assuming a lower bound of 0 (for illustration purposes):
54 /// ```
55 /// scf.for %iv = %c0 to %ub step %c4 {
56 /// (loop body)
57 /// }
58 /// ```
59 /// is rewritten into the following pseudo IR:
60 /// ```
61 /// %newUb = %ub - (%ub mod %c4)
62 /// scf.for %iv = %c0 to %newUb step %c4 {
63 /// (loop body)
64 /// }
65 /// scf.for %iv2 = %newUb to %ub {
66 /// (loop body)
67 /// }
68 /// ```
69 ///
70 /// After loop peeling, this function tries to simplify affine.min and
71 /// affine.max ops in the body of the peeled loop and in the body of the partial
72 /// iteration loop, taking advantage of the fact that the peeled loop has only
73 /// "full" iterations. This simplification is expected to enable further
74 /// canonicalization opportunities through other patterns.
75 ///
76 /// The return value indicates whether the loop was rewritten or not. Loops are
77 /// not rewritten if:
78 /// * Loop step size is 1 or
79 /// * Loop bounds and step size are static, and step already divides the
80 /// iteration space evenly.
81 ///
82 /// Note: This function rewrites the given scf.for loop in-place and creates a
83 /// new scf.for operation for the last iteration. It replaces all uses of the
84 /// unpeeled loop with the results of the newly generated scf.for.
86  scf::ForOp &partialIteration);
87 
88 /// Peel the first iteration out of the scf.for loop. If there is only one
89 /// iteration, return the original loop.
91  scf::ForOp &partialIteration);
92 
93 /// Tile a parallel loop of the form
94 /// scf.parallel (%i0, %i1) = (%arg0, %arg1) to (%arg2, %arg3)
95 /// step (%arg4, %arg5)
96 ///
97 /// into
98 /// scf.parallel (%i0, %i1) = (%arg0, %arg1) to (%arg2, %arg3)
99 /// step (%arg4*tileSize[0],
100 /// %arg5*tileSize[1])
101 /// scf.parallel (%j0, %j1) = (0, 0) to (min(tileSize[0], %arg2-%j0)
102 /// min(tileSize[1], %arg3-%j1))
103 /// step (%arg4, %arg5)
104 /// The old loop is replaced with the new one.
105 ///
106 /// The function returns the resulting ParallelOps, i.e. {outer_loop_op,
107 /// inner_loop_op}.
108 std::pair<ParallelOp, ParallelOp>
109 tileParallelLoop(ParallelOp op, llvm::ArrayRef<int64_t> tileSizes,
110  bool noMinMaxBounds);
111 
112 /// Options to dictate how loops should be pipelined.
114  /// Lambda returning all the operation in the forOp, with their stage, in the
115  /// order picked for the pipelined loop.
116  using GetScheduleFnType = std::function<void(
117  scf::ForOp, std::vector<std::pair<Operation *, unsigned>> &)>;
119  enum class PipelinerPart {
120  Prologue,
121  Kernel,
122  Epilogue,
123  };
124  /// Lambda called by the pipeliner to allow the user to annotate the IR while
125  /// it is generated.
126  /// The callback passes the operation created along with the part of the
127  /// pipeline and the iteration index. The iteration index is always 0 for the
128  /// kernel. For the prologue and epilogue, it corresponds to the iteration
129  /// peeled out of the loop in the range [0, maxStage[.
131  std::function<void(Operation *, PipelinerPart, unsigned)>;
133 
134  /// Control whether the epilogue should be peeled out of the loop or
135  /// operations should be predicated to skip the early stages in the last loop
136  /// iterations. If the epilogue is predicated; the user needs to provide a
137  /// lambda to generate the predicated version of operations.
138  bool peelEpilogue = true;
139 
140  /// Control whether the transformation checks that the number of iterations is
141  /// greater or equal to the number of stages and skip the transformation if
142  /// this is not the case. If the loop is dynamic and this is set to true and
143  /// the loop bounds are not static the pipeliner will have to predicate
144  /// operations in the the prologue/epilogue.
145  bool supportDynamicLoops = false;
146 
147  // Callback to predicate operations when the prologue or epilogue are not
148  // peeled. This takes the original operation, an i1 predicate value and the
149  // pattern rewriter. It is expected to replace the given operation with
150  // the predicated equivalent and return it, or return nullptr if the
151  // predication is impossible. In the latter case, pipelining will fail and
152  // may leave IR in a partially transformed state.
154  std::function<Operation *(RewriterBase &, Operation *, Value)>;
156 
157  // TODO: add option to decide if the prologue should be peeled.
158 };
159 
160 /// Generate a pipelined version of the scf.for loop based on the schedule given
161 /// as option. This applies the mechanical transformation of changing the loop
162 /// and generating the prologue/epilogue for the pipelining and doesn't make any
163 /// decision regarding the schedule.
164 /// Based on the options the loop is split into several stages.
165 /// The transformation assumes that the scheduling given by user is valid.
166 /// For example if we break a loop into 3 stages named S0, S1, S2 we would
167 /// generate the following code with the number in parenthesis as the iteration
168 /// index:
169 ///
170 /// S0(0) // Prologue
171 /// S0(1) S1(0) // Prologue
172 /// scf.for %I = %C0 to %N - 2 {
173 /// S0(I+2) S1(I+1) S2(I) // Pipelined kernel
174 /// }
175 /// S1(N) S2(N-1) // Epilogue
176 /// S2(N) // Epilogue
177 ///
178 /// If `modifiedIR` is provided, it will be set to a value that indicates
179 /// whether pipelining modified the IR before failing, signaling to the caller
180 /// whether they can proceed with different transformations.
181 FailureOr<ForOp> pipelineForLoop(RewriterBase &rewriter, ForOp forOp,
182  const PipeliningOption &options,
183  bool *modifiedIR = nullptr);
184 
185 /// Create zero-trip-check around a `while` op and return the new loop op in the
186 /// check. The while loop is rotated to avoid evaluating the condition twice
187 ///
188 /// By default the check won't be created for do-while loop as it is not
189 /// required. `forceCreateCheck` can force the creation.
190 ///
191 /// It turns:
192 ///
193 /// scf.while (%arg0 = %init) : (i32) -> i64 {
194 /// %val = .., %arg0 : i64
195 /// %cond = arith.cmpi .., %arg0 : i32
196 /// scf.condition(%cond) %val : i64
197 /// } do {
198 /// ^bb0(%arg1: i64):
199 /// %next = .., %arg1 : i32
200 /// scf.yield %next : i32
201 /// }
202 ///
203 /// into:
204 ///
205 /// %pre_val = .., %init : i64
206 /// %pre_cond = arith.cmpi .., %init : i32
207 /// scf.if %pre_cond -> i64 {
208 /// %res = scf.while (%arg1 = %va0) : (i64) -> i64 {
209 /// %next = .., %arg1 : i32
210 /// %val = .., %next : i64
211 /// %cond = arith.cmpi .., %next : i32
212 /// scf.condition(%cond) %val : i64
213 /// } do {
214 /// ^bb0(%arg2: i64):
215 /// %scf.yield %arg2 : i32
216 /// }
217 /// scf.yield %res : i64
218 /// } else {
219 /// scf.yield %pre_val : i64
220 /// }
222  RewriterBase &rewriter,
223  bool forceCreateCheck = false);
224 
225 } // namespace scf
226 } // namespace mlir
227 
228 #endif // MLIR_DIALECT_SCF_TRANSFORMS_TRANSFORMS_H_
static bool mayAlias(Value first, Value second)
Returns true if two values may be referencing aliasing memory.
static llvm::ManagedStatic< PassManagerOptions > options
This class provides support for representing a failure result, or a valid value of type T.
Definition: LogicalResult.h:78
Operation is the basic unit of execution within MLIR.
Definition: Operation.h:88
This class coordinates the application of a rewrite on a set of IR, providing a way for clients to tr...
Definition: PatternMatch.h:400
This class represents an instance of an SSA value in the MLIR system, representing a computable value...
Definition: Value.h:96
void naivelyFuseParallelOps(Region &region, llvm::function_ref< bool(Value, Value)> mayAlias)
Fuses all adjacent scf.parallel operations with identical bounds and step into one scf....
std::pair< ParallelOp, ParallelOp > tileParallelLoop(ParallelOp op, llvm::ArrayRef< int64_t > tileSizes, bool noMinMaxBounds)
Tile a parallel loop of the form scf.parallel (i0, i1) = (arg0, arg1) to (arg2, arg3) step (arg4,...
LogicalResult peelForLoopAndSimplifyBounds(RewriterBase &rewriter, ForOp forOp, scf::ForOp &partialIteration)
Rewrite a for loop with bounds/step that potentially do not divide evenly into a for loop where the s...
LogicalResult peelForLoopFirstIteration(RewriterBase &rewriter, ForOp forOp, scf::ForOp &partialIteration)
Peel the first iteration out of the scf.for loop.
FailureOr< WhileOp > wrapWhileLoopInZeroTripCheck(WhileOp whileOp, RewriterBase &rewriter, bool forceCreateCheck=false)
Create zero-trip-check around a while op and return the new loop op in the check.
FailureOr< ForOp > pipelineForLoop(RewriterBase &rewriter, ForOp forOp, const PipeliningOption &options, bool *modifiedIR=nullptr)
Generate a pipelined version of the scf.for loop based on the schedule given as option.
Include the generated interface declarations.
This class represents an efficient way to signal success or failure.
Definition: LogicalResult.h:26
Options to dictate how loops should be pipelined.
Definition: Transforms.h:113
std::function< void(scf::ForOp, std::vector< std::pair< Operation *, unsigned > > &)> GetScheduleFnType
Lambda returning all the operation in the forOp, with their stage, in the order picked for the pipeli...
Definition: Transforms.h:117
GetScheduleFnType getScheduleFn
Definition: Transforms.h:118
std::function< void(Operation *, PipelinerPart, unsigned)> AnnotationlFnType
Lambda called by the pipeliner to allow the user to annotate the IR while it is generated.
Definition: Transforms.h:131
bool peelEpilogue
Control whether the epilogue should be peeled out of the loop or operations should be predicated to s...
Definition: Transforms.h:138
AnnotationlFnType annotateFn
Definition: Transforms.h:132
bool supportDynamicLoops
Control whether the transformation checks that the number of iterations is greater or equal to the nu...
Definition: Transforms.h:145
std::function< Operation *(RewriterBase &, Operation *, Value)> PredicateOpFn
Definition: Transforms.h:154