MLIR  16.0.0git
LoopUtils.h
Go to the documentation of this file.
1 //===- LoopUtils.h - Loop 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 prototypes for various loop transformation utility
10 // methods: these are not passes by themselves but are used either by passes,
11 // optimization sequences, or in turn by other transformation utilities.
12 //
13 //===----------------------------------------------------------------------===//
14 
15 #ifndef MLIR_DIALECT_AFFINE_LOOPUTILS_H
16 #define MLIR_DIALECT_AFFINE_LOOPUTILS_H
17 
18 #include "mlir/IR/Block.h"
19 #include "mlir/Support/LLVM.h"
21 
22 namespace mlir {
23 class AffineForOp;
24 class AffineMap;
25 class LoopLikeOpInterface;
26 struct MemRefRegion;
27 class OpBuilder;
28 class Value;
29 class ValueRange;
30 
31 namespace func {
32 class FuncOp;
33 } // namespace func
34 
35 namespace scf {
36 class ForOp;
37 class ParallelOp;
38 } // namespace scf
39 
40 /// Unrolls this for operation completely if the trip count is known to be
41 /// constant. Returns failure otherwise.
42 LogicalResult loopUnrollFull(AffineForOp forOp);
43 
44 /// Unrolls this for operation by the specified unroll factor. Returns failure
45 /// if the loop cannot be unrolled either due to restrictions or due to invalid
46 /// unroll factors. Requires positive loop bounds and step. If specified,
47 /// annotates the Ops in each unrolled iteration by applying `annotateFn`.
48 /// When `cleanUpUnroll` is true, we can ensure the cleanup loop is unrolled
49 /// regardless of the unroll factor.
51  AffineForOp forOp, uint64_t unrollFactor,
52  function_ref<void(unsigned, Operation *, OpBuilder)> annotateFn = nullptr,
53  bool cleanUpUnroll = false);
54 
55 /// Unrolls this loop by the specified unroll factor or its trip count,
56 /// whichever is lower.
57 LogicalResult loopUnrollUpToFactor(AffineForOp forOp, uint64_t unrollFactor);
58 
59 /// Returns true if `loops` is a perfectly nested loop nest, where loops appear
60 /// in it from outermost to innermost.
61 bool LLVM_ATTRIBUTE_UNUSED isPerfectlyNested(ArrayRef<AffineForOp> loops);
62 
63 /// Get perfectly nested sequence of loops starting at root of loop nest
64 /// (the first op being another AffineFor, and the second op - a terminator).
65 /// A loop is perfectly nested iff: the first op in the loop's body is another
66 /// AffineForOp, and the second op is a terminator).
68  AffineForOp root);
69 
70 /// Unrolls and jams this loop by the specified factor. `forOp` can be a loop
71 /// with iteration arguments performing supported reductions and its inner loops
72 /// can have iteration arguments. Returns success if the loop is successfully
73 /// unroll-jammed.
74 LogicalResult loopUnrollJamByFactor(AffineForOp forOp,
75  uint64_t unrollJamFactor);
76 
77 /// Unrolls and jams this loop by the specified factor or by the trip count (if
78 /// constant), whichever is lower.
79 LogicalResult loopUnrollJamUpToFactor(AffineForOp forOp,
80  uint64_t unrollJamFactor);
81 
82 /// Promotes the loop body of a AffineForOp to its containing block if the loop
83 /// was known to have a single iteration.
84 LogicalResult promoteIfSingleIteration(AffineForOp forOp);
85 
86 /// Promotes all single iteration AffineForOp's in the Function, i.e., moves
87 /// their body into the containing Block.
88 void promoteSingleIterationLoops(func::FuncOp f);
89 
90 /// Skew the operations in an affine.for's body with the specified
91 /// operation-wise shifts. The shifts are with respect to the original execution
92 /// order, and are multiplied by the loop 'step' before being applied. If
93 /// `unrollPrologueEpilogue` is set, fully unroll the prologue and epilogue
94 /// loops when possible.
95 LogicalResult affineForOpBodySkew(AffineForOp forOp, ArrayRef<uint64_t> shifts,
96  bool unrollPrologueEpilogue = false);
97 
98 /// Identify valid and profitable bands of loops to tile. This is currently just
99 /// a temporary placeholder to test the mechanics of tiled code generation.
100 /// Returns all maximal outermost perfect loop nests to tile.
101 void getTileableBands(func::FuncOp f,
102  std::vector<SmallVector<AffineForOp, 6>> *bands);
103 
104 /// Tiles the specified band of perfectly nested loops creating tile-space loops
105 /// and intra-tile loops. A band is a contiguous set of loops.
108  ArrayRef<unsigned> tileSizes,
109  SmallVectorImpl<AffineForOp> *tiledNest = nullptr);
110 
111 /// Tiles the specified band of perfectly nested loops creating tile-space
112 /// loops and intra-tile loops, using SSA values as tiling parameters. A band
113 /// is a contiguous set of loops.
116  SmallVectorImpl<AffineForOp> *tiledNest = nullptr);
117 
118 /// Performs loop interchange on 'forOpA' and 'forOpB'. Requires that 'forOpA'
119 /// and 'forOpB' are part of a perfectly nested sequence of loops.
120 void interchangeLoops(AffineForOp forOpA, AffineForOp forOpB);
121 
122 /// Checks if the loop interchange permutation 'loopPermMap', of the perfectly
123 /// nested sequence of loops in 'loops', would violate dependences (loop 'i' in
124 /// 'loops' is mapped to location 'j = 'loopPermMap[i]' in the interchange).
126  ArrayRef<unsigned> loopPermMap);
127 
128 /// Performs a loop permutation on a perfectly nested loop nest `inputNest`
129 /// (where the contained loops appear from outer to inner) as specified by the
130 /// permutation `permMap`: loop 'i' in `inputNest` is mapped to location
131 /// 'loopPermMap[i]', where positions 0, 1, ... are from the outermost position
132 /// to inner. Returns the position in `inputNest` of the AffineForOp that
133 /// becomes the new outermost loop of this nest. This method always succeeds,
134 /// asserts out on invalid input / specifications.
135 unsigned permuteLoops(MutableArrayRef<AffineForOp> inputNest,
136  ArrayRef<unsigned> permMap);
137 
138 // Sinks all sequential loops to the innermost levels (while preserving
139 // relative order among them) and moves all parallel loops to the
140 // outermost (while again preserving relative order among them).
141 // Returns AffineForOp of the root of the new loop nest after loop interchanges.
142 AffineForOp sinkSequentialLoops(AffineForOp forOp);
143 
144 /// Performs tiling fo imperfectly nested loops (with interchange) by
145 /// strip-mining the `forOps` by `sizes` and sinking them, in their order of
146 /// occurrence in `forOps`, under each of the `targets`.
147 /// Returns the new AffineForOps, one per each of (`forOps`, `targets`) pair,
148 /// nested immediately under each of `targets`.
150  ArrayRef<uint64_t> sizes,
151  ArrayRef<AffineForOp> targets);
152 
153 /// Performs tiling (with interchange) by strip-mining the `forOps` by `sizes`
154 /// and sinking them, in their order of occurrence in `forOps`, under `target`.
155 /// Returns the new AffineForOps, one per `forOps`, nested immediately under
156 /// `target`.
158  ArrayRef<uint64_t> sizes, AffineForOp target);
159 
160 /// Explicit copy / DMA generation options for mlir::affineDataCopyGenerate.
162  // True if DMAs should be generated instead of point-wise copies.
164  // The slower memory space from which data is to be moved.
165  unsigned slowMemorySpace;
166  // Memory space of the faster one (typically a scratchpad).
167  unsigned fastMemorySpace;
168  // Memory space to place tags in: only meaningful for DMAs.
169  unsigned tagMemorySpace;
170  // Capacity of the fast memory space in bytes.
172 };
173 
174 /// Performs explicit copying for the contiguous sequence of operations in the
175 /// block iterator range [`begin', `end'), where `end' can't be past the
176 /// terminator of the block (since additional operations are potentially
177 /// inserted right before `end`. `copyOptions` provides various parameters, and
178 /// the output argument `copyNests` is the set of all copy nests inserted, each
179 /// represented by its root affine.for. Since we generate alloc's and dealloc's
180 /// for all fast buffers (before and after the range of operations resp. or at a
181 /// hoisted position), all of the fast memory capacity is assumed to be
182 /// available for processing this block range. When 'filterMemRef' is specified,
183 /// copies are only generated for the provided MemRef. Returns success if the
184 /// explicit copying succeeded for all memrefs on which affine load/stores were
185 /// encountered.
187  const AffineCopyOptions &copyOptions,
188  Optional<Value> filterMemRef,
189  DenseSet<Operation *> &copyNests);
190 
191 /// A convenience version of affineDataCopyGenerate for all ops in the body of
192 /// an AffineForOp.
193 LogicalResult affineDataCopyGenerate(AffineForOp forOp,
194  const AffineCopyOptions &copyOptions,
195  Optional<Value> filterMemRef,
196  DenseSet<Operation *> &copyNests);
197 
198 /// Result for calling generateCopyForMemRegion.
200  // Number of bytes used by alloc.
201  uint64_t sizeInBytes;
202 
203  // The newly created buffer allocation.
205 
206  // Generated loop nest for copying data between the allocated buffer and the
207  // original memref.
209 };
210 
211 /// generateCopyForMemRegion is similar to affineDataCopyGenerate, but works
212 /// with a single memref region. `memrefRegion` is supposed to contain analysis
213 /// information within analyzedOp. The generated prologue and epilogue always
214 /// surround `analyzedOp`.
215 ///
216 /// Note that `analyzedOp` is a single op for API convenience, and the
217 /// [begin, end) version can be added as needed.
218 ///
219 /// Also note that certain options in `copyOptions` aren't looked at anymore,
220 /// like slowMemorySpace.
222  Operation *analyzedOp,
223  const AffineCopyOptions &copyOptions,
224  CopyGenerateResult &result);
225 
226 /// Replace a perfect nest of "for" loops with a single linearized loop. Assumes
227 /// `loops` contains a list of perfectly nested loops outermost to innermost
228 /// that are normalized (step one and lower bound of zero) and with bounds and
229 /// steps independent of any loop induction variable involved in the nest.
230 /// Coalescing affine.for loops is not always possible, i.e., the result may not
231 /// be representable using affine.for.
233 
234 /// Maps `forOp` for execution on a parallel grid of virtual `processorIds` of
235 /// size given by `numProcessors`. This is achieved by embedding the SSA values
236 /// corresponding to `processorIds` and `numProcessors` into the bounds and step
237 /// of the `forOp`. No check is performed on the legality of the rewrite, it is
238 /// the caller's responsibility to ensure legality.
239 ///
240 /// Requires that `processorIds` and `numProcessors` have the same size and that
241 /// for each idx, `processorIds`[idx] takes, at runtime, all values between 0
242 /// and `numProcessors`[idx] - 1. This corresponds to traditional use cases for:
243 /// 1. GPU (threadIdx, get_local_id(), ...)
244 /// 2. MPI (MPI_Comm_rank)
245 /// 3. OpenMP (omp_get_thread_num)
246 ///
247 /// Example:
248 /// Assuming a 2-d grid with processorIds = [blockIdx.x, threadIdx.x] and
249 /// numProcessors = [gridDim.x, blockDim.x], the loop:
250 ///
251 /// ```
252 /// scf.for %i = %lb to %ub step %step {
253 /// ...
254 /// }
255 /// ```
256 ///
257 /// is rewritten into a version resembling the following pseudo-IR:
258 ///
259 /// ```
260 /// scf.for %i = %lb + %step * (threadIdx.x + blockIdx.x * blockDim.x)
261 /// to %ub step %gridDim.x * blockDim.x * %step {
262 /// ...
263 /// }
264 /// ```
265 void mapLoopToProcessorIds(scf::ForOp forOp, ArrayRef<Value> processorId,
266  ArrayRef<Value> numProcessors);
267 
268 /// Gathers all AffineForOps in 'func.func' grouped by loop depth.
269 void gatherLoops(func::FuncOp func,
270  std::vector<SmallVector<AffineForOp, 2>> &depthToLoops);
271 
272 /// Creates an AffineForOp while ensuring that the lower and upper bounds are
273 /// canonicalized, i.e., unused and duplicate operands are removed, any constant
274 /// operands propagated/folded in, and duplicate bound maps dropped.
276  ValueRange lbOperands,
277  AffineMap lbMap,
278  ValueRange ubOperands,
279  AffineMap ubMap, int64_t step = 1);
280 
281 /// Separates full tiles from partial tiles for a perfect nest `nest` by
282 /// generating a conditional guard that selects between the full tile version
283 /// and the partial tile version using an AffineIfOp. The original loop nest
284 /// is replaced by this guarded two version form.
285 ///
286 /// affine.if (cond)
287 /// // full_tile
288 /// else
289 /// // partial tile
290 ///
293  SmallVectorImpl<AffineForOp> *fullTileNest = nullptr);
294 
295 } // namespace mlir
296 
297 #endif // MLIR_DIALECT_AFFINE_LOOPUTILS_H
Include the generated interface declarations.
LogicalResult affineForOpBodySkew(AffineForOp forOp, ArrayRef< uint64_t > shifts, bool unrollPrologueEpilogue=false)
Skew the operations in an affine.for&#39;s body with the specified operation-wise shifts.
Definition: LoopUtils.cpp:235
Operation is a basic unit of execution within MLIR.
Definition: Operation.h:28
bool LLVM_ATTRIBUTE_UNUSED isPerfectlyNested(ArrayRef< AffineForOp > loops)
Returns true if loops is a perfectly nested loop nest, where loops appear in it from outermost to inn...
Definition: LoopUtils.cpp:1474
unsigned permuteLoops(MutableArrayRef< AffineForOp > inputNest, ArrayRef< unsigned > permMap)
Performs a loop permutation on a perfectly nested loop nest inputNest (where the contained loops appe...
Definition: LoopUtils.cpp:1496
LogicalResult coalesceLoops(MutableArrayRef< AffineForOp > loops)
Replace a perfect nest of "for" loops with a single linearized loop.
Definition: LoopUtils.cpp:1724
LogicalResult promoteIfSingleIteration(AffineForOp forOp)
Promotes the loop body of a AffineForOp to its containing block if the loop was known to have a singl...
Definition: LoopUtils.cpp:131
LogicalResult loopUnrollByFactor(AffineForOp forOp, uint64_t unrollFactor, function_ref< void(unsigned, Operation *, OpBuilder)> annotateFn=nullptr, bool cleanUpUnroll=false)
Unrolls this for operation by the specified unroll factor.
Definition: LoopUtils.cpp:1091
void gatherLoops(func::FuncOp func, std::vector< SmallVector< AffineForOp, 2 >> &depthToLoops)
Gathers all AffineForOps in &#39;func.func&#39; grouped by loop depth.
Definition: LoopUtils.cpp:2620
This class defines the main interface for locations in MLIR and acts as a non-nullable wrapper around...
Definition: Location.h:48
LogicalResult tilePerfectlyNested(MutableArrayRef< AffineForOp > input, ArrayRef< unsigned > tileSizes, SmallVectorImpl< AffineForOp > *tiledNest=nullptr)
Tiles the specified band of perfectly nested loops creating tile-space loops and intra-tile loops...
Definition: LoopUtils.cpp:859
Result for calling generateCopyForMemRegion.
Definition: LoopUtils.h:199
LogicalResult loopUnrollJamUpToFactor(AffineForOp forOp, uint64_t unrollJamFactor)
Unrolls and jams this loop by the specified factor or by the trip count (if constant), whichever is lower.
Definition: LoopUtils.cpp:1158
This class represents an efficient way to signal success or failure.
Definition: LogicalResult.h:26
LogicalResult loopUnrollFull(AffineForOp forOp)
Unrolls this for operation completely if the trip count is known to be constant.
Definition: LoopUtils.cpp:974
OpListType::iterator iterator
Definition: Block.h:131
LogicalResult tilePerfectlyNestedParametric(MutableArrayRef< AffineForOp > input, ArrayRef< Value > tileSizes, SmallVectorImpl< AffineForOp > *tiledNest=nullptr)
Tiles the specified band of perfectly nested loops creating tile-space loops and intra-tile loops...
Definition: LoopUtils.cpp:903
SmallVector< SmallVector< AffineForOp, 8 >, 8 > tile(ArrayRef< AffineForOp > forOps, ArrayRef< uint64_t > sizes, ArrayRef< AffineForOp > targets)
Performs tiling fo imperfectly nested loops (with interchange) by strip-mining the forOps by sizes an...
Definition: LoopUtils.cpp:1701
void getPerfectlyNestedLoops(SmallVectorImpl< AffineForOp > &nestedLoops, AffineForOp root)
Get perfectly nested sequence of loops starting at root of loop nest (the first op being another Affi...
Definition: LoopUtils.cpp:945
AffineForOp sinkSequentialLoops(AffineForOp forOp)
Definition: LoopUtils.cpp:1566
LogicalResult loopUnrollUpToFactor(AffineForOp forOp, uint64_t unrollFactor)
Unrolls this loop by the specified unroll factor or its trip count, whichever is lower.
Definition: LoopUtils.cpp:989
uint64_t fastMemCapacityBytes
Definition: LoopUtils.h:171
A multi-dimensional affine map Affine map&#39;s are immutable like Type&#39;s, and they are uniqued...
Definition: AffineMap.h:42
LogicalResult separateFullTiles(MutableArrayRef< AffineForOp > nest, SmallVectorImpl< AffineForOp > *fullTileNest=nullptr)
Separates full tiles from partial tiles for a perfect nest nest by generating a conditional guard tha...
Definition: LoopUtils.cpp:2793
Explicit copy / DMA generation options for mlir::affineDataCopyGenerate.
Definition: LoopUtils.h:161
LogicalResult affineDataCopyGenerate(Block::iterator begin, Block::iterator end, const AffineCopyOptions &copyOptions, Optional< Value > filterMemRef, DenseSet< Operation *> &copyNests)
Performs explicit copying for the contiguous sequence of operations in the block iterator range [`beg...
Definition: LoopUtils.cpp:2377
LogicalResult generateCopyForMemRegion(const MemRefRegion &memrefRegion, Operation *analyzedOp, const AffineCopyOptions &copyOptions, CopyGenerateResult &result)
generateCopyForMemRegion is similar to affineDataCopyGenerate, but works with a single memref region...
Definition: LoopUtils.cpp:2576
void getTileableBands(func::FuncOp f, std::vector< SmallVector< AffineForOp, 6 >> *bands)
Identify valid and profitable bands of loops to tile.
Definition: LoopUtils.cpp:962
void promoteSingleIterationLoops(func::FuncOp f)
Promotes all single iteration AffineForOp&#39;s in the Function, i.e., moves their body into the containi...
void interchangeLoops(AffineForOp forOpA, AffineForOp forOpB)
Performs loop interchange on &#39;forOpA&#39; and &#39;forOpB&#39;.
Definition: LoopUtils.cpp:1403
LogicalResult loopUnrollJamByFactor(AffineForOp forOp, uint64_t unrollJamFactor)
Unrolls and jams this loop by the specified factor.
Definition: LoopUtils.cpp:1209
bool isValidLoopInterchangePermutation(ArrayRef< AffineForOp > loops, ArrayRef< unsigned > loopPermMap)
Checks if the loop interchange permutation &#39;loopPermMap&#39;, of the perfectly nested sequence of loops i...
Definition: LoopUtils.cpp:1460
A region of a memref&#39;s data space; this is typically constructed by analyzing load/store op&#39;s on this...
Definition: Utils.h:247
void mapLoopToProcessorIds(scf::ForOp forOp, ArrayRef< Value > processorId, ArrayRef< Value > numProcessors)
Maps forOp for execution on a parallel grid of virtual processorIds of size given by numProcessors...
Definition: LoopUtils.cpp:1838
This class helps build Operations.
Definition: Builders.h:193
This class provides an abstraction over the different types of ranges over Values.
Definition: ValueRange.h:345
AffineForOp createCanonicalizedAffineForOp(OpBuilder b, Location loc, ValueRange lbOperands, AffineMap lbMap, ValueRange ubOperands, AffineMap ubMap, int64_t step=1)
Creates an AffineForOp while ensuring that the lower and upper bounds are canonicalized, i.e., unused and duplicate operands are removed, any constant operands propagated/folded in, and duplicate bound maps dropped.
Definition: LoopUtils.cpp:2636