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