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 
16 #include "mlir/Dialect/SCF/Utils/AffineCanonicalizationUtils.h"
17 #include "mlir/Support/LLVM.h"
18 #include "llvm/ADT/ArrayRef.h"
19 
20 namespace mlir {
21 
22 class AffineMap;
23 class ConversionTarget;
24 struct LogicalResult;
25 class MLIRContext;
26 class Region;
27 class RewriterBase;
28 class TypeConverter;
29 class RewritePatternSet;
30 class Operation;
31 class Value;
32 class ValueRange;
33 class PatternRewriter;
34 
35 namespace scf {
36 
37 class IfOp;
38 class ForOp;
39 class ParallelOp;
40 
41 /// Fuses all adjacent scf.parallel operations with identical bounds and step
42 /// into one scf.parallel operations. Uses a naive aliasing and dependency
43 /// analysis.
44 void naivelyFuseParallelOps(Region &region);
45 
46 /// Rewrite a for loop with bounds/step that potentially do not divide evenly
47 /// into a for loop where the step divides the iteration space evenly, followed
48 /// by another scf.for for the last (partial) iteration (if any; returned via
49 /// `partialIteration`). This transformation is called "loop peeling".
50 ///
51 /// This transformation is beneficial for a wide range of transformations such
52 /// as vectorization or loop tiling: It enables additional canonicalizations
53 /// inside the peeled loop body such as rewriting masked loads into unmaked
54 /// loads.
55 ///
56 /// E.g., assuming a lower bound of 0 (for illustration purposes):
57 /// ```
58 /// scf.for %iv = %c0 to %ub step %c4 {
59 ///   (loop body)
60 /// }
61 /// ```
62 /// is rewritten into the following pseudo IR:
63 /// ```
64 /// %newUb = %ub - (%ub mod %c4)
65 /// scf.for %iv = %c0 to %newUb step %c4 {
66 ///   (loop body)
67 /// }
68 /// scf.for %iv2 = %newUb to %ub {
69 ///   (loop body)
70 /// }
71 /// ```
72 ///
73 /// After loop peeling, this function tries to simplify/canonicalize affine.min
74 /// and affine.max ops in the body of the peeled loop and in the body of the
75 /// partial iteration loop, taking advantage of the fact that the peeled loop
76 /// has only "full" iterations. This canonicalization is expected to enable
77 /// further canonicalization opportunities through other patterns.
78 ///
79 /// The return value indicates whether the loop was rewritten or not. Loops are
80 /// not rewritten if:
81 /// * Loop step size is 1 or
82 /// * Loop bounds and step size are static, and step already divides the
83 ///   iteration space evenly.
84 ///
85 /// Note: This function rewrites the given scf.for loop in-place and creates a
86 /// new scf.for operation for the last iteration. It replaces all uses of the
87 /// unpeeled loop with the results of the newly generated scf.for.
88 LogicalResult peelAndCanonicalizeForLoop(RewriterBase &rewriter, ForOp forOp,
89                                          scf::ForOp &partialIteration);
90 
91 /// Tile a parallel loop of the form
92 ///   scf.parallel (%i0, %i1) = (%arg0, %arg1) to (%arg2, %arg3)
93 ///                                             step (%arg4, %arg5)
94 ///
95 /// into
96 ///   scf.parallel (%i0, %i1) = (%arg0, %arg1) to (%arg2, %arg3)
97 ///                                             step (%arg4*tileSize[0],
98 ///                                                   %arg5*tileSize[1])
99 ///     scf.parallel (%j0, %j1) = (0, 0) to (min(tileSize[0], %arg2-%j0)
100 ///                                           min(tileSize[1], %arg3-%j1))
101 ///                                        step (%arg4, %arg5)
102 /// The old loop is replaced with the new one.
103 ///
104 /// The function returns the resulting ParallelOps, i.e. {outer_loop_op,
105 /// inner_loop_op}.
106 std::pair<ParallelOp, ParallelOp>
107 tileParallelLoop(ParallelOp op, llvm::ArrayRef<int64_t> tileSizes,
108                  bool noMinMaxBounds);
109 
110 /// Populates patterns for SCF structural type conversions and sets up the
111 /// provided ConversionTarget with the appropriate legality configuration for
112 /// the ops to get converted properly.
113 ///
114 /// A "structural" type conversion is one where the underlying ops are
115 /// completely agnostic to the actual types involved and simply need to update
116 /// their types. An example of this is scf.if -- the scf.if op and the
117 /// corresponding scf.yield ops need to update their types accordingly to the
118 /// TypeConverter, but otherwise don't care what type conversions are happening.
119 void populateSCFStructuralTypeConversionsAndLegality(
120     TypeConverter &typeConverter, RewritePatternSet &patterns,
121     ConversionTarget &target);
122 
123 /// Options to dictate how loops should be pipelined.
124 struct PipeliningOption {
125   /// Lambda returning all the operation in the forOp, with their stage, in the
126   /// order picked for the pipelined loop.
127   using GetScheduleFnType = std::function<void(
128       scf::ForOp, std::vector<std::pair<Operation *, unsigned>> &)>;
129   GetScheduleFnType getScheduleFn = nullptr;
130   enum class PipelinerPart {
131     Prologue,
132     Kernel,
133     Epilogue,
134   };
135   /// Lambda called by the pipeliner to allow the user to annotate the IR while
136   /// it is generated.
137   /// The callback passes the operation created along with the part of the
138   /// pipeline and the iteration index. The iteration index is always 0 for the
139   /// kernel. For the prologue and epilogue, it corresponds to the iteration
140   /// peeled out of the loop in the range [0, maxStage[.
141   using AnnotationlFnType =
142       std::function<void(Operation *, PipelinerPart, unsigned)>;
143   AnnotationlFnType annotateFn = nullptr;
144 
145   /// Control whether the epilogue should be peeled out of the loop or
146   /// operations should be predicated to skip the early stages in the last loop
147   /// iterations. If the epilogue is predicated; the user needs to provide a
148   /// lambda to generate the predicated version of operations.
149   bool peelEpilogue = true;
150 
151   // Lamdba to predicate operations when the prologue or epilogue are not
152   // peeled. This takes the original operation, an i1 predicate value and the
153   // pattern rewriter.
154   using PredicateOpFn =
155       std::function<Operation *(Operation *, Value, PatternRewriter &)>;
156   PredicateOpFn predicateFn = nullptr;
157 
158   // TODO: add option to decide if the prologue should be peeled.
159 };
160 
161 /// Populate patterns for SCF software pipelining transformation. See the
162 /// ForLoopPipeliningPattern for the transformation details.
163 void populateSCFLoopPipeliningPatterns(RewritePatternSet &patterns,
164                                        const PipeliningOption &options);
165 
166 /// Populate patterns for canonicalizing operations inside SCF loop bodies.
167 /// At the moment, only affine.min/max computations with iteration variables,
168 /// loop bounds and loop steps are canonicalized.
169 void populateSCFForLoopCanonicalizationPatterns(RewritePatternSet &patterns);
170 
171 } // namespace scf
172 } // namespace mlir
173 
174 #endif // MLIR_DIALECT_SCF_TRANSFORMS_TRANSFORMS_H_
175