MLIR  20.0.0git
Utils.h
Go to the documentation of this file.
1 //===- Utils.h - SCF dialect 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 prototypes for various SCF utilities.
10 //
11 //===----------------------------------------------------------------------===//
12 
13 #ifndef MLIR_DIALECT_SCF_UTILS_UTILS_H_
14 #define MLIR_DIALECT_SCF_UTILS_UTILS_H_
15 
17 #include "mlir/IR/PatternMatch.h"
18 #include "mlir/Support/LLVM.h"
19 #include "llvm/ADT/STLExtras.h"
20 #include <optional>
21 
22 namespace mlir {
23 class Location;
24 class Operation;
25 class OpBuilder;
26 class Region;
27 class RewriterBase;
28 class ValueRange;
29 class Value;
30 
31 namespace func {
32 class CallOp;
33 class FuncOp;
34 } // namespace func
35 
36 /// Update a perfectly nested loop nest to yield new values from the innermost
37 /// loop and propagating it up through the loop nest. This function
38 /// - Expects `loopNest` to be a perfectly nested loop with outer most loop
39 /// first and innermost loop last.
40 /// - `newIterOperands` are the initialization values to be used for the
41 /// outermost loop
42 /// - `newYielValueFn` is the callback that generates the new values to be
43 /// yielded from within the innermost loop.
44 /// - The original loops are not erased, but are left in a "no-op" state where
45 /// the body of the loop just yields the basic block arguments that correspond
46 /// to the initialization values of a loop. The original loops are dead after
47 /// this method.
48 /// - If `replaceIterOperandsUsesInLoop` is true, all uses of the
49 /// `newIterOperands` within the generated new loop are replaced with the
50 /// corresponding `BlockArgument` in the loop body.
51 SmallVector<scf::ForOp> replaceLoopNestWithNewYields(
52  RewriterBase &rewriter, MutableArrayRef<scf::ForOp> loopNest,
53  ValueRange newIterOperands, const NewYieldValuesFn &newYieldValuesFn,
54  bool replaceIterOperandsUsesInLoop = true);
55 
56 /// Outline a region with a single block into a new FuncOp.
57 /// Assumes the FuncOp result types is the type of the yielded operands of the
58 /// single block. This constraint makes it easy to determine the result.
59 /// This method also clones the `arith::ConstantIndexOp` at the start of
60 /// `outlinedFuncBody` to alloc simple canonicalizations.
61 /// Creates a new FuncOp and thus cannot be used in a FuncOp pass.
62 /// The client is responsible for providing a unique `funcName` that will not
63 /// collide with another FuncOp name. If `callOp` is provided, it will be set
64 /// to point to the operation that calls the outlined function.
65 // TODO: support more than single-block regions.
66 // TODO: more flexible constant handling.
67 FailureOr<func::FuncOp>
68 outlineSingleBlockRegion(RewriterBase &rewriter, Location loc, Region &region,
69  StringRef funcName, func::CallOp *callOp = nullptr);
70 
71 /// Outline the then and/or else regions of `ifOp` as follows:
72 /// - if `thenFn` is not null, `thenFnName` must be specified and the `then`
73 /// region is inlined into a new FuncOp that is captured by the pointer.
74 /// - if `elseFn` is not null, `elseFnName` must be specified and the `else`
75 /// region is inlined into a new FuncOp that is captured by the pointer.
76 /// Creates new FuncOps and thus cannot be used in a FuncOp pass.
77 /// The client is responsible for providing a unique `thenFnName`/`elseFnName`
78 /// that will not collide with another FuncOp name.
79 LogicalResult outlineIfOp(RewriterBase &b, scf::IfOp ifOp, func::FuncOp *thenFn,
80  StringRef thenFnName, func::FuncOp *elseFn,
81  StringRef elseFnName);
82 
83 /// Get a list of innermost parallel loops contained in `rootOp`. Innermost
84 /// parallel loops are those that do not contain further parallel loops
85 /// themselves.
86 bool getInnermostParallelLoops(Operation *rootOp,
87  SmallVectorImpl<scf::ParallelOp> &result);
88 
89 /// Return the min/max expressions for `value` if it is an induction variable
90 /// from scf.for or scf.parallel loop.
91 /// if `loopFilter` is passed, the filter determines which loop to consider.
92 /// Other induction variables are ignored.
93 std::optional<std::pair<AffineExpr, AffineExpr>>
95  SmallVectorImpl<Value> &symbols,
96  llvm::function_ref<bool(Operation *)> loopFilter = nullptr);
97 
98 /// Replace a perfect nest of "for" loops with a single linearized loop. Assumes
99 /// `loops` contains a list of perfectly nested loops with bounds and steps
100 /// independent of any loop induction variable involved in the nest.
101 LogicalResult coalesceLoops(MutableArrayRef<scf::ForOp> loops);
102 LogicalResult coalesceLoops(RewriterBase &rewriter,
104 
105 /// Walk an affine.for to find a band to coalesce.
106 LogicalResult coalescePerfectlyNestedSCFForLoops(scf::ForOp op);
107 
108 /// Take the ParallelLoop and for each set of dimension indices, combine them
109 /// into a single dimension. combinedDimensions must contain each index into
110 /// loops exactly once.
111 void collapseParallelLoops(RewriterBase &rewriter, scf::ParallelOp loops,
112  ArrayRef<std::vector<unsigned>> combinedDimensions);
113 
114 /// Unrolls this for operation by the specified unroll factor. Returns failure
115 /// if the loop cannot be unrolled either due to restrictions or due to invalid
116 /// unroll factors. Requires positive loop bounds and step. If specified,
117 /// annotates the Ops in each unrolled iteration by applying `annotateFn`.
118 LogicalResult loopUnrollByFactor(
119  scf::ForOp forOp, uint64_t unrollFactor,
120  function_ref<void(unsigned, Operation *, OpBuilder)> annotateFn = nullptr);
121 
122 /// Unrolls and jams this `scf.for` operation by the specified unroll factor.
123 /// Returns failure if the loop cannot be unrolled either due to restrictions or
124 /// due to invalid unroll factors. In case of unroll factor of 1, the function
125 /// bails out without doing anything (returns success). Currently, only constant
126 /// trip count that are divided by the unroll factor is supported. Currently,
127 /// for operations with results are not supported.
128 LogicalResult loopUnrollJamByFactor(scf::ForOp forOp, uint64_t unrollFactor);
129 
130 /// Materialize bounds and step of a zero-based and unit-step loop derived by
131 /// normalizing the specified bounds and step.
133  OpFoldResult lb, OpFoldResult ub,
134  OpFoldResult step);
135 
136 /// Get back the original induction variable values after loop normalization.
138  Value normalizedIv, OpFoldResult origLb,
139  OpFoldResult origStep);
140 
141 /// Tile a nest of standard for loops rooted at `rootForOp` by finding such
142 /// parametric tile sizes that the outer loops have a fixed number of iterations
143 /// as defined in `sizes`.
145 using TileLoops = std::pair<Loops, Loops>;
146 TileLoops extractFixedOuterLoops(scf::ForOp rootFOrOp, ArrayRef<int64_t> sizes);
147 
148 /// Performs tiling fo imperfectly nested loops (with interchange) by
149 /// strip-mining the `forOps` by `sizes` and sinking them, in their order of
150 /// occurrence in `forOps`, under each of the `targets`.
151 /// Returns the new AffineForOps, one per each of (`forOps`, `targets`) pair,
152 /// nested immediately under each of `targets`.
154  ArrayRef<scf::ForOp> targets);
155 
156 /// Performs tiling (with interchange) by strip-mining the `forOps` by `sizes`
157 /// and sinking them, in their order of occurrence in `forOps`, under `target`.
158 /// Returns the new AffineForOps, one per `forOps`, nested immediately under
159 /// `target`.
161  scf::ForOp target);
162 
163 /// Tile a nest of scf::ForOp loops rooted at `rootForOp` with the given
164 /// (parametric) sizes. Sizes are expected to be strictly positive values at
165 /// runtime. If more sizes than loops are provided, discard the trailing values
166 /// in sizes. Assumes the loop nest is permutable.
167 /// Returns the newly created intra-tile loops.
168 Loops tilePerfectlyNested(scf::ForOp rootForOp, ArrayRef<Value> sizes);
169 
170 /// Get perfectly nested sequence of loops starting at root of loop nest
171 /// (the first op being another AffineFor, and the second op - a terminator).
172 /// A loop is perfectly nested iff: the first op in the loop's body is another
173 /// AffineForOp, and the second op is a terminator).
175  scf::ForOp root);
176 
177 /// Given two scf.forall loops, `target` and `source`, fuses `target` into
178 /// `source`. Assumes that the given loops are siblings and are independent of
179 /// each other.
180 ///
181 /// This function does not perform any legality checks and simply fuses the
182 /// loops. The caller is responsible for ensuring that the loops are legal to
183 /// fuse.
184 scf::ForallOp fuseIndependentSiblingForallLoops(scf::ForallOp target,
185  scf::ForallOp source,
186  RewriterBase &rewriter);
187 
188 /// Given two scf.for loops, `target` and `source`, fuses `target` into
189 /// `source`. Assumes that the given loops are siblings and are independent of
190 /// each other.
191 ///
192 /// This function does not perform any legality checks and simply fuses the
193 /// loops. The caller is responsible for ensuring that the loops are legal to
194 /// fuse.
195 scf::ForOp fuseIndependentSiblingForLoops(scf::ForOp target, scf::ForOp source,
196  RewriterBase &rewriter);
197 
198 /// Normalize an `scf.forall` operation. Returns `failure()`if normalization
199 /// fails.
200 // On `success()` returns the
201 /// newly created operation with all uses of the original operation replaced
202 /// with results of the new operation.
203 FailureOr<scf::ForallOp> normalizeForallOp(RewriterBase &rewriter,
204  scf::ForallOp forallOp);
205 
206 } // namespace mlir
207 
208 #endif // MLIR_DIALECT_SCF_UTILS_UTILS_H_
This class defines the main interface for locations in MLIR and acts as a non-nullable wrapper around...
Definition: Location.h:66
This class helps build Operations.
Definition: Builders.h:215
This class represents a single result from folding an operation.
Definition: OpDefinition.h:268
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
Include the generated interface declarations.
void getPerfectlyNestedLoops(SmallVectorImpl< scf::ForOp > &nestedLoops, scf::ForOp root)
Get perfectly nested sequence of loops starting at root of loop nest (the first op being another Affi...
Definition: Utils.cpp:1313
LogicalResult loopUnrollByFactor(scf::ForOp forOp, uint64_t unrollFactor, function_ref< void(unsigned, Operation *, OpBuilder)> annotateFn=nullptr)
Unrolls this for operation by the specified unroll factor.
Definition: Utils.cpp:376
LogicalResult outlineIfOp(RewriterBase &b, scf::IfOp ifOp, func::FuncOp *thenFn, StringRef thenFnName, func::FuncOp *elseFn, StringRef elseFnName)
Outline the then and/or else regions of ifOp as follows:
Definition: Utils.cpp:223
SmallVector< scf::ForOp > replaceLoopNestWithNewYields(RewriterBase &rewriter, MutableArrayRef< scf::ForOp > loopNest, ValueRange newIterOperands, const NewYieldValuesFn &newYieldValuesFn, bool replaceIterOperandsUsesInLoop=true)
Update a perfectly nested loop nest to yield new values from the innermost loop and propagating it up...
Definition: Utils.cpp:40
LogicalResult coalescePerfectlyNestedSCFForLoops(scf::ForOp op)
Walk an affine.for to find a band to coalesce.
Definition: Utils.cpp:982
std::pair< Loops, Loops > TileLoops
Definition: Utils.h:145
void collapseParallelLoops(RewriterBase &rewriter, scf::ParallelOp loops, ArrayRef< std::vector< unsigned >> combinedDimensions)
Take the ParallelLoop and for each set of dimension indices, combine them into a single dimension.
Definition: Utils.cpp:1057
std::function< SmallVector< Value >(OpBuilder &b, Location loc, ArrayRef< BlockArgument > newBbArgs)> NewYieldValuesFn
A function that returns the additional yielded values during replaceWithAdditionalYields.
Loops tilePerfectlyNested(scf::ForOp rootForOp, ArrayRef< Value > sizes)
Tile a nest of scf::ForOp loops rooted at rootForOp with the given (parametric) sizes.
Definition: Utils.cpp:1301
LogicalResult loopUnrollJamByFactor(scf::ForOp forOp, uint64_t unrollFactor)
Unrolls and jams this scf.for operation by the specified unroll factor.
Definition: Utils.cpp:511
bool getInnermostParallelLoops(Operation *rootOp, SmallVectorImpl< scf::ParallelOp > &result)
Get a list of innermost parallel loops contained in rootOp.
Definition: Utils.cpp:246
SmallVector< Loops, 8 > tile(ArrayRef< scf::ForOp > forOps, ArrayRef< Value > sizes, ArrayRef< scf::ForOp > targets)
Performs tiling fo imperfectly nested loops (with interchange) by strip-mining the forOps by sizes an...
Definition: Utils.cpp:1278
std::optional< std::pair< AffineExpr, AffineExpr > > getSCFMinMaxExpr(Value value, SmallVectorImpl< Value > &dims, SmallVectorImpl< Value > &symbols, llvm::function_ref< bool(Operation *)> loopFilter=nullptr)
Return the min/max expressions for value if it is an induction variable from scf.for or scf....
FailureOr< func::FuncOp > outlineSingleBlockRegion(RewriterBase &rewriter, Location loc, Region &region, StringRef funcName, func::CallOp *callOp=nullptr)
Outline a region with a single block into a new FuncOp.
Definition: Utils.cpp:119
void denormalizeInductionVariable(RewriterBase &rewriter, Location loc, Value normalizedIv, OpFoldResult origLb, OpFoldResult origStep)
Get back the original induction variable values after loop normalization.
Definition: Utils.cpp:760
scf::ForallOp fuseIndependentSiblingForallLoops(scf::ForallOp target, scf::ForallOp source, RewriterBase &rewriter)
Given two scf.forall loops, target and source, fuses target into source.
Definition: Utils.cpp:1360
LogicalResult coalesceLoops(MutableArrayRef< scf::ForOp > loops)
Replace a perfect nest of "for" loops with a single linearized loop.
Definition: Utils.cpp:974
scf::ForOp fuseIndependentSiblingForLoops(scf::ForOp target, scf::ForOp source, RewriterBase &rewriter)
Given two scf.for loops, target and source, fuses target into source.
Definition: Utils.cpp:1412
TileLoops extractFixedOuterLoops(scf::ForOp rootFOrOp, ArrayRef< int64_t > sizes)
Definition: Utils.cpp:1318
Range emitNormalizedLoopBounds(RewriterBase &rewriter, Location loc, OpFoldResult lb, OpFoldResult ub, OpFoldResult step)
Materialize bounds and step of a zero-based and unit-step loop derived by normalizing the specified b...
Definition: Utils.cpp:689
FailureOr< scf::ForallOp > normalizeForallOp(RewriterBase &rewriter, scf::ForallOp forallOp)
Normalize an scf.forall operation.
Definition: Utils.cpp:1462
Represents a range (offset, size, and stride) where each element of the triple may be dynamic or stat...