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 struct LogicalResult; 22 class Operation; 23 class OpFoldResult; 24 class RewriterBase; 25 class Value; 26 class ValueRange; 27 28 namespace scf { 29 class IfOp; 30 31 /// Match "for loop"-like operations: If the first parameter is an iteration 32 /// variable, return lower/upper bounds via the second/third parameter and the 33 /// step size via the last parameter. The function should return `success` in 34 /// that case. If the first parameter is not an iteration variable, return 35 /// `failure`. 36 using LoopMatcherFn = function_ref<LogicalResult( 37 Value, OpFoldResult &, OpFoldResult &, OpFoldResult &)>; 38 39 /// Try to canonicalize an min/max operations in the context of for `loops` with 40 /// a known range. 41 /// 42 /// `map` is the body of the min/max operation and `operands` are the SSA values 43 /// that the dimensions and symbols are bound to; dimensions are listed first. 44 /// If `isMin`, the operation is a min operation; otherwise, a max operation. 45 /// `loopMatcher` is used to retrieve loop bounds and the step size for a given 46 /// iteration variable. 47 /// 48 /// Note: `loopMatcher` allows this function to be used with any "for loop"-like 49 /// operation (scf.for, scf.parallel and even ops defined in other dialects). 50 LogicalResult canonicalizeMinMaxOpInLoop(RewriterBase &rewriter, Operation *op, 51 AffineMap map, ValueRange operands, 52 bool isMin, LoopMatcherFn loopMatcher); 53 54 /// Try to simplify a min/max operation `op` after loop peeling. This function 55 /// can simplify min/max operations such as (ub is the previous upper bound of 56 /// the unpeeled loop): 57 /// ``` 58 /// #map = affine_map<(d0)[s0, s1] -> (s0, -d0 + s1)> 59 /// %r = affine.min #affine.min #map(%iv)[%step, %ub] 60 /// ``` 61 /// and rewrites them into (in the case the peeled loop): 62 /// ``` 63 /// %r = %step 64 /// ``` 65 /// min/max operations inside the partial iteration are rewritten in a similar 66 /// way. 67 LogicalResult rewritePeeledMinMaxOp(RewriterBase &rewriter, Operation *op, 68 AffineMap map, ValueRange operands, 69 bool isMin, Value iv, Value ub, Value step, 70 bool insideLoop); 71 72 } // namespace scf 73 } // namespace mlir 74 75 #endif // MLIR_DIALECT_SCF_UTILS_AFFINECANONICALIZATIONUTILS_H_ 76