|
| 1 | +//===- SimplifyAffineMinMax.cpp - Simplify affine min/max ops -------------===// |
| 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 file implements a transform to simplify mix/max affine operations. |
| 10 | +// |
| 11 | +//===----------------------------------------------------------------------===// |
| 12 | + |
| 13 | +#include "mlir/Dialect/Affine/IR/AffineOps.h" |
| 14 | +#include "mlir/Dialect/Affine/Transforms/Transforms.h" |
| 15 | +#include "mlir/IR/PatternMatch.h" |
| 16 | +#include "mlir/Interfaces/ValueBoundsOpInterface.h" |
| 17 | +#include "mlir/Transforms/GreedyPatternRewriteDriver.h" |
| 18 | +#include "llvm/ADT/IntEqClasses.h" |
| 19 | +#include "llvm/Support/Debug.h" |
| 20 | + |
| 21 | +#define DEBUG_TYPE "affine-min-max" |
| 22 | +#define DBGS() (llvm::dbgs() << "[" DEBUG_TYPE << "]: ") |
| 23 | + |
| 24 | +using namespace mlir; |
| 25 | +using namespace mlir::affine; |
| 26 | + |
| 27 | +/// Simplifies an affine min/max operation by proving there's a lower or upper |
| 28 | +/// bound. |
| 29 | +template <typename AffineOp> |
| 30 | +static bool simplifyAffineMinMaxOp(RewriterBase &rewriter, AffineOp affineOp) { |
| 31 | + using Variable = ValueBoundsConstraintSet::Variable; |
| 32 | + using ComparisonOperator = ValueBoundsConstraintSet::ComparisonOperator; |
| 33 | + |
| 34 | + AffineMap affineMap = affineOp.getMap(); |
| 35 | + ValueRange operands = affineOp.getOperands(); |
| 36 | + static constexpr bool isMin = std::is_same_v<AffineOp, AffineMinOp>; |
| 37 | + |
| 38 | + LLVM_DEBUG({ DBGS() << "analyzing value: `" << affineOp << "`\n"; }); |
| 39 | + |
| 40 | + // Create a `Variable` list with values corresponding to each of the results |
| 41 | + // in the affine affineMap. |
| 42 | + SmallVector<Variable> variables = llvm::map_to_vector( |
| 43 | + llvm::iota_range<unsigned>(0u, affineMap.getNumResults(), false), |
| 44 | + [&](unsigned i) { |
| 45 | + return Variable(affineMap.getSliceMap(i, 1), operands); |
| 46 | + }); |
| 47 | + |
| 48 | + // Get the comparison operation. |
| 49 | + ComparisonOperator cmpOp = |
| 50 | + isMin ? ComparisonOperator::LT : ComparisonOperator::GT; |
| 51 | + |
| 52 | + // Find disjoint sets bounded by a common value. |
| 53 | + llvm::IntEqClasses boundedClasses(variables.size()); |
| 54 | + DenseMap<unsigned, Variable *> bounds; |
| 55 | + for (auto &&[i, v] : llvm::enumerate(variables)) { |
| 56 | + unsigned eqClass = boundedClasses.findLeader(i); |
| 57 | + |
| 58 | + // If the class already has a bound continue. |
| 59 | + if (bounds.contains(eqClass)) |
| 60 | + continue; |
| 61 | + |
| 62 | + // Initialize the bound. |
| 63 | + Variable *bound = &v; |
| 64 | + |
| 65 | + LLVM_DEBUG({ |
| 66 | + DBGS() << "- inspecting variable: #" << i << ", with map: `" << v.getMap() |
| 67 | + << "`\n"; |
| 68 | + }); |
| 69 | + |
| 70 | + // Check against the other variables. |
| 71 | + for (size_t j = i + 1; j < variables.size(); ++j) { |
| 72 | + unsigned jEqClass = boundedClasses.findLeader(j); |
| 73 | + // Skip if the class is the same. |
| 74 | + if (jEqClass == eqClass) |
| 75 | + continue; |
| 76 | + |
| 77 | + // Get the bound of the equivalence class or itself. |
| 78 | + Variable *nv = bounds.lookup_or(jEqClass, &variables[j]); |
| 79 | + |
| 80 | + LLVM_DEBUG({ |
| 81 | + DBGS() << "- comparing with variable: #" << jEqClass |
| 82 | + << ", with map: " << nv->getMap() << "\n"; |
| 83 | + }); |
| 84 | + |
| 85 | + // Compare the variables. |
| 86 | + FailureOr<bool> cmpResult = |
| 87 | + ValueBoundsConstraintSet::strongCompare(*bound, cmpOp, *nv); |
| 88 | + |
| 89 | + // The variables cannot be compared. |
| 90 | + if (failed(cmpResult)) { |
| 91 | + LLVM_DEBUG({ |
| 92 | + DBGS() << "-- classes: #" << i << ", #" << jEqClass |
| 93 | + << " cannot be merged\n"; |
| 94 | + }); |
| 95 | + continue; |
| 96 | + } |
| 97 | + |
| 98 | + // Join the equivalent classes and update the bound if necessary. |
| 99 | + LLVM_DEBUG({ |
| 100 | + DBGS() << "-- merging classes: #" << i << ", #" << jEqClass |
| 101 | + << ", is cmp(lhs, rhs): " << *cmpResult << "`\n"; |
| 102 | + }); |
| 103 | + if (*cmpResult) { |
| 104 | + boundedClasses.join(eqClass, jEqClass); |
| 105 | + } else { |
| 106 | + // In this case we have lhs > rhs if isMin == true, or lhs < rhs if |
| 107 | + // isMin == false. |
| 108 | + bound = nv; |
| 109 | + boundedClasses.join(eqClass, jEqClass); |
| 110 | + } |
| 111 | + } |
| 112 | + bounds[boundedClasses.findLeader(i)] = bound; |
| 113 | + } |
| 114 | + |
| 115 | + // Return if there's no simplification. |
| 116 | + if (bounds.size() >= affineMap.getNumResults()) { |
| 117 | + LLVM_DEBUG( |
| 118 | + { DBGS() << "- the affine operation couldn't get simplified\n"; }); |
| 119 | + return false; |
| 120 | + } |
| 121 | + |
| 122 | + // Construct the new affine affineMap. |
| 123 | + SmallVector<AffineExpr> results; |
| 124 | + results.reserve(bounds.size()); |
| 125 | + for (auto [k, bound] : bounds) |
| 126 | + results.push_back(bound->getMap().getResult(0)); |
| 127 | + |
| 128 | + affineMap = AffineMap::get(affineMap.getNumDims(), affineMap.getNumSymbols(), |
| 129 | + results, rewriter.getContext()); |
| 130 | + |
| 131 | + // Update the affine op. |
| 132 | + rewriter.modifyOpInPlace(affineOp, [&]() { affineOp.setMap(affineMap); }); |
| 133 | + LLVM_DEBUG({ DBGS() << "- simplified affine op: `" << affineOp << "`\n"; }); |
| 134 | + return true; |
| 135 | +} |
| 136 | + |
| 137 | +bool mlir::affine::simplifyAffineMinOp(RewriterBase &rewriter, AffineMinOp op) { |
| 138 | + return simplifyAffineMinMaxOp(rewriter, op); |
| 139 | +} |
| 140 | + |
| 141 | +bool mlir::affine::simplifyAffineMaxOp(RewriterBase &rewriter, AffineMaxOp op) { |
| 142 | + return simplifyAffineMinMaxOp(rewriter, op); |
| 143 | +} |
| 144 | + |
| 145 | +LogicalResult mlir::affine::simplifyAffineMinMaxOps(RewriterBase &rewriter, |
| 146 | + ArrayRef<Operation *> ops, |
| 147 | + bool *modified) { |
| 148 | + bool changed = false; |
| 149 | + for (Operation *op : ops) { |
| 150 | + if ( auto minOp = dyn_cast<AffineMinOp>(op)) |
| 151 | + changed = simplifyAffineMinOp(rewriter, minOp) || changed; |
| 152 | + else if ( auto maxOp = cast<AffineMaxOp>(op)) |
| 153 | + changed = simplifyAffineMaxOp(rewriter, maxOp) || changed; |
| 154 | + } |
| 155 | + RewritePatternSet patterns(rewriter.getContext()); |
| 156 | + AffineMaxOp::getCanonicalizationPatterns(patterns, rewriter.getContext()); |
| 157 | + AffineMinOp::getCanonicalizationPatterns(patterns, rewriter.getContext()); |
| 158 | + FrozenRewritePatternSet frozenPatterns(std::move(patterns)); |
| 159 | + if (modified) |
| 160 | + *modified = changed; |
| 161 | + // Canonicalize to a fixpoint. |
| 162 | + if (failed(applyOpPatternsGreedily( |
| 163 | + ops, frozenPatterns, |
| 164 | + GreedyRewriteConfig() |
| 165 | + .setListener( |
| 166 | + static_cast<RewriterBase::Listener *>(rewriter.getListener())) |
| 167 | + .setStrictness(GreedyRewriteStrictness::ExistingAndNewOps), |
| 168 | + &changed))) { |
| 169 | + return failure(); |
| 170 | + } |
| 171 | + if (modified) |
| 172 | + *modified = changed; |
| 173 | + return success(); |
| 174 | +} |
0 commit comments