xref: /llvm-project/mlir/include/mlir/Dialect/SCF/Utils/AffineCanonicalizationUtils.h (revision db791b278a414fb6df1acc1799adcf11d8fb9169)
1 //===- AffineCanonicalizationUtils.h ----------------------------*- 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 utility functions to canonicalize affine ops
10 // within SCF op regions.
11 //
12 //===----------------------------------------------------------------------===//
13 
14 #ifndef MLIR_DIALECT_SCF_UTILS_AFFINECANONICALIZATIONUTILS_H_
15 #define MLIR_DIALECT_SCF_UTILS_AFFINECANONICALIZATIONUTILS_H_
16 
17 #include "mlir/Support/LLVM.h"
18 
19 namespace mlir {
20 class AffineMap;
21 class Operation;
22 class OpFoldResult;
23 class RewriterBase;
24 class Value;
25 class ValueRange;
26 
27 namespace affine {
28 class FlatAffineValueConstraints;
29 } // namespace affine
30 
31 namespace scf {
32 class IfOp;
33 
34 /// Match "for loop"-like operations: If the first parameter is an iteration
35 /// variable, return lower/upper bounds via the second/third parameter and the
36 /// step size via the last parameter. The function should return `success` in
37 /// that case. If the first parameter is not an iteration variable, return
38 /// `failure`.
39 using LoopMatcherFn = function_ref<LogicalResult(
40     Value, OpFoldResult &, OpFoldResult &, OpFoldResult &)>;
41 
42 /// Match "for loop"-like operations from the SCF dialect.
43 LogicalResult matchForLikeLoop(Value iv, OpFoldResult &lb, OpFoldResult &ub,
44                                OpFoldResult &step);
45 
46 /// Populate the given constraint set with induction variable constraints of a
47 /// "for" loop with the given range and step.
48 LogicalResult addLoopRangeConstraints(affine::FlatAffineValueConstraints &cstr,
49                                       Value iv, OpFoldResult lb,
50                                       OpFoldResult ub, OpFoldResult step);
51 
52 /// Try to canonicalize the given affine.min/max operation in the context of
53 /// for `loops` with a known range.
54 ///
55 /// `loopMatcher` is used to retrieve loop bounds and the step size for a given
56 /// iteration variable.
57 ///
58 /// Note: `loopMatcher` allows this function to be used with any "for loop"-like
59 /// operation (scf.for, scf.parallel and even ops defined in other dialects).
60 LogicalResult canonicalizeMinMaxOpInLoop(RewriterBase &rewriter, Operation *op,
61                                          LoopMatcherFn loopMatcher);
62 
63 /// Try to simplify the given affine.min/max operation `op` after loop peeling.
64 /// This function can simplify min/max operations such as (ub is the previous
65 /// upper bound of the unpeeled loop):
66 /// ```
67 /// #map = affine_map<(d0)[s0, s1] -> (s0, -d0 + s1)>
68 /// %r = affine.min #map(%iv)[%step, %ub]
69 /// ```
70 /// and rewrites them into (in the case the peeled loop):
71 /// ```
72 /// %r = %step
73 /// ```
74 /// min/max operations inside the partial iteration are rewritten in a similar
75 /// way.
76 LogicalResult rewritePeeledMinMaxOp(RewriterBase &rewriter, Operation *op,
77                                     Value iv, Value ub, Value step,
78                                     bool insideLoop);
79 
80 } // namespace scf
81 } // namespace mlir
82 
83 #endif // MLIR_DIALECT_SCF_UTILS_AFFINECANONICALIZATIONUTILS_H_
84