| File: | build/source/mlir/lib/IR/AffineExpr.cpp |
| Warning: | line 1072, column 7 Value stored to 'rhsPos' is never read |
Press '?' to see keyboard shortcuts
Keyboard shortcuts:
| 1 | //===- AffineExpr.cpp - MLIR Affine Expr Classes --------------------------===// |
| 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 | #include <utility> |
| 10 | |
| 11 | #include "AffineExprDetail.h" |
| 12 | #include "mlir/IR/AffineExpr.h" |
| 13 | #include "mlir/IR/AffineExprVisitor.h" |
| 14 | #include "mlir/IR/AffineMap.h" |
| 15 | #include "mlir/IR/IntegerSet.h" |
| 16 | #include "mlir/Support/MathExtras.h" |
| 17 | #include "mlir/Support/TypeID.h" |
| 18 | #include "llvm/ADT/STLExtras.h" |
| 19 | #include <numeric> |
| 20 | #include <optional> |
| 21 | |
| 22 | using namespace mlir; |
| 23 | using namespace mlir::detail; |
| 24 | |
| 25 | MLIRContext *AffineExpr::getContext() const { return expr->context; } |
| 26 | |
| 27 | AffineExprKind AffineExpr::getKind() const { return expr->kind; } |
| 28 | |
| 29 | /// Walk all of the AffineExprs in this subgraph in postorder. |
| 30 | void AffineExpr::walk(std::function<void(AffineExpr)> callback) const { |
| 31 | struct AffineExprWalker : public AffineExprVisitor<AffineExprWalker> { |
| 32 | std::function<void(AffineExpr)> callback; |
| 33 | |
| 34 | AffineExprWalker(std::function<void(AffineExpr)> callback) |
| 35 | : callback(std::move(callback)) {} |
| 36 | |
| 37 | void visitAffineBinaryOpExpr(AffineBinaryOpExpr expr) { callback(expr); } |
| 38 | void visitConstantExpr(AffineConstantExpr expr) { callback(expr); } |
| 39 | void visitDimExpr(AffineDimExpr expr) { callback(expr); } |
| 40 | void visitSymbolExpr(AffineSymbolExpr expr) { callback(expr); } |
| 41 | }; |
| 42 | |
| 43 | AffineExprWalker(std::move(callback)).walkPostOrder(*this); |
| 44 | } |
| 45 | |
| 46 | // Dispatch affine expression construction based on kind. |
| 47 | AffineExpr mlir::getAffineBinaryOpExpr(AffineExprKind kind, AffineExpr lhs, |
| 48 | AffineExpr rhs) { |
| 49 | if (kind == AffineExprKind::Add) |
| 50 | return lhs + rhs; |
| 51 | if (kind == AffineExprKind::Mul) |
| 52 | return lhs * rhs; |
| 53 | if (kind == AffineExprKind::FloorDiv) |
| 54 | return lhs.floorDiv(rhs); |
| 55 | if (kind == AffineExprKind::CeilDiv) |
| 56 | return lhs.ceilDiv(rhs); |
| 57 | if (kind == AffineExprKind::Mod) |
| 58 | return lhs % rhs; |
| 59 | |
| 60 | llvm_unreachable("unknown binary operation on affine expressions")::llvm::llvm_unreachable_internal("unknown binary operation on affine expressions" , "mlir/lib/IR/AffineExpr.cpp", 60); |
| 61 | } |
| 62 | |
| 63 | /// This method substitutes any uses of dimensions and symbols (e.g. |
| 64 | /// dim#0 with dimReplacements[0]) and returns the modified expression tree. |
| 65 | AffineExpr |
| 66 | AffineExpr::replaceDimsAndSymbols(ArrayRef<AffineExpr> dimReplacements, |
| 67 | ArrayRef<AffineExpr> symReplacements) const { |
| 68 | switch (getKind()) { |
| 69 | case AffineExprKind::Constant: |
| 70 | return *this; |
| 71 | case AffineExprKind::DimId: { |
| 72 | unsigned dimId = cast<AffineDimExpr>().getPosition(); |
| 73 | if (dimId >= dimReplacements.size()) |
| 74 | return *this; |
| 75 | return dimReplacements[dimId]; |
| 76 | } |
| 77 | case AffineExprKind::SymbolId: { |
| 78 | unsigned symId = cast<AffineSymbolExpr>().getPosition(); |
| 79 | if (symId >= symReplacements.size()) |
| 80 | return *this; |
| 81 | return symReplacements[symId]; |
| 82 | } |
| 83 | case AffineExprKind::Add: |
| 84 | case AffineExprKind::Mul: |
| 85 | case AffineExprKind::FloorDiv: |
| 86 | case AffineExprKind::CeilDiv: |
| 87 | case AffineExprKind::Mod: |
| 88 | auto binOp = cast<AffineBinaryOpExpr>(); |
| 89 | auto lhs = binOp.getLHS(), rhs = binOp.getRHS(); |
| 90 | auto newLHS = lhs.replaceDimsAndSymbols(dimReplacements, symReplacements); |
| 91 | auto newRHS = rhs.replaceDimsAndSymbols(dimReplacements, symReplacements); |
| 92 | if (newLHS == lhs && newRHS == rhs) |
| 93 | return *this; |
| 94 | return getAffineBinaryOpExpr(getKind(), newLHS, newRHS); |
| 95 | } |
| 96 | llvm_unreachable("Unknown AffineExpr")::llvm::llvm_unreachable_internal("Unknown AffineExpr", "mlir/lib/IR/AffineExpr.cpp" , 96); |
| 97 | } |
| 98 | |
| 99 | AffineExpr AffineExpr::replaceDims(ArrayRef<AffineExpr> dimReplacements) const { |
| 100 | return replaceDimsAndSymbols(dimReplacements, {}); |
| 101 | } |
| 102 | |
| 103 | AffineExpr |
| 104 | AffineExpr::replaceSymbols(ArrayRef<AffineExpr> symReplacements) const { |
| 105 | return replaceDimsAndSymbols({}, symReplacements); |
| 106 | } |
| 107 | |
| 108 | /// Replace dims[offset ... numDims) |
| 109 | /// by dims[offset + shift ... shift + numDims). |
| 110 | AffineExpr AffineExpr::shiftDims(unsigned numDims, unsigned shift, |
| 111 | unsigned offset) const { |
| 112 | SmallVector<AffineExpr, 4> dims; |
| 113 | for (unsigned idx = 0; idx < offset; ++idx) |
| 114 | dims.push_back(getAffineDimExpr(idx, getContext())); |
| 115 | for (unsigned idx = offset; idx < numDims; ++idx) |
| 116 | dims.push_back(getAffineDimExpr(idx + shift, getContext())); |
| 117 | return replaceDimsAndSymbols(dims, {}); |
| 118 | } |
| 119 | |
| 120 | /// Replace symbols[offset ... numSymbols) |
| 121 | /// by symbols[offset + shift ... shift + numSymbols). |
| 122 | AffineExpr AffineExpr::shiftSymbols(unsigned numSymbols, unsigned shift, |
| 123 | unsigned offset) const { |
| 124 | SmallVector<AffineExpr, 4> symbols; |
| 125 | for (unsigned idx = 0; idx < offset; ++idx) |
| 126 | symbols.push_back(getAffineSymbolExpr(idx, getContext())); |
| 127 | for (unsigned idx = offset; idx < numSymbols; ++idx) |
| 128 | symbols.push_back(getAffineSymbolExpr(idx + shift, getContext())); |
| 129 | return replaceDimsAndSymbols({}, symbols); |
| 130 | } |
| 131 | |
| 132 | /// Sparse replace method. Return the modified expression tree. |
| 133 | AffineExpr |
| 134 | AffineExpr::replace(const DenseMap<AffineExpr, AffineExpr> &map) const { |
| 135 | auto it = map.find(*this); |
| 136 | if (it != map.end()) |
| 137 | return it->second; |
| 138 | switch (getKind()) { |
| 139 | default: |
| 140 | return *this; |
| 141 | case AffineExprKind::Add: |
| 142 | case AffineExprKind::Mul: |
| 143 | case AffineExprKind::FloorDiv: |
| 144 | case AffineExprKind::CeilDiv: |
| 145 | case AffineExprKind::Mod: |
| 146 | auto binOp = cast<AffineBinaryOpExpr>(); |
| 147 | auto lhs = binOp.getLHS(), rhs = binOp.getRHS(); |
| 148 | auto newLHS = lhs.replace(map); |
| 149 | auto newRHS = rhs.replace(map); |
| 150 | if (newLHS == lhs && newRHS == rhs) |
| 151 | return *this; |
| 152 | return getAffineBinaryOpExpr(getKind(), newLHS, newRHS); |
| 153 | } |
| 154 | llvm_unreachable("Unknown AffineExpr")::llvm::llvm_unreachable_internal("Unknown AffineExpr", "mlir/lib/IR/AffineExpr.cpp" , 154); |
| 155 | } |
| 156 | |
| 157 | /// Sparse replace method. Return the modified expression tree. |
| 158 | AffineExpr AffineExpr::replace(AffineExpr expr, AffineExpr replacement) const { |
| 159 | DenseMap<AffineExpr, AffineExpr> map; |
| 160 | map.insert(std::make_pair(expr, replacement)); |
| 161 | return replace(map); |
| 162 | } |
| 163 | /// Returns true if this expression is made out of only symbols and |
| 164 | /// constants (no dimensional identifiers). |
| 165 | bool AffineExpr::isSymbolicOrConstant() const { |
| 166 | switch (getKind()) { |
| 167 | case AffineExprKind::Constant: |
| 168 | return true; |
| 169 | case AffineExprKind::DimId: |
| 170 | return false; |
| 171 | case AffineExprKind::SymbolId: |
| 172 | return true; |
| 173 | |
| 174 | case AffineExprKind::Add: |
| 175 | case AffineExprKind::Mul: |
| 176 | case AffineExprKind::FloorDiv: |
| 177 | case AffineExprKind::CeilDiv: |
| 178 | case AffineExprKind::Mod: { |
| 179 | auto expr = this->cast<AffineBinaryOpExpr>(); |
| 180 | return expr.getLHS().isSymbolicOrConstant() && |
| 181 | expr.getRHS().isSymbolicOrConstant(); |
| 182 | } |
| 183 | } |
| 184 | llvm_unreachable("Unknown AffineExpr")::llvm::llvm_unreachable_internal("Unknown AffineExpr", "mlir/lib/IR/AffineExpr.cpp" , 184); |
| 185 | } |
| 186 | |
| 187 | /// Returns true if this is a pure affine expression, i.e., multiplication, |
| 188 | /// floordiv, ceildiv, and mod is only allowed w.r.t constants. |
| 189 | bool AffineExpr::isPureAffine() const { |
| 190 | switch (getKind()) { |
| 191 | case AffineExprKind::SymbolId: |
| 192 | case AffineExprKind::DimId: |
| 193 | case AffineExprKind::Constant: |
| 194 | return true; |
| 195 | case AffineExprKind::Add: { |
| 196 | auto op = cast<AffineBinaryOpExpr>(); |
| 197 | return op.getLHS().isPureAffine() && op.getRHS().isPureAffine(); |
| 198 | } |
| 199 | |
| 200 | case AffineExprKind::Mul: { |
| 201 | // TODO: Canonicalize the constants in binary operators to the RHS when |
| 202 | // possible, allowing this to merge into the next case. |
| 203 | auto op = cast<AffineBinaryOpExpr>(); |
| 204 | return op.getLHS().isPureAffine() && op.getRHS().isPureAffine() && |
| 205 | (op.getLHS().template isa<AffineConstantExpr>() || |
| 206 | op.getRHS().template isa<AffineConstantExpr>()); |
| 207 | } |
| 208 | case AffineExprKind::FloorDiv: |
| 209 | case AffineExprKind::CeilDiv: |
| 210 | case AffineExprKind::Mod: { |
| 211 | auto op = cast<AffineBinaryOpExpr>(); |
| 212 | return op.getLHS().isPureAffine() && |
| 213 | op.getRHS().template isa<AffineConstantExpr>(); |
| 214 | } |
| 215 | } |
| 216 | llvm_unreachable("Unknown AffineExpr")::llvm::llvm_unreachable_internal("Unknown AffineExpr", "mlir/lib/IR/AffineExpr.cpp" , 216); |
| 217 | } |
| 218 | |
| 219 | // Returns the greatest known integral divisor of this affine expression. |
| 220 | int64_t AffineExpr::getLargestKnownDivisor() const { |
| 221 | AffineBinaryOpExpr binExpr(nullptr); |
| 222 | switch (getKind()) { |
| 223 | case AffineExprKind::DimId: |
| 224 | [[fallthrough]]; |
| 225 | case AffineExprKind::SymbolId: |
| 226 | return 1; |
| 227 | case AffineExprKind::CeilDiv: |
| 228 | [[fallthrough]]; |
| 229 | case AffineExprKind::FloorDiv: { |
| 230 | // If the RHS is a constant and divides the known divisor on the LHS, the |
| 231 | // quotient is a known divisor of the expression. |
| 232 | binExpr = this->cast<AffineBinaryOpExpr>(); |
| 233 | auto rhs = binExpr.getRHS().dyn_cast<AffineConstantExpr>(); |
| 234 | // Leave alone undefined expressions. |
| 235 | if (rhs && rhs.getValue() != 0) { |
| 236 | int64_t lhsDiv = binExpr.getLHS().getLargestKnownDivisor(); |
| 237 | if (lhsDiv % rhs.getValue() == 0) |
| 238 | return lhsDiv / rhs.getValue(); |
| 239 | } |
| 240 | return 1; |
| 241 | } |
| 242 | case AffineExprKind::Constant: |
| 243 | return std::abs(this->cast<AffineConstantExpr>().getValue()); |
| 244 | case AffineExprKind::Mul: { |
| 245 | binExpr = this->cast<AffineBinaryOpExpr>(); |
| 246 | return binExpr.getLHS().getLargestKnownDivisor() * |
| 247 | binExpr.getRHS().getLargestKnownDivisor(); |
| 248 | } |
| 249 | case AffineExprKind::Add: |
| 250 | [[fallthrough]]; |
| 251 | case AffineExprKind::Mod: { |
| 252 | binExpr = cast<AffineBinaryOpExpr>(); |
| 253 | return std::gcd((uint64_t)binExpr.getLHS().getLargestKnownDivisor(), |
| 254 | (uint64_t)binExpr.getRHS().getLargestKnownDivisor()); |
| 255 | } |
| 256 | } |
| 257 | llvm_unreachable("Unknown AffineExpr")::llvm::llvm_unreachable_internal("Unknown AffineExpr", "mlir/lib/IR/AffineExpr.cpp" , 257); |
| 258 | } |
| 259 | |
| 260 | bool AffineExpr::isMultipleOf(int64_t factor) const { |
| 261 | AffineBinaryOpExpr binExpr(nullptr); |
| 262 | uint64_t l, u; |
| 263 | switch (getKind()) { |
| 264 | case AffineExprKind::SymbolId: |
| 265 | [[fallthrough]]; |
| 266 | case AffineExprKind::DimId: |
| 267 | return factor * factor == 1; |
| 268 | case AffineExprKind::Constant: |
| 269 | return cast<AffineConstantExpr>().getValue() % factor == 0; |
| 270 | case AffineExprKind::Mul: { |
| 271 | binExpr = cast<AffineBinaryOpExpr>(); |
| 272 | // It's probably not worth optimizing this further (to not traverse the |
| 273 | // whole sub-tree under - it that would require a version of isMultipleOf |
| 274 | // that on a 'false' return also returns the largest known divisor). |
| 275 | return (l = binExpr.getLHS().getLargestKnownDivisor()) % factor == 0 || |
| 276 | (u = binExpr.getRHS().getLargestKnownDivisor()) % factor == 0 || |
| 277 | (l * u) % factor == 0; |
| 278 | } |
| 279 | case AffineExprKind::Add: |
| 280 | case AffineExprKind::FloorDiv: |
| 281 | case AffineExprKind::CeilDiv: |
| 282 | case AffineExprKind::Mod: { |
| 283 | binExpr = cast<AffineBinaryOpExpr>(); |
| 284 | return std::gcd((uint64_t)binExpr.getLHS().getLargestKnownDivisor(), |
| 285 | (uint64_t)binExpr.getRHS().getLargestKnownDivisor()) % |
| 286 | factor == |
| 287 | 0; |
| 288 | } |
| 289 | } |
| 290 | llvm_unreachable("Unknown AffineExpr")::llvm::llvm_unreachable_internal("Unknown AffineExpr", "mlir/lib/IR/AffineExpr.cpp" , 290); |
| 291 | } |
| 292 | |
| 293 | bool AffineExpr::isFunctionOfDim(unsigned position) const { |
| 294 | if (getKind() == AffineExprKind::DimId) { |
| 295 | return *this == mlir::getAffineDimExpr(position, getContext()); |
| 296 | } |
| 297 | if (auto expr = this->dyn_cast<AffineBinaryOpExpr>()) { |
| 298 | return expr.getLHS().isFunctionOfDim(position) || |
| 299 | expr.getRHS().isFunctionOfDim(position); |
| 300 | } |
| 301 | return false; |
| 302 | } |
| 303 | |
| 304 | bool AffineExpr::isFunctionOfSymbol(unsigned position) const { |
| 305 | if (getKind() == AffineExprKind::SymbolId) { |
| 306 | return *this == mlir::getAffineSymbolExpr(position, getContext()); |
| 307 | } |
| 308 | if (auto expr = this->dyn_cast<AffineBinaryOpExpr>()) { |
| 309 | return expr.getLHS().isFunctionOfSymbol(position) || |
| 310 | expr.getRHS().isFunctionOfSymbol(position); |
| 311 | } |
| 312 | return false; |
| 313 | } |
| 314 | |
| 315 | AffineBinaryOpExpr::AffineBinaryOpExpr(AffineExpr::ImplType *ptr) |
| 316 | : AffineExpr(ptr) {} |
| 317 | AffineExpr AffineBinaryOpExpr::getLHS() const { |
| 318 | return static_cast<ImplType *>(expr)->lhs; |
| 319 | } |
| 320 | AffineExpr AffineBinaryOpExpr::getRHS() const { |
| 321 | return static_cast<ImplType *>(expr)->rhs; |
| 322 | } |
| 323 | |
| 324 | AffineDimExpr::AffineDimExpr(AffineExpr::ImplType *ptr) : AffineExpr(ptr) {} |
| 325 | unsigned AffineDimExpr::getPosition() const { |
| 326 | return static_cast<ImplType *>(expr)->position; |
| 327 | } |
| 328 | |
| 329 | /// Returns true if the expression is divisible by the given symbol with |
| 330 | /// position `symbolPos`. The argument `opKind` specifies here what kind of |
| 331 | /// division or mod operation called this division. It helps in implementing the |
| 332 | /// commutative property of the floordiv and ceildiv operations. If the argument |
| 333 | ///`exprKind` is floordiv and `expr` is also a binary expression of a floordiv |
| 334 | /// operation, then the commutative property can be used otherwise, the floordiv |
| 335 | /// operation is not divisible. The same argument holds for ceildiv operation. |
| 336 | static bool isDivisibleBySymbol(AffineExpr expr, unsigned symbolPos, |
| 337 | AffineExprKind opKind) { |
| 338 | // The argument `opKind` can either be Modulo, Floordiv or Ceildiv only. |
| 339 | assert((opKind == AffineExprKind::Mod || opKind == AffineExprKind::FloorDiv ||(static_cast <bool> ((opKind == AffineExprKind::Mod || opKind == AffineExprKind::FloorDiv || opKind == AffineExprKind::CeilDiv ) && "unexpected opKind") ? void (0) : __assert_fail ( "(opKind == AffineExprKind::Mod || opKind == AffineExprKind::FloorDiv || opKind == AffineExprKind::CeilDiv) && \"unexpected opKind\"" , "mlir/lib/IR/AffineExpr.cpp", 341, __extension__ __PRETTY_FUNCTION__ )) |
| 340 | opKind == AffineExprKind::CeilDiv) &&(static_cast <bool> ((opKind == AffineExprKind::Mod || opKind == AffineExprKind::FloorDiv || opKind == AffineExprKind::CeilDiv ) && "unexpected opKind") ? void (0) : __assert_fail ( "(opKind == AffineExprKind::Mod || opKind == AffineExprKind::FloorDiv || opKind == AffineExprKind::CeilDiv) && \"unexpected opKind\"" , "mlir/lib/IR/AffineExpr.cpp", 341, __extension__ __PRETTY_FUNCTION__ )) |
| 341 | "unexpected opKind")(static_cast <bool> ((opKind == AffineExprKind::Mod || opKind == AffineExprKind::FloorDiv || opKind == AffineExprKind::CeilDiv ) && "unexpected opKind") ? void (0) : __assert_fail ( "(opKind == AffineExprKind::Mod || opKind == AffineExprKind::FloorDiv || opKind == AffineExprKind::CeilDiv) && \"unexpected opKind\"" , "mlir/lib/IR/AffineExpr.cpp", 341, __extension__ __PRETTY_FUNCTION__ )); |
| 342 | switch (expr.getKind()) { |
| 343 | case AffineExprKind::Constant: |
| 344 | return expr.cast<AffineConstantExpr>().getValue() == 0; |
| 345 | case AffineExprKind::DimId: |
| 346 | return false; |
| 347 | case AffineExprKind::SymbolId: |
| 348 | return (expr.cast<AffineSymbolExpr>().getPosition() == symbolPos); |
| 349 | // Checks divisibility by the given symbol for both operands. |
| 350 | case AffineExprKind::Add: { |
| 351 | AffineBinaryOpExpr binaryExpr = expr.cast<AffineBinaryOpExpr>(); |
| 352 | return isDivisibleBySymbol(binaryExpr.getLHS(), symbolPos, opKind) && |
| 353 | isDivisibleBySymbol(binaryExpr.getRHS(), symbolPos, opKind); |
| 354 | } |
| 355 | // Checks divisibility by the given symbol for both operands. Consider the |
| 356 | // expression `(((s1*s0) floordiv w) mod ((s1 * s2) floordiv p)) floordiv s1`, |
| 357 | // this is a division by s1 and both the operands of modulo are divisible by |
| 358 | // s1 but it is not divisible by s1 always. The third argument is |
| 359 | // `AffineExprKind::Mod` for this reason. |
| 360 | case AffineExprKind::Mod: { |
| 361 | AffineBinaryOpExpr binaryExpr = expr.cast<AffineBinaryOpExpr>(); |
| 362 | return isDivisibleBySymbol(binaryExpr.getLHS(), symbolPos, |
| 363 | AffineExprKind::Mod) && |
| 364 | isDivisibleBySymbol(binaryExpr.getRHS(), symbolPos, |
| 365 | AffineExprKind::Mod); |
| 366 | } |
| 367 | // Checks if any of the operand divisible by the given symbol. |
| 368 | case AffineExprKind::Mul: { |
| 369 | AffineBinaryOpExpr binaryExpr = expr.cast<AffineBinaryOpExpr>(); |
| 370 | return isDivisibleBySymbol(binaryExpr.getLHS(), symbolPos, opKind) || |
| 371 | isDivisibleBySymbol(binaryExpr.getRHS(), symbolPos, opKind); |
| 372 | } |
| 373 | // Floordiv and ceildiv are divisible by the given symbol when the first |
| 374 | // operand is divisible, and the affine expression kind of the argument expr |
| 375 | // is same as the argument `opKind`. This can be inferred from commutative |
| 376 | // property of floordiv and ceildiv operations and are as follow: |
| 377 | // (exp1 floordiv exp2) floordiv exp3 = (exp1 floordiv exp3) floordiv exp2 |
| 378 | // (exp1 ceildiv exp2) ceildiv exp3 = (exp1 ceildiv exp3) ceildiv expr2 |
| 379 | // It will fail if operations are not same. For example: |
| 380 | // (exps1 ceildiv exp2) floordiv exp3 can not be simplified. |
| 381 | case AffineExprKind::FloorDiv: |
| 382 | case AffineExprKind::CeilDiv: { |
| 383 | AffineBinaryOpExpr binaryExpr = expr.cast<AffineBinaryOpExpr>(); |
| 384 | if (opKind != expr.getKind()) |
| 385 | return false; |
| 386 | return isDivisibleBySymbol(binaryExpr.getLHS(), symbolPos, expr.getKind()); |
| 387 | } |
| 388 | } |
| 389 | llvm_unreachable("Unknown AffineExpr")::llvm::llvm_unreachable_internal("Unknown AffineExpr", "mlir/lib/IR/AffineExpr.cpp" , 389); |
| 390 | } |
| 391 | |
| 392 | /// Divides the given expression by the given symbol at position `symbolPos`. It |
| 393 | /// considers the divisibility condition is checked before calling itself. A |
| 394 | /// null expression is returned whenever the divisibility condition fails. |
| 395 | static AffineExpr symbolicDivide(AffineExpr expr, unsigned symbolPos, |
| 396 | AffineExprKind opKind) { |
| 397 | // THe argument `opKind` can either be Modulo, Floordiv or Ceildiv only. |
| 398 | assert((opKind == AffineExprKind::Mod || opKind == AffineExprKind::FloorDiv ||(static_cast <bool> ((opKind == AffineExprKind::Mod || opKind == AffineExprKind::FloorDiv || opKind == AffineExprKind::CeilDiv ) && "unexpected opKind") ? void (0) : __assert_fail ( "(opKind == AffineExprKind::Mod || opKind == AffineExprKind::FloorDiv || opKind == AffineExprKind::CeilDiv) && \"unexpected opKind\"" , "mlir/lib/IR/AffineExpr.cpp", 400, __extension__ __PRETTY_FUNCTION__ )) |
| 399 | opKind == AffineExprKind::CeilDiv) &&(static_cast <bool> ((opKind == AffineExprKind::Mod || opKind == AffineExprKind::FloorDiv || opKind == AffineExprKind::CeilDiv ) && "unexpected opKind") ? void (0) : __assert_fail ( "(opKind == AffineExprKind::Mod || opKind == AffineExprKind::FloorDiv || opKind == AffineExprKind::CeilDiv) && \"unexpected opKind\"" , "mlir/lib/IR/AffineExpr.cpp", 400, __extension__ __PRETTY_FUNCTION__ )) |
| 400 | "unexpected opKind")(static_cast <bool> ((opKind == AffineExprKind::Mod || opKind == AffineExprKind::FloorDiv || opKind == AffineExprKind::CeilDiv ) && "unexpected opKind") ? void (0) : __assert_fail ( "(opKind == AffineExprKind::Mod || opKind == AffineExprKind::FloorDiv || opKind == AffineExprKind::CeilDiv) && \"unexpected opKind\"" , "mlir/lib/IR/AffineExpr.cpp", 400, __extension__ __PRETTY_FUNCTION__ )); |
| 401 | switch (expr.getKind()) { |
| 402 | case AffineExprKind::Constant: |
| 403 | if (expr.cast<AffineConstantExpr>().getValue() != 0) |
| 404 | return nullptr; |
| 405 | return getAffineConstantExpr(0, expr.getContext()); |
| 406 | case AffineExprKind::DimId: |
| 407 | return nullptr; |
| 408 | case AffineExprKind::SymbolId: |
| 409 | return getAffineConstantExpr(1, expr.getContext()); |
| 410 | // Dividing both operands by the given symbol. |
| 411 | case AffineExprKind::Add: { |
| 412 | AffineBinaryOpExpr binaryExpr = expr.cast<AffineBinaryOpExpr>(); |
| 413 | return getAffineBinaryOpExpr( |
| 414 | expr.getKind(), symbolicDivide(binaryExpr.getLHS(), symbolPos, opKind), |
| 415 | symbolicDivide(binaryExpr.getRHS(), symbolPos, opKind)); |
| 416 | } |
| 417 | // Dividing both operands by the given symbol. |
| 418 | case AffineExprKind::Mod: { |
| 419 | AffineBinaryOpExpr binaryExpr = expr.cast<AffineBinaryOpExpr>(); |
| 420 | return getAffineBinaryOpExpr( |
| 421 | expr.getKind(), |
| 422 | symbolicDivide(binaryExpr.getLHS(), symbolPos, expr.getKind()), |
| 423 | symbolicDivide(binaryExpr.getRHS(), symbolPos, expr.getKind())); |
| 424 | } |
| 425 | // Dividing any of the operand by the given symbol. |
| 426 | case AffineExprKind::Mul: { |
| 427 | AffineBinaryOpExpr binaryExpr = expr.cast<AffineBinaryOpExpr>(); |
| 428 | if (!isDivisibleBySymbol(binaryExpr.getLHS(), symbolPos, opKind)) |
| 429 | return binaryExpr.getLHS() * |
| 430 | symbolicDivide(binaryExpr.getRHS(), symbolPos, opKind); |
| 431 | return symbolicDivide(binaryExpr.getLHS(), symbolPos, opKind) * |
| 432 | binaryExpr.getRHS(); |
| 433 | } |
| 434 | // Dividing first operand only by the given symbol. |
| 435 | case AffineExprKind::FloorDiv: |
| 436 | case AffineExprKind::CeilDiv: { |
| 437 | AffineBinaryOpExpr binaryExpr = expr.cast<AffineBinaryOpExpr>(); |
| 438 | return getAffineBinaryOpExpr( |
| 439 | expr.getKind(), |
| 440 | symbolicDivide(binaryExpr.getLHS(), symbolPos, expr.getKind()), |
| 441 | binaryExpr.getRHS()); |
| 442 | } |
| 443 | } |
| 444 | llvm_unreachable("Unknown AffineExpr")::llvm::llvm_unreachable_internal("Unknown AffineExpr", "mlir/lib/IR/AffineExpr.cpp" , 444); |
| 445 | } |
| 446 | |
| 447 | /// Simplify a semi-affine expression by handling modulo, floordiv, or ceildiv |
| 448 | /// operations when the second operand simplifies to a symbol and the first |
| 449 | /// operand is divisible by that symbol. It can be applied to any semi-affine |
| 450 | /// expression. Returned expression can either be a semi-affine or pure affine |
| 451 | /// expression. |
| 452 | static AffineExpr simplifySemiAffine(AffineExpr expr) { |
| 453 | switch (expr.getKind()) { |
| 454 | case AffineExprKind::Constant: |
| 455 | case AffineExprKind::DimId: |
| 456 | case AffineExprKind::SymbolId: |
| 457 | return expr; |
| 458 | case AffineExprKind::Add: |
| 459 | case AffineExprKind::Mul: { |
| 460 | AffineBinaryOpExpr binaryExpr = expr.cast<AffineBinaryOpExpr>(); |
| 461 | return getAffineBinaryOpExpr(expr.getKind(), |
| 462 | simplifySemiAffine(binaryExpr.getLHS()), |
| 463 | simplifySemiAffine(binaryExpr.getRHS())); |
| 464 | } |
| 465 | // Check if the simplification of the second operand is a symbol, and the |
| 466 | // first operand is divisible by it. If the operation is a modulo, a constant |
| 467 | // zero expression is returned. In the case of floordiv and ceildiv, the |
| 468 | // symbol from the simplification of the second operand divides the first |
| 469 | // operand. Otherwise, simplification is not possible. |
| 470 | case AffineExprKind::FloorDiv: |
| 471 | case AffineExprKind::CeilDiv: |
| 472 | case AffineExprKind::Mod: { |
| 473 | AffineBinaryOpExpr binaryExpr = expr.cast<AffineBinaryOpExpr>(); |
| 474 | AffineExpr sLHS = simplifySemiAffine(binaryExpr.getLHS()); |
| 475 | AffineExpr sRHS = simplifySemiAffine(binaryExpr.getRHS()); |
| 476 | AffineSymbolExpr symbolExpr = |
| 477 | simplifySemiAffine(binaryExpr.getRHS()).dyn_cast<AffineSymbolExpr>(); |
| 478 | if (!symbolExpr) |
| 479 | return getAffineBinaryOpExpr(expr.getKind(), sLHS, sRHS); |
| 480 | unsigned symbolPos = symbolExpr.getPosition(); |
| 481 | if (!isDivisibleBySymbol(binaryExpr.getLHS(), symbolPos, expr.getKind())) |
| 482 | return getAffineBinaryOpExpr(expr.getKind(), sLHS, sRHS); |
| 483 | if (expr.getKind() == AffineExprKind::Mod) |
| 484 | return getAffineConstantExpr(0, expr.getContext()); |
| 485 | return symbolicDivide(sLHS, symbolPos, expr.getKind()); |
| 486 | } |
| 487 | } |
| 488 | llvm_unreachable("Unknown AffineExpr")::llvm::llvm_unreachable_internal("Unknown AffineExpr", "mlir/lib/IR/AffineExpr.cpp" , 488); |
| 489 | } |
| 490 | |
| 491 | static AffineExpr getAffineDimOrSymbol(AffineExprKind kind, unsigned position, |
| 492 | MLIRContext *context) { |
| 493 | auto assignCtx = [context](AffineDimExprStorage *storage) { |
| 494 | storage->context = context; |
| 495 | }; |
| 496 | |
| 497 | StorageUniquer &uniquer = context->getAffineUniquer(); |
| 498 | return uniquer.get<AffineDimExprStorage>( |
| 499 | assignCtx, static_cast<unsigned>(kind), position); |
| 500 | } |
| 501 | |
| 502 | AffineExpr mlir::getAffineDimExpr(unsigned position, MLIRContext *context) { |
| 503 | return getAffineDimOrSymbol(AffineExprKind::DimId, position, context); |
| 504 | } |
| 505 | |
| 506 | AffineSymbolExpr::AffineSymbolExpr(AffineExpr::ImplType *ptr) |
| 507 | : AffineExpr(ptr) {} |
| 508 | unsigned AffineSymbolExpr::getPosition() const { |
| 509 | return static_cast<ImplType *>(expr)->position; |
| 510 | } |
| 511 | |
| 512 | AffineExpr mlir::getAffineSymbolExpr(unsigned position, MLIRContext *context) { |
| 513 | return getAffineDimOrSymbol(AffineExprKind::SymbolId, position, context); |
| 514 | ; |
| 515 | } |
| 516 | |
| 517 | AffineConstantExpr::AffineConstantExpr(AffineExpr::ImplType *ptr) |
| 518 | : AffineExpr(ptr) {} |
| 519 | int64_t AffineConstantExpr::getValue() const { |
| 520 | return static_cast<ImplType *>(expr)->constant; |
| 521 | } |
| 522 | |
| 523 | bool AffineExpr::operator==(int64_t v) const { |
| 524 | return *this == getAffineConstantExpr(v, getContext()); |
| 525 | } |
| 526 | |
| 527 | AffineExpr mlir::getAffineConstantExpr(int64_t constant, MLIRContext *context) { |
| 528 | auto assignCtx = [context](AffineConstantExprStorage *storage) { |
| 529 | storage->context = context; |
| 530 | }; |
| 531 | |
| 532 | StorageUniquer &uniquer = context->getAffineUniquer(); |
| 533 | return uniquer.get<AffineConstantExprStorage>(assignCtx, constant); |
| 534 | } |
| 535 | |
| 536 | /// Simplify add expression. Return nullptr if it can't be simplified. |
| 537 | static AffineExpr simplifyAdd(AffineExpr lhs, AffineExpr rhs) { |
| 538 | auto lhsConst = lhs.dyn_cast<AffineConstantExpr>(); |
| 539 | auto rhsConst = rhs.dyn_cast<AffineConstantExpr>(); |
| 540 | // Fold if both LHS, RHS are a constant. |
| 541 | if (lhsConst && rhsConst) |
| 542 | return getAffineConstantExpr(lhsConst.getValue() + rhsConst.getValue(), |
| 543 | lhs.getContext()); |
| 544 | |
| 545 | // Canonicalize so that only the RHS is a constant. (4 + d0 becomes d0 + 4). |
| 546 | // If only one of them is a symbolic expressions, make it the RHS. |
| 547 | if (lhs.isa<AffineConstantExpr>() || |
| 548 | (lhs.isSymbolicOrConstant() && !rhs.isSymbolicOrConstant())) { |
| 549 | return rhs + lhs; |
| 550 | } |
| 551 | |
| 552 | // At this point, if there was a constant, it would be on the right. |
| 553 | |
| 554 | // Addition with a zero is a noop, return the other input. |
| 555 | if (rhsConst) { |
| 556 | if (rhsConst.getValue() == 0) |
| 557 | return lhs; |
| 558 | } |
| 559 | // Fold successive additions like (d0 + 2) + 3 into d0 + 5. |
| 560 | auto lBin = lhs.dyn_cast<AffineBinaryOpExpr>(); |
| 561 | if (lBin && rhsConst && lBin.getKind() == AffineExprKind::Add) { |
| 562 | if (auto lrhs = lBin.getRHS().dyn_cast<AffineConstantExpr>()) |
| 563 | return lBin.getLHS() + (lrhs.getValue() + rhsConst.getValue()); |
| 564 | } |
| 565 | |
| 566 | // Detect "c1 * expr + c_2 * expr" as "(c1 + c2) * expr". |
| 567 | // c1 is rRhsConst, c2 is rLhsConst; firstExpr, secondExpr are their |
| 568 | // respective multiplicands. |
| 569 | std::optional<int64_t> rLhsConst, rRhsConst; |
| 570 | AffineExpr firstExpr, secondExpr; |
| 571 | AffineConstantExpr rLhsConstExpr; |
| 572 | auto lBinOpExpr = lhs.dyn_cast<AffineBinaryOpExpr>(); |
| 573 | if (lBinOpExpr && lBinOpExpr.getKind() == AffineExprKind::Mul && |
| 574 | (rLhsConstExpr = lBinOpExpr.getRHS().dyn_cast<AffineConstantExpr>())) { |
| 575 | rLhsConst = rLhsConstExpr.getValue(); |
| 576 | firstExpr = lBinOpExpr.getLHS(); |
| 577 | } else { |
| 578 | rLhsConst = 1; |
| 579 | firstExpr = lhs; |
| 580 | } |
| 581 | |
| 582 | auto rBinOpExpr = rhs.dyn_cast<AffineBinaryOpExpr>(); |
| 583 | AffineConstantExpr rRhsConstExpr; |
| 584 | if (rBinOpExpr && rBinOpExpr.getKind() == AffineExprKind::Mul && |
| 585 | (rRhsConstExpr = rBinOpExpr.getRHS().dyn_cast<AffineConstantExpr>())) { |
| 586 | rRhsConst = rRhsConstExpr.getValue(); |
| 587 | secondExpr = rBinOpExpr.getLHS(); |
| 588 | } else { |
| 589 | rRhsConst = 1; |
| 590 | secondExpr = rhs; |
| 591 | } |
| 592 | |
| 593 | if (rLhsConst && rRhsConst && firstExpr == secondExpr) |
| 594 | return getAffineBinaryOpExpr( |
| 595 | AffineExprKind::Mul, firstExpr, |
| 596 | getAffineConstantExpr(*rLhsConst + *rRhsConst, lhs.getContext())); |
| 597 | |
| 598 | // When doing successive additions, bring constant to the right: turn (d0 + 2) |
| 599 | // + d1 into (d0 + d1) + 2. |
| 600 | if (lBin && lBin.getKind() == AffineExprKind::Add) { |
| 601 | if (auto lrhs = lBin.getRHS().dyn_cast<AffineConstantExpr>()) { |
| 602 | return lBin.getLHS() + rhs + lrhs; |
| 603 | } |
| 604 | } |
| 605 | |
| 606 | // Detect and transform "expr - q * (expr floordiv q)" to "expr mod q", where |
| 607 | // q may be a constant or symbolic expression. This leads to a much more |
| 608 | // efficient form when 'c' is a power of two, and in general a more compact |
| 609 | // and readable form. |
| 610 | |
| 611 | // Process '(expr floordiv c) * (-c)'. |
| 612 | if (!rBinOpExpr) |
| 613 | return nullptr; |
| 614 | |
| 615 | auto lrhs = rBinOpExpr.getLHS(); |
| 616 | auto rrhs = rBinOpExpr.getRHS(); |
| 617 | |
| 618 | AffineExpr llrhs, rlrhs; |
| 619 | |
| 620 | // Check if lrhsBinOpExpr is of the form (expr floordiv q) * q, where q is a |
| 621 | // symbolic expression. |
| 622 | auto lrhsBinOpExpr = lrhs.dyn_cast<AffineBinaryOpExpr>(); |
| 623 | // Check rrhsConstOpExpr = -1. |
| 624 | auto rrhsConstOpExpr = rrhs.dyn_cast<AffineConstantExpr>(); |
| 625 | if (rrhsConstOpExpr && rrhsConstOpExpr.getValue() == -1 && lrhsBinOpExpr && |
| 626 | lrhsBinOpExpr.getKind() == AffineExprKind::Mul) { |
| 627 | // Check llrhs = expr floordiv q. |
| 628 | llrhs = lrhsBinOpExpr.getLHS(); |
| 629 | // Check rlrhs = q. |
| 630 | rlrhs = lrhsBinOpExpr.getRHS(); |
| 631 | auto llrhsBinOpExpr = llrhs.dyn_cast<AffineBinaryOpExpr>(); |
| 632 | if (!llrhsBinOpExpr || llrhsBinOpExpr.getKind() != AffineExprKind::FloorDiv) |
| 633 | return nullptr; |
| 634 | if (llrhsBinOpExpr.getRHS() == rlrhs && lhs == llrhsBinOpExpr.getLHS()) |
| 635 | return lhs % rlrhs; |
| 636 | } |
| 637 | |
| 638 | // Process lrhs, which is 'expr floordiv c'. |
| 639 | AffineBinaryOpExpr lrBinOpExpr = lrhs.dyn_cast<AffineBinaryOpExpr>(); |
| 640 | if (!lrBinOpExpr || lrBinOpExpr.getKind() != AffineExprKind::FloorDiv) |
| 641 | return nullptr; |
| 642 | |
| 643 | llrhs = lrBinOpExpr.getLHS(); |
| 644 | rlrhs = lrBinOpExpr.getRHS(); |
| 645 | |
| 646 | if (lhs == llrhs && rlrhs == -rrhs) { |
| 647 | return lhs % rlrhs; |
| 648 | } |
| 649 | return nullptr; |
| 650 | } |
| 651 | |
| 652 | AffineExpr AffineExpr::operator+(int64_t v) const { |
| 653 | return *this + getAffineConstantExpr(v, getContext()); |
| 654 | } |
| 655 | AffineExpr AffineExpr::operator+(AffineExpr other) const { |
| 656 | if (auto simplified = simplifyAdd(*this, other)) |
| 657 | return simplified; |
| 658 | |
| 659 | StorageUniquer &uniquer = getContext()->getAffineUniquer(); |
| 660 | return uniquer.get<AffineBinaryOpExprStorage>( |
| 661 | /*initFn=*/{}, static_cast<unsigned>(AffineExprKind::Add), *this, other); |
| 662 | } |
| 663 | |
| 664 | /// Simplify a multiply expression. Return nullptr if it can't be simplified. |
| 665 | static AffineExpr simplifyMul(AffineExpr lhs, AffineExpr rhs) { |
| 666 | auto lhsConst = lhs.dyn_cast<AffineConstantExpr>(); |
| 667 | auto rhsConst = rhs.dyn_cast<AffineConstantExpr>(); |
| 668 | |
| 669 | if (lhsConst && rhsConst) |
| 670 | return getAffineConstantExpr(lhsConst.getValue() * rhsConst.getValue(), |
| 671 | lhs.getContext()); |
| 672 | |
| 673 | assert(lhs.isSymbolicOrConstant() || rhs.isSymbolicOrConstant())(static_cast <bool> (lhs.isSymbolicOrConstant() || rhs. isSymbolicOrConstant()) ? void (0) : __assert_fail ("lhs.isSymbolicOrConstant() || rhs.isSymbolicOrConstant()" , "mlir/lib/IR/AffineExpr.cpp", 673, __extension__ __PRETTY_FUNCTION__ )); |
| 674 | |
| 675 | // Canonicalize the mul expression so that the constant/symbolic term is the |
| 676 | // RHS. If both the lhs and rhs are symbolic, swap them if the lhs is a |
| 677 | // constant. (Note that a constant is trivially symbolic). |
| 678 | if (!rhs.isSymbolicOrConstant() || lhs.isa<AffineConstantExpr>()) { |
| 679 | // At least one of them has to be symbolic. |
| 680 | return rhs * lhs; |
| 681 | } |
| 682 | |
| 683 | // At this point, if there was a constant, it would be on the right. |
| 684 | |
| 685 | // Multiplication with a one is a noop, return the other input. |
| 686 | if (rhsConst) { |
| 687 | if (rhsConst.getValue() == 1) |
| 688 | return lhs; |
| 689 | // Multiplication with zero. |
| 690 | if (rhsConst.getValue() == 0) |
| 691 | return rhsConst; |
| 692 | } |
| 693 | |
| 694 | // Fold successive multiplications: eg: (d0 * 2) * 3 into d0 * 6. |
| 695 | auto lBin = lhs.dyn_cast<AffineBinaryOpExpr>(); |
| 696 | if (lBin && rhsConst && lBin.getKind() == AffineExprKind::Mul) { |
| 697 | if (auto lrhs = lBin.getRHS().dyn_cast<AffineConstantExpr>()) |
| 698 | return lBin.getLHS() * (lrhs.getValue() * rhsConst.getValue()); |
| 699 | } |
| 700 | |
| 701 | // When doing successive multiplication, bring constant to the right: turn (d0 |
| 702 | // * 2) * d1 into (d0 * d1) * 2. |
| 703 | if (lBin && lBin.getKind() == AffineExprKind::Mul) { |
| 704 | if (auto lrhs = lBin.getRHS().dyn_cast<AffineConstantExpr>()) { |
| 705 | return (lBin.getLHS() * rhs) * lrhs; |
| 706 | } |
| 707 | } |
| 708 | |
| 709 | return nullptr; |
| 710 | } |
| 711 | |
| 712 | AffineExpr AffineExpr::operator*(int64_t v) const { |
| 713 | return *this * getAffineConstantExpr(v, getContext()); |
| 714 | } |
| 715 | AffineExpr AffineExpr::operator*(AffineExpr other) const { |
| 716 | if (auto simplified = simplifyMul(*this, other)) |
| 717 | return simplified; |
| 718 | |
| 719 | StorageUniquer &uniquer = getContext()->getAffineUniquer(); |
| 720 | return uniquer.get<AffineBinaryOpExprStorage>( |
| 721 | /*initFn=*/{}, static_cast<unsigned>(AffineExprKind::Mul), *this, other); |
| 722 | } |
| 723 | |
| 724 | // Unary minus, delegate to operator*. |
| 725 | AffineExpr AffineExpr::operator-() const { |
| 726 | return *this * getAffineConstantExpr(-1, getContext()); |
| 727 | } |
| 728 | |
| 729 | // Delegate to operator+. |
| 730 | AffineExpr AffineExpr::operator-(int64_t v) const { return *this + (-v); } |
| 731 | AffineExpr AffineExpr::operator-(AffineExpr other) const { |
| 732 | return *this + (-other); |
| 733 | } |
| 734 | |
| 735 | static AffineExpr simplifyFloorDiv(AffineExpr lhs, AffineExpr rhs) { |
| 736 | auto lhsConst = lhs.dyn_cast<AffineConstantExpr>(); |
| 737 | auto rhsConst = rhs.dyn_cast<AffineConstantExpr>(); |
| 738 | |
| 739 | // mlir floordiv by zero or negative numbers is undefined and preserved as is. |
| 740 | if (!rhsConst || rhsConst.getValue() < 1) |
| 741 | return nullptr; |
| 742 | |
| 743 | if (lhsConst) |
| 744 | return getAffineConstantExpr( |
| 745 | floorDiv(lhsConst.getValue(), rhsConst.getValue()), lhs.getContext()); |
| 746 | |
| 747 | // Fold floordiv of a multiply with a constant that is a multiple of the |
| 748 | // divisor. Eg: (i * 128) floordiv 64 = i * 2. |
| 749 | if (rhsConst == 1) |
| 750 | return lhs; |
| 751 | |
| 752 | // Simplify (expr * const) floordiv divConst when expr is known to be a |
| 753 | // multiple of divConst. |
| 754 | auto lBin = lhs.dyn_cast<AffineBinaryOpExpr>(); |
| 755 | if (lBin && lBin.getKind() == AffineExprKind::Mul) { |
| 756 | if (auto lrhs = lBin.getRHS().dyn_cast<AffineConstantExpr>()) { |
| 757 | // rhsConst is known to be a positive constant. |
| 758 | if (lrhs.getValue() % rhsConst.getValue() == 0) |
| 759 | return lBin.getLHS() * (lrhs.getValue() / rhsConst.getValue()); |
| 760 | } |
| 761 | } |
| 762 | |
| 763 | // Simplify (expr1 + expr2) floordiv divConst when either expr1 or expr2 is |
| 764 | // known to be a multiple of divConst. |
| 765 | if (lBin && lBin.getKind() == AffineExprKind::Add) { |
| 766 | int64_t llhsDiv = lBin.getLHS().getLargestKnownDivisor(); |
| 767 | int64_t lrhsDiv = lBin.getRHS().getLargestKnownDivisor(); |
| 768 | // rhsConst is known to be a positive constant. |
| 769 | if (llhsDiv % rhsConst.getValue() == 0 || |
| 770 | lrhsDiv % rhsConst.getValue() == 0) |
| 771 | return lBin.getLHS().floorDiv(rhsConst.getValue()) + |
| 772 | lBin.getRHS().floorDiv(rhsConst.getValue()); |
| 773 | } |
| 774 | |
| 775 | return nullptr; |
| 776 | } |
| 777 | |
| 778 | AffineExpr AffineExpr::floorDiv(uint64_t v) const { |
| 779 | return floorDiv(getAffineConstantExpr(v, getContext())); |
| 780 | } |
| 781 | AffineExpr AffineExpr::floorDiv(AffineExpr other) const { |
| 782 | if (auto simplified = simplifyFloorDiv(*this, other)) |
| 783 | return simplified; |
| 784 | |
| 785 | StorageUniquer &uniquer = getContext()->getAffineUniquer(); |
| 786 | return uniquer.get<AffineBinaryOpExprStorage>( |
| 787 | /*initFn=*/{}, static_cast<unsigned>(AffineExprKind::FloorDiv), *this, |
| 788 | other); |
| 789 | } |
| 790 | |
| 791 | static AffineExpr simplifyCeilDiv(AffineExpr lhs, AffineExpr rhs) { |
| 792 | auto lhsConst = lhs.dyn_cast<AffineConstantExpr>(); |
| 793 | auto rhsConst = rhs.dyn_cast<AffineConstantExpr>(); |
| 794 | |
| 795 | if (!rhsConst || rhsConst.getValue() < 1) |
| 796 | return nullptr; |
| 797 | |
| 798 | if (lhsConst) |
| 799 | return getAffineConstantExpr( |
| 800 | ceilDiv(lhsConst.getValue(), rhsConst.getValue()), lhs.getContext()); |
| 801 | |
| 802 | // Fold ceildiv of a multiply with a constant that is a multiple of the |
| 803 | // divisor. Eg: (i * 128) ceildiv 64 = i * 2. |
| 804 | if (rhsConst.getValue() == 1) |
| 805 | return lhs; |
| 806 | |
| 807 | // Simplify (expr * const) ceildiv divConst when const is known to be a |
| 808 | // multiple of divConst. |
| 809 | auto lBin = lhs.dyn_cast<AffineBinaryOpExpr>(); |
| 810 | if (lBin && lBin.getKind() == AffineExprKind::Mul) { |
| 811 | if (auto lrhs = lBin.getRHS().dyn_cast<AffineConstantExpr>()) { |
| 812 | // rhsConst is known to be a positive constant. |
| 813 | if (lrhs.getValue() % rhsConst.getValue() == 0) |
| 814 | return lBin.getLHS() * (lrhs.getValue() / rhsConst.getValue()); |
| 815 | } |
| 816 | } |
| 817 | |
| 818 | return nullptr; |
| 819 | } |
| 820 | |
| 821 | AffineExpr AffineExpr::ceilDiv(uint64_t v) const { |
| 822 | return ceilDiv(getAffineConstantExpr(v, getContext())); |
| 823 | } |
| 824 | AffineExpr AffineExpr::ceilDiv(AffineExpr other) const { |
| 825 | if (auto simplified = simplifyCeilDiv(*this, other)) |
| 826 | return simplified; |
| 827 | |
| 828 | StorageUniquer &uniquer = getContext()->getAffineUniquer(); |
| 829 | return uniquer.get<AffineBinaryOpExprStorage>( |
| 830 | /*initFn=*/{}, static_cast<unsigned>(AffineExprKind::CeilDiv), *this, |
| 831 | other); |
| 832 | } |
| 833 | |
| 834 | static AffineExpr simplifyMod(AffineExpr lhs, AffineExpr rhs) { |
| 835 | auto lhsConst = lhs.dyn_cast<AffineConstantExpr>(); |
| 836 | auto rhsConst = rhs.dyn_cast<AffineConstantExpr>(); |
| 837 | |
| 838 | // mod w.r.t zero or negative numbers is undefined and preserved as is. |
| 839 | if (!rhsConst || rhsConst.getValue() < 1) |
| 840 | return nullptr; |
| 841 | |
| 842 | if (lhsConst) |
| 843 | return getAffineConstantExpr(mod(lhsConst.getValue(), rhsConst.getValue()), |
| 844 | lhs.getContext()); |
| 845 | |
| 846 | // Fold modulo of an expression that is known to be a multiple of a constant |
| 847 | // to zero if that constant is a multiple of the modulo factor. Eg: (i * 128) |
| 848 | // mod 64 is folded to 0, and less trivially, (i*(j*4*(k*32))) mod 128 = 0. |
| 849 | if (lhs.getLargestKnownDivisor() % rhsConst.getValue() == 0) |
| 850 | return getAffineConstantExpr(0, lhs.getContext()); |
| 851 | |
| 852 | // Simplify (expr1 + expr2) mod divConst when either expr1 or expr2 is |
| 853 | // known to be a multiple of divConst. |
| 854 | auto lBin = lhs.dyn_cast<AffineBinaryOpExpr>(); |
| 855 | if (lBin && lBin.getKind() == AffineExprKind::Add) { |
| 856 | int64_t llhsDiv = lBin.getLHS().getLargestKnownDivisor(); |
| 857 | int64_t lrhsDiv = lBin.getRHS().getLargestKnownDivisor(); |
| 858 | // rhsConst is known to be a positive constant. |
| 859 | if (llhsDiv % rhsConst.getValue() == 0) |
| 860 | return lBin.getRHS() % rhsConst.getValue(); |
| 861 | if (lrhsDiv % rhsConst.getValue() == 0) |
| 862 | return lBin.getLHS() % rhsConst.getValue(); |
| 863 | } |
| 864 | |
| 865 | // Simplify (e % a) % b to e % b when b evenly divides a |
| 866 | if (lBin && lBin.getKind() == AffineExprKind::Mod) { |
| 867 | auto intermediate = lBin.getRHS().dyn_cast<AffineConstantExpr>(); |
| 868 | if (intermediate && intermediate.getValue() >= 1 && |
| 869 | mod(intermediate.getValue(), rhsConst.getValue()) == 0) { |
| 870 | return lBin.getLHS() % rhsConst.getValue(); |
| 871 | } |
| 872 | } |
| 873 | |
| 874 | return nullptr; |
| 875 | } |
| 876 | |
| 877 | AffineExpr AffineExpr::operator%(uint64_t v) const { |
| 878 | return *this % getAffineConstantExpr(v, getContext()); |
| 879 | } |
| 880 | AffineExpr AffineExpr::operator%(AffineExpr other) const { |
| 881 | if (auto simplified = simplifyMod(*this, other)) |
| 882 | return simplified; |
| 883 | |
| 884 | StorageUniquer &uniquer = getContext()->getAffineUniquer(); |
| 885 | return uniquer.get<AffineBinaryOpExprStorage>( |
| 886 | /*initFn=*/{}, static_cast<unsigned>(AffineExprKind::Mod), *this, other); |
| 887 | } |
| 888 | |
| 889 | AffineExpr AffineExpr::compose(AffineMap map) const { |
| 890 | SmallVector<AffineExpr, 8> dimReplacements(map.getResults().begin(), |
| 891 | map.getResults().end()); |
| 892 | return replaceDimsAndSymbols(dimReplacements, {}); |
| 893 | } |
| 894 | raw_ostream &mlir::operator<<(raw_ostream &os, AffineExpr expr) { |
| 895 | expr.print(os); |
| 896 | return os; |
| 897 | } |
| 898 | |
| 899 | /// Constructs an affine expression from a flat ArrayRef. If there are local |
| 900 | /// identifiers (neither dimensional nor symbolic) that appear in the sum of |
| 901 | /// products expression, `localExprs` is expected to have the AffineExpr |
| 902 | /// for it, and is substituted into. The ArrayRef `flatExprs` is expected to be |
| 903 | /// in the format [dims, symbols, locals, constant term]. |
| 904 | AffineExpr mlir::getAffineExprFromFlatForm(ArrayRef<int64_t> flatExprs, |
| 905 | unsigned numDims, |
| 906 | unsigned numSymbols, |
| 907 | ArrayRef<AffineExpr> localExprs, |
| 908 | MLIRContext *context) { |
| 909 | // Assert expected numLocals = flatExprs.size() - numDims - numSymbols - 1. |
| 910 | assert(flatExprs.size() - numDims - numSymbols - 1 == localExprs.size() &&(static_cast <bool> (flatExprs.size() - numDims - numSymbols - 1 == localExprs.size() && "unexpected number of local expressions" ) ? void (0) : __assert_fail ("flatExprs.size() - numDims - numSymbols - 1 == localExprs.size() && \"unexpected number of local expressions\"" , "mlir/lib/IR/AffineExpr.cpp", 911, __extension__ __PRETTY_FUNCTION__ )) |
| 911 | "unexpected number of local expressions")(static_cast <bool> (flatExprs.size() - numDims - numSymbols - 1 == localExprs.size() && "unexpected number of local expressions" ) ? void (0) : __assert_fail ("flatExprs.size() - numDims - numSymbols - 1 == localExprs.size() && \"unexpected number of local expressions\"" , "mlir/lib/IR/AffineExpr.cpp", 911, __extension__ __PRETTY_FUNCTION__ )); |
| 912 | |
| 913 | auto expr = getAffineConstantExpr(0, context); |
| 914 | // Dimensions and symbols. |
| 915 | for (unsigned j = 0; j < numDims + numSymbols; j++) { |
| 916 | if (flatExprs[j] == 0) |
| 917 | continue; |
| 918 | auto id = j < numDims ? getAffineDimExpr(j, context) |
| 919 | : getAffineSymbolExpr(j - numDims, context); |
| 920 | expr = expr + id * flatExprs[j]; |
| 921 | } |
| 922 | |
| 923 | // Local identifiers. |
| 924 | for (unsigned j = numDims + numSymbols, e = flatExprs.size() - 1; j < e; |
| 925 | j++) { |
| 926 | if (flatExprs[j] == 0) |
| 927 | continue; |
| 928 | auto term = localExprs[j - numDims - numSymbols] * flatExprs[j]; |
| 929 | expr = expr + term; |
| 930 | } |
| 931 | |
| 932 | // Constant term. |
| 933 | int64_t constTerm = flatExprs[flatExprs.size() - 1]; |
| 934 | if (constTerm != 0) |
| 935 | expr = expr + constTerm; |
| 936 | return expr; |
| 937 | } |
| 938 | |
| 939 | /// Constructs a semi-affine expression from a flat ArrayRef. If there are |
| 940 | /// local identifiers (neither dimensional nor symbolic) that appear in the sum |
| 941 | /// of products expression, `localExprs` is expected to have the AffineExprs for |
| 942 | /// it, and is substituted into. The ArrayRef `flatExprs` is expected to be in |
| 943 | /// the format [dims, symbols, locals, constant term]. The semi-affine |
| 944 | /// expression is constructed in the sorted order of dimension and symbol |
| 945 | /// position numbers. Note: local expressions/ids are used for mod, div as well |
| 946 | /// as symbolic RHS terms for terms that are not pure affine. |
| 947 | static AffineExpr getSemiAffineExprFromFlatForm(ArrayRef<int64_t> flatExprs, |
| 948 | unsigned numDims, |
| 949 | unsigned numSymbols, |
| 950 | ArrayRef<AffineExpr> localExprs, |
| 951 | MLIRContext *context) { |
| 952 | assert(!flatExprs.empty() && "flatExprs cannot be empty")(static_cast <bool> (!flatExprs.empty() && "flatExprs cannot be empty" ) ? void (0) : __assert_fail ("!flatExprs.empty() && \"flatExprs cannot be empty\"" , "mlir/lib/IR/AffineExpr.cpp", 952, __extension__ __PRETTY_FUNCTION__ )); |
| 953 | |
| 954 | // Assert expected numLocals = flatExprs.size() - numDims - numSymbols - 1. |
| 955 | assert(flatExprs.size() - numDims - numSymbols - 1 == localExprs.size() &&(static_cast <bool> (flatExprs.size() - numDims - numSymbols - 1 == localExprs.size() && "unexpected number of local expressions" ) ? void (0) : __assert_fail ("flatExprs.size() - numDims - numSymbols - 1 == localExprs.size() && \"unexpected number of local expressions\"" , "mlir/lib/IR/AffineExpr.cpp", 956, __extension__ __PRETTY_FUNCTION__ )) |
| 956 | "unexpected number of local expressions")(static_cast <bool> (flatExprs.size() - numDims - numSymbols - 1 == localExprs.size() && "unexpected number of local expressions" ) ? void (0) : __assert_fail ("flatExprs.size() - numDims - numSymbols - 1 == localExprs.size() && \"unexpected number of local expressions\"" , "mlir/lib/IR/AffineExpr.cpp", 956, __extension__ __PRETTY_FUNCTION__ )); |
| 957 | |
| 958 | AffineExpr expr = getAffineConstantExpr(0, context); |
| 959 | |
| 960 | // We design indices as a pair which help us present the semi-affine map as |
| 961 | // sum of product where terms are sorted based on dimension or symbol |
| 962 | // position: <keyA, keyB> for expressions of the form dimension * symbol, |
| 963 | // where keyA is the position number of the dimension and keyB is the |
| 964 | // position number of the symbol. For dimensional expressions we set the index |
| 965 | // as (position number of the dimension, -1), as we want dimensional |
| 966 | // expressions to appear before symbolic and product of dimensional and |
| 967 | // symbolic expressions having the dimension with the same position number. |
| 968 | // For symbolic expression set the index as (position number of the symbol, |
| 969 | // maximum of last dimension and symbol position) number. For example, we want |
| 970 | // the expression we are constructing to look something like: d0 + d0 * s0 + |
| 971 | // s0 + d1*s1 + s1. |
| 972 | |
| 973 | // Stores the affine expression corresponding to a given index. |
| 974 | DenseMap<std::pair<unsigned, signed>, AffineExpr> indexToExprMap; |
| 975 | // Stores the constant coefficient value corresponding to a given |
| 976 | // dimension, symbol or a non-pure affine expression stored in `localExprs`. |
| 977 | DenseMap<std::pair<unsigned, signed>, int64_t> coefficients; |
| 978 | // Stores the indices as defined above, and later sorted to produce |
| 979 | // the semi-affine expression in the desired form. |
| 980 | SmallVector<std::pair<unsigned, signed>, 8> indices; |
| 981 | |
| 982 | // Example: expression = d0 + d0 * s0 + 2 * s0. |
| 983 | // indices = [{0,-1}, {0, 0}, {0, 1}] |
| 984 | // coefficients = [{{0, -1}, 1}, {{0, 0}, 1}, {{0, 1}, 2}] |
| 985 | // indexToExprMap = [{{0, -1}, d0}, {{0, 0}, d0 * s0}, {{0, 1}, s0}] |
| 986 | |
| 987 | // Adds entries to `indexToExprMap`, `coefficients` and `indices`. |
| 988 | auto addEntry = [&](std::pair<unsigned, signed> index, int64_t coefficient, |
| 989 | AffineExpr expr) { |
| 990 | assert(!llvm::is_contained(indices, index) &&(static_cast <bool> (!llvm::is_contained(indices, index ) && "Key is already present in indices vector and overwriting will " "happen in `indexToExprMap` and `coefficients`!") ? void (0) : __assert_fail ("!llvm::is_contained(indices, index) && \"Key is already present in indices vector and overwriting will \" \"happen in `indexToExprMap` and `coefficients`!\"" , "mlir/lib/IR/AffineExpr.cpp", 992, __extension__ __PRETTY_FUNCTION__ )) |
| 991 | "Key is already present in indices vector and overwriting will "(static_cast <bool> (!llvm::is_contained(indices, index ) && "Key is already present in indices vector and overwriting will " "happen in `indexToExprMap` and `coefficients`!") ? void (0) : __assert_fail ("!llvm::is_contained(indices, index) && \"Key is already present in indices vector and overwriting will \" \"happen in `indexToExprMap` and `coefficients`!\"" , "mlir/lib/IR/AffineExpr.cpp", 992, __extension__ __PRETTY_FUNCTION__ )) |
| 992 | "happen in `indexToExprMap` and `coefficients`!")(static_cast <bool> (!llvm::is_contained(indices, index ) && "Key is already present in indices vector and overwriting will " "happen in `indexToExprMap` and `coefficients`!") ? void (0) : __assert_fail ("!llvm::is_contained(indices, index) && \"Key is already present in indices vector and overwriting will \" \"happen in `indexToExprMap` and `coefficients`!\"" , "mlir/lib/IR/AffineExpr.cpp", 992, __extension__ __PRETTY_FUNCTION__ )); |
| 993 | |
| 994 | indices.push_back(index); |
| 995 | coefficients.insert({index, coefficient}); |
| 996 | indexToExprMap.insert({index, expr}); |
| 997 | }; |
| 998 | |
| 999 | // Design indices for dimensional or symbolic terms, and store the indices, |
| 1000 | // constant coefficient corresponding to the indices in `coefficients` map, |
| 1001 | // and affine expression corresponding to indices in `indexToExprMap` map. |
| 1002 | |
| 1003 | // Ensure we do not have duplicate keys in `indexToExpr` map. |
| 1004 | unsigned offsetSym = 0; |
| 1005 | signed offsetDim = -1; |
| 1006 | for (unsigned j = numDims; j < numDims + numSymbols; ++j) { |
| 1007 | if (flatExprs[j] == 0) |
| 1008 | continue; |
| 1009 | // For symbolic expression set the index as <position number |
| 1010 | // of the symbol, max(dimCount, symCount)> number, |
| 1011 | // as we want symbolic expressions with the same positional number to |
| 1012 | // appear after dimensional expressions having the same positional number. |
| 1013 | std::pair<unsigned, signed> indexEntry( |
| 1014 | j - numDims, std::max(numDims, numSymbols) + offsetSym++); |
| 1015 | addEntry(indexEntry, flatExprs[j], |
| 1016 | getAffineSymbolExpr(j - numDims, context)); |
| 1017 | } |
| 1018 | |
| 1019 | // Denotes semi-affine product, modulo or division terms, which has been added |
| 1020 | // to the `indexToExpr` map. |
| 1021 | SmallVector<bool, 4> addedToMap(flatExprs.size() - numDims - numSymbols - 1, |
| 1022 | false); |
| 1023 | unsigned lhsPos, rhsPos; |
| 1024 | // Construct indices for product terms involving dimension, symbol or constant |
| 1025 | // as lhs/rhs, and store the indices, constant coefficient corresponding to |
| 1026 | // the indices in `coefficients` map, and affine expression corresponding to |
| 1027 | // in indices in `indexToExprMap` map. |
| 1028 | for (const auto &it : llvm::enumerate(localExprs)) { |
| 1029 | AffineExpr expr = it.value(); |
| 1030 | if (flatExprs[numDims + numSymbols + it.index()] == 0) |
| 1031 | continue; |
| 1032 | AffineExpr lhs = expr.cast<AffineBinaryOpExpr>().getLHS(); |
| 1033 | AffineExpr rhs = expr.cast<AffineBinaryOpExpr>().getRHS(); |
| 1034 | if (!((lhs.isa<AffineDimExpr>() || lhs.isa<AffineSymbolExpr>()) && |
| 1035 | (rhs.isa<AffineDimExpr>() || rhs.isa<AffineSymbolExpr>() || |
| 1036 | rhs.isa<AffineConstantExpr>()))) { |
| 1037 | continue; |
| 1038 | } |
| 1039 | if (rhs.isa<AffineConstantExpr>()) { |
| 1040 | // For product/modulo/division expressions, when rhs of modulo/division |
| 1041 | // expression is constant, we put 0 in place of keyB, because we want |
| 1042 | // them to appear earlier in the semi-affine expression we are |
| 1043 | // constructing. When rhs is constant, we place 0 in place of keyB. |
| 1044 | if (lhs.isa<AffineDimExpr>()) { |
| 1045 | lhsPos = lhs.cast<AffineDimExpr>().getPosition(); |
| 1046 | std::pair<unsigned, signed> indexEntry(lhsPos, offsetDim--); |
| 1047 | addEntry(indexEntry, flatExprs[numDims + numSymbols + it.index()], |
| 1048 | expr); |
| 1049 | } else { |
| 1050 | lhsPos = lhs.cast<AffineSymbolExpr>().getPosition(); |
| 1051 | std::pair<unsigned, signed> indexEntry( |
| 1052 | lhsPos, std::max(numDims, numSymbols) + offsetSym++); |
| 1053 | addEntry(indexEntry, flatExprs[numDims + numSymbols + it.index()], |
| 1054 | expr); |
| 1055 | } |
| 1056 | } else if (lhs.isa<AffineDimExpr>()) { |
| 1057 | // For product/modulo/division expressions having lhs as dimension and rhs |
| 1058 | // as symbol, we order the terms in the semi-affine expression based on |
| 1059 | // the pair: <keyA, keyB> for expressions of the form dimension * symbol, |
| 1060 | // where keyA is the position number of the dimension and keyB is the |
| 1061 | // position number of the symbol. |
| 1062 | lhsPos = lhs.cast<AffineDimExpr>().getPosition(); |
| 1063 | rhsPos = rhs.cast<AffineSymbolExpr>().getPosition(); |
| 1064 | std::pair<unsigned, signed> indexEntry(lhsPos, rhsPos); |
| 1065 | addEntry(indexEntry, flatExprs[numDims + numSymbols + it.index()], expr); |
| 1066 | } else { |
| 1067 | // For product/modulo/division expressions having both lhs and rhs as |
| 1068 | // symbol, we design indices as a pair: <keyA, keyB> for expressions |
| 1069 | // of the form dimension * symbol, where keyA is the position number of |
| 1070 | // the dimension and keyB is the position number of the symbol. |
| 1071 | lhsPos = lhs.cast<AffineSymbolExpr>().getPosition(); |
| 1072 | rhsPos = rhs.cast<AffineSymbolExpr>().getPosition(); |
Value stored to 'rhsPos' is never read | |
| 1073 | std::pair<unsigned, signed> indexEntry( |
| 1074 | lhsPos, std::max(numDims, numSymbols) + offsetSym++); |
| 1075 | addEntry(indexEntry, flatExprs[numDims + numSymbols + it.index()], expr); |
| 1076 | } |
| 1077 | addedToMap[it.index()] = true; |
| 1078 | } |
| 1079 | |
| 1080 | for (unsigned j = 0; j < numDims; ++j) { |
| 1081 | if (flatExprs[j] == 0) |
| 1082 | continue; |
| 1083 | // For dimensional expressions we set the index as <position number of the |
| 1084 | // dimension, 0>, as we want dimensional expressions to appear before |
| 1085 | // symbolic ones and products of dimensional and symbolic expressions |
| 1086 | // having the dimension with the same position number. |
| 1087 | std::pair<unsigned, signed> indexEntry(j, offsetDim--); |
| 1088 | addEntry(indexEntry, flatExprs[j], getAffineDimExpr(j, context)); |
| 1089 | } |
| 1090 | |
| 1091 | // Constructing the simplified semi-affine sum of product/division/mod |
| 1092 | // expression from the flattened form in the desired sorted order of indices |
| 1093 | // of the various individual product/division/mod expressions. |
| 1094 | llvm::sort(indices); |
| 1095 | for (const std::pair<unsigned, unsigned> index : indices) { |
| 1096 | assert(indexToExprMap.lookup(index) &&(static_cast <bool> (indexToExprMap.lookup(index) && "cannot find key in `indexToExprMap` map") ? void (0) : __assert_fail ("indexToExprMap.lookup(index) && \"cannot find key in `indexToExprMap` map\"" , "mlir/lib/IR/AffineExpr.cpp", 1097, __extension__ __PRETTY_FUNCTION__ )) |
| 1097 | "cannot find key in `indexToExprMap` map")(static_cast <bool> (indexToExprMap.lookup(index) && "cannot find key in `indexToExprMap` map") ? void (0) : __assert_fail ("indexToExprMap.lookup(index) && \"cannot find key in `indexToExprMap` map\"" , "mlir/lib/IR/AffineExpr.cpp", 1097, __extension__ __PRETTY_FUNCTION__ )); |
| 1098 | expr = expr + indexToExprMap.lookup(index) * coefficients.lookup(index); |
| 1099 | } |
| 1100 | |
| 1101 | // Local identifiers. |
| 1102 | for (unsigned j = numDims + numSymbols, e = flatExprs.size() - 1; j < e; |
| 1103 | j++) { |
| 1104 | // If the coefficient of the local expression is 0, continue as we need not |
| 1105 | // add it in out final expression. |
| 1106 | if (flatExprs[j] == 0 || addedToMap[j - numDims - numSymbols]) |
| 1107 | continue; |
| 1108 | auto term = localExprs[j - numDims - numSymbols] * flatExprs[j]; |
| 1109 | expr = expr + term; |
| 1110 | } |
| 1111 | |
| 1112 | // Constant term. |
| 1113 | int64_t constTerm = flatExprs.back(); |
| 1114 | if (constTerm != 0) |
| 1115 | expr = expr + constTerm; |
| 1116 | return expr; |
| 1117 | } |
| 1118 | |
| 1119 | SimpleAffineExprFlattener::SimpleAffineExprFlattener(unsigned numDims, |
| 1120 | unsigned numSymbols) |
| 1121 | : numDims(numDims), numSymbols(numSymbols), numLocals(0) { |
| 1122 | operandExprStack.reserve(8); |
| 1123 | } |
| 1124 | |
| 1125 | // In pure affine t = expr * c, we multiply each coefficient of lhs with c. |
| 1126 | // |
| 1127 | // In case of semi affine multiplication expressions, t = expr * symbolic_expr, |
| 1128 | // introduce a local variable p (= expr * symbolic_expr), and the affine |
| 1129 | // expression expr * symbolic_expr is added to `localExprs`. |
| 1130 | void SimpleAffineExprFlattener::visitMulExpr(AffineBinaryOpExpr expr) { |
| 1131 | assert(operandExprStack.size() >= 2)(static_cast <bool> (operandExprStack.size() >= 2) ? void (0) : __assert_fail ("operandExprStack.size() >= 2", "mlir/lib/IR/AffineExpr.cpp", 1131, __extension__ __PRETTY_FUNCTION__ )); |
| 1132 | SmallVector<int64_t, 8> rhs = operandExprStack.back(); |
| 1133 | operandExprStack.pop_back(); |
| 1134 | SmallVector<int64_t, 8> &lhs = operandExprStack.back(); |
| 1135 | |
| 1136 | // Flatten semi-affine multiplication expressions by introducing a local |
| 1137 | // variable in place of the product; the affine expression |
| 1138 | // corresponding to the quantifier is added to `localExprs`. |
| 1139 | if (!expr.getRHS().isa<AffineConstantExpr>()) { |
| 1140 | MLIRContext *context = expr.getContext(); |
| 1141 | AffineExpr a = getAffineExprFromFlatForm(lhs, numDims, numSymbols, |
| 1142 | localExprs, context); |
| 1143 | AffineExpr b = getAffineExprFromFlatForm(rhs, numDims, numSymbols, |
| 1144 | localExprs, context); |
| 1145 | addLocalVariableSemiAffine(a * b, lhs, lhs.size()); |
| 1146 | return; |
| 1147 | } |
| 1148 | |
| 1149 | // Get the RHS constant. |
| 1150 | auto rhsConst = rhs[getConstantIndex()]; |
| 1151 | for (unsigned i = 0, e = lhs.size(); i < e; i++) { |
| 1152 | lhs[i] *= rhsConst; |
| 1153 | } |
| 1154 | } |
| 1155 | |
| 1156 | void SimpleAffineExprFlattener::visitAddExpr(AffineBinaryOpExpr expr) { |
| 1157 | assert(operandExprStack.size() >= 2)(static_cast <bool> (operandExprStack.size() >= 2) ? void (0) : __assert_fail ("operandExprStack.size() >= 2", "mlir/lib/IR/AffineExpr.cpp", 1157, __extension__ __PRETTY_FUNCTION__ )); |
| 1158 | const auto &rhs = operandExprStack.back(); |
| 1159 | auto &lhs = operandExprStack[operandExprStack.size() - 2]; |
| 1160 | assert(lhs.size() == rhs.size())(static_cast <bool> (lhs.size() == rhs.size()) ? void ( 0) : __assert_fail ("lhs.size() == rhs.size()", "mlir/lib/IR/AffineExpr.cpp" , 1160, __extension__ __PRETTY_FUNCTION__)); |
| 1161 | // Update the LHS in place. |
| 1162 | for (unsigned i = 0, e = rhs.size(); i < e; i++) { |
| 1163 | lhs[i] += rhs[i]; |
| 1164 | } |
| 1165 | // Pop off the RHS. |
| 1166 | operandExprStack.pop_back(); |
| 1167 | } |
| 1168 | |
| 1169 | // |
| 1170 | // t = expr mod c <=> t = expr - c*q and c*q <= expr <= c*q + c - 1 |
| 1171 | // |
| 1172 | // A mod expression "expr mod c" is thus flattened by introducing a new local |
| 1173 | // variable q (= expr floordiv c), such that expr mod c is replaced with |
| 1174 | // 'expr - c * q' and c * q <= expr <= c * q + c - 1 are added to localVarCst. |
| 1175 | // |
| 1176 | // In case of semi-affine modulo expressions, t = expr mod symbolic_expr, |
| 1177 | // introduce a local variable m (= expr mod symbolic_expr), and the affine |
| 1178 | // expression expr mod symbolic_expr is added to `localExprs`. |
| 1179 | void SimpleAffineExprFlattener::visitModExpr(AffineBinaryOpExpr expr) { |
| 1180 | assert(operandExprStack.size() >= 2)(static_cast <bool> (operandExprStack.size() >= 2) ? void (0) : __assert_fail ("operandExprStack.size() >= 2", "mlir/lib/IR/AffineExpr.cpp", 1180, __extension__ __PRETTY_FUNCTION__ )); |
| 1181 | |
| 1182 | SmallVector<int64_t, 8> rhs = operandExprStack.back(); |
| 1183 | operandExprStack.pop_back(); |
| 1184 | SmallVector<int64_t, 8> &lhs = operandExprStack.back(); |
| 1185 | MLIRContext *context = expr.getContext(); |
| 1186 | |
| 1187 | // Flatten semi affine modulo expressions by introducing a local |
| 1188 | // variable in place of the modulo value, and the affine expression |
| 1189 | // corresponding to the quantifier is added to `localExprs`. |
| 1190 | if (!expr.getRHS().isa<AffineConstantExpr>()) { |
| 1191 | AffineExpr dividendExpr = getAffineExprFromFlatForm( |
| 1192 | lhs, numDims, numSymbols, localExprs, context); |
| 1193 | AffineExpr divisorExpr = getAffineExprFromFlatForm(rhs, numDims, numSymbols, |
| 1194 | localExprs, context); |
| 1195 | AffineExpr modExpr = dividendExpr % divisorExpr; |
| 1196 | addLocalVariableSemiAffine(modExpr, lhs, lhs.size()); |
| 1197 | return; |
| 1198 | } |
| 1199 | |
| 1200 | int64_t rhsConst = rhs[getConstantIndex()]; |
| 1201 | // TODO: handle modulo by zero case when this issue is fixed |
| 1202 | // at the other places in the IR. |
| 1203 | assert(rhsConst > 0 && "RHS constant has to be positive")(static_cast <bool> (rhsConst > 0 && "RHS constant has to be positive" ) ? void (0) : __assert_fail ("rhsConst > 0 && \"RHS constant has to be positive\"" , "mlir/lib/IR/AffineExpr.cpp", 1203, __extension__ __PRETTY_FUNCTION__ )); |
| 1204 | |
| 1205 | // Check if the LHS expression is a multiple of modulo factor. |
| 1206 | unsigned i, e; |
| 1207 | for (i = 0, e = lhs.size(); i < e; i++) |
| 1208 | if (lhs[i] % rhsConst != 0) |
| 1209 | break; |
| 1210 | // If yes, modulo expression here simplifies to zero. |
| 1211 | if (i == lhs.size()) { |
| 1212 | std::fill(lhs.begin(), lhs.end(), 0); |
| 1213 | return; |
| 1214 | } |
| 1215 | |
| 1216 | // Add a local variable for the quotient, i.e., expr % c is replaced by |
| 1217 | // (expr - q * c) where q = expr floordiv c. Do this while canceling out |
| 1218 | // the GCD of expr and c. |
| 1219 | SmallVector<int64_t, 8> floorDividend(lhs); |
| 1220 | uint64_t gcd = rhsConst; |
| 1221 | for (unsigned i = 0, e = lhs.size(); i < e; i++) |
| 1222 | gcd = std::gcd(gcd, (uint64_t)std::abs(lhs[i])); |
| 1223 | // Simplify the numerator and the denominator. |
| 1224 | if (gcd != 1) { |
| 1225 | for (unsigned i = 0, e = floorDividend.size(); i < e; i++) |
| 1226 | floorDividend[i] = floorDividend[i] / static_cast<int64_t>(gcd); |
| 1227 | } |
| 1228 | int64_t floorDivisor = rhsConst / static_cast<int64_t>(gcd); |
| 1229 | |
| 1230 | // Construct the AffineExpr form of the floordiv to store in localExprs. |
| 1231 | |
| 1232 | AffineExpr dividendExpr = getAffineExprFromFlatForm( |
| 1233 | floorDividend, numDims, numSymbols, localExprs, context); |
| 1234 | AffineExpr divisorExpr = getAffineConstantExpr(floorDivisor, context); |
| 1235 | AffineExpr floorDivExpr = dividendExpr.floorDiv(divisorExpr); |
| 1236 | int loc; |
| 1237 | if ((loc = findLocalId(floorDivExpr)) == -1) { |
| 1238 | addLocalFloorDivId(floorDividend, floorDivisor, floorDivExpr); |
| 1239 | // Set result at top of stack to "lhs - rhsConst * q". |
| 1240 | lhs[getLocalVarStartIndex() + numLocals - 1] = -rhsConst; |
| 1241 | } else { |
| 1242 | // Reuse the existing local id. |
| 1243 | lhs[getLocalVarStartIndex() + loc] = -rhsConst; |
| 1244 | } |
| 1245 | } |
| 1246 | |
| 1247 | void SimpleAffineExprFlattener::visitCeilDivExpr(AffineBinaryOpExpr expr) { |
| 1248 | visitDivExpr(expr, /*isCeil=*/true); |
| 1249 | } |
| 1250 | void SimpleAffineExprFlattener::visitFloorDivExpr(AffineBinaryOpExpr expr) { |
| 1251 | visitDivExpr(expr, /*isCeil=*/false); |
| 1252 | } |
| 1253 | |
| 1254 | void SimpleAffineExprFlattener::visitDimExpr(AffineDimExpr expr) { |
| 1255 | operandExprStack.emplace_back(SmallVector<int64_t, 32>(getNumCols(), 0)); |
| 1256 | auto &eq = operandExprStack.back(); |
| 1257 | assert(expr.getPosition() < numDims && "Inconsistent number of dims")(static_cast <bool> (expr.getPosition() < numDims && "Inconsistent number of dims") ? void (0) : __assert_fail ("expr.getPosition() < numDims && \"Inconsistent number of dims\"" , "mlir/lib/IR/AffineExpr.cpp", 1257, __extension__ __PRETTY_FUNCTION__ )); |
| 1258 | eq[getDimStartIndex() + expr.getPosition()] = 1; |
| 1259 | } |
| 1260 | |
| 1261 | void SimpleAffineExprFlattener::visitSymbolExpr(AffineSymbolExpr expr) { |
| 1262 | operandExprStack.emplace_back(SmallVector<int64_t, 32>(getNumCols(), 0)); |
| 1263 | auto &eq = operandExprStack.back(); |
| 1264 | assert(expr.getPosition() < numSymbols && "inconsistent number of symbols")(static_cast <bool> (expr.getPosition() < numSymbols && "inconsistent number of symbols") ? void (0) : __assert_fail ("expr.getPosition() < numSymbols && \"inconsistent number of symbols\"" , "mlir/lib/IR/AffineExpr.cpp", 1264, __extension__ __PRETTY_FUNCTION__ )); |
| 1265 | eq[getSymbolStartIndex() + expr.getPosition()] = 1; |
| 1266 | } |
| 1267 | |
| 1268 | void SimpleAffineExprFlattener::visitConstantExpr(AffineConstantExpr expr) { |
| 1269 | operandExprStack.emplace_back(SmallVector<int64_t, 32>(getNumCols(), 0)); |
| 1270 | auto &eq = operandExprStack.back(); |
| 1271 | eq[getConstantIndex()] = expr.getValue(); |
| 1272 | } |
| 1273 | |
| 1274 | void SimpleAffineExprFlattener::addLocalVariableSemiAffine( |
| 1275 | AffineExpr expr, SmallVectorImpl<int64_t> &result, |
| 1276 | unsigned long resultSize) { |
| 1277 | assert(result.size() == resultSize &&(static_cast <bool> (result.size() == resultSize && "`result` vector passed is not of correct size") ? void (0) : __assert_fail ("result.size() == resultSize && \"`result` vector passed is not of correct size\"" , "mlir/lib/IR/AffineExpr.cpp", 1278, __extension__ __PRETTY_FUNCTION__ )) |
| 1278 | "`result` vector passed is not of correct size")(static_cast <bool> (result.size() == resultSize && "`result` vector passed is not of correct size") ? void (0) : __assert_fail ("result.size() == resultSize && \"`result` vector passed is not of correct size\"" , "mlir/lib/IR/AffineExpr.cpp", 1278, __extension__ __PRETTY_FUNCTION__ )); |
| 1279 | int loc; |
| 1280 | if ((loc = findLocalId(expr)) == -1) |
| 1281 | addLocalIdSemiAffine(expr); |
| 1282 | std::fill(result.begin(), result.end(), 0); |
| 1283 | if (loc == -1) |
| 1284 | result[getLocalVarStartIndex() + numLocals - 1] = 1; |
| 1285 | else |
| 1286 | result[getLocalVarStartIndex() + loc] = 1; |
| 1287 | } |
| 1288 | |
| 1289 | // t = expr floordiv c <=> t = q, c * q <= expr <= c * q + c - 1 |
| 1290 | // A floordiv is thus flattened by introducing a new local variable q, and |
| 1291 | // replacing that expression with 'q' while adding the constraints |
| 1292 | // c * q <= expr <= c * q + c - 1 to localVarCst (done by |
| 1293 | // IntegerRelation::addLocalFloorDiv). |
| 1294 | // |
| 1295 | // A ceildiv is similarly flattened: |
| 1296 | // t = expr ceildiv c <=> t = (expr + c - 1) floordiv c |
| 1297 | // |
| 1298 | // In case of semi affine division expressions, t = expr floordiv symbolic_expr |
| 1299 | // or t = expr ceildiv symbolic_expr, introduce a local variable q (= expr |
| 1300 | // floordiv/ceildiv symbolic_expr), and the affine floordiv/ceildiv is added to |
| 1301 | // `localExprs`. |
| 1302 | void SimpleAffineExprFlattener::visitDivExpr(AffineBinaryOpExpr expr, |
| 1303 | bool isCeil) { |
| 1304 | assert(operandExprStack.size() >= 2)(static_cast <bool> (operandExprStack.size() >= 2) ? void (0) : __assert_fail ("operandExprStack.size() >= 2", "mlir/lib/IR/AffineExpr.cpp", 1304, __extension__ __PRETTY_FUNCTION__ )); |
| 1305 | |
| 1306 | MLIRContext *context = expr.getContext(); |
| 1307 | SmallVector<int64_t, 8> rhs = operandExprStack.back(); |
| 1308 | operandExprStack.pop_back(); |
| 1309 | SmallVector<int64_t, 8> &lhs = operandExprStack.back(); |
| 1310 | |
| 1311 | // Flatten semi affine division expressions by introducing a local |
| 1312 | // variable in place of the quotient, and the affine expression corresponding |
| 1313 | // to the quantifier is added to `localExprs`. |
| 1314 | if (!expr.getRHS().isa<AffineConstantExpr>()) { |
| 1315 | AffineExpr a = getAffineExprFromFlatForm(lhs, numDims, numSymbols, |
| 1316 | localExprs, context); |
| 1317 | AffineExpr b = getAffineExprFromFlatForm(rhs, numDims, numSymbols, |
| 1318 | localExprs, context); |
| 1319 | AffineExpr divExpr = isCeil ? a.ceilDiv(b) : a.floorDiv(b); |
| 1320 | addLocalVariableSemiAffine(divExpr, lhs, lhs.size()); |
| 1321 | return; |
| 1322 | } |
| 1323 | |
| 1324 | // This is a pure affine expr; the RHS is a positive constant. |
| 1325 | int64_t rhsConst = rhs[getConstantIndex()]; |
| 1326 | // TODO: handle division by zero at the same time the issue is |
| 1327 | // fixed at other places. |
| 1328 | assert(rhsConst > 0 && "RHS constant has to be positive")(static_cast <bool> (rhsConst > 0 && "RHS constant has to be positive" ) ? void (0) : __assert_fail ("rhsConst > 0 && \"RHS constant has to be positive\"" , "mlir/lib/IR/AffineExpr.cpp", 1328, __extension__ __PRETTY_FUNCTION__ )); |
| 1329 | |
| 1330 | // Simplify the floordiv, ceildiv if possible by canceling out the greatest |
| 1331 | // common divisors of the numerator and denominator. |
| 1332 | uint64_t gcd = std::abs(rhsConst); |
| 1333 | for (unsigned i = 0, e = lhs.size(); i < e; i++) |
| 1334 | gcd = std::gcd(gcd, (uint64_t)std::abs(lhs[i])); |
| 1335 | // Simplify the numerator and the denominator. |
| 1336 | if (gcd != 1) { |
| 1337 | for (unsigned i = 0, e = lhs.size(); i < e; i++) |
| 1338 | lhs[i] = lhs[i] / static_cast<int64_t>(gcd); |
| 1339 | } |
| 1340 | int64_t divisor = rhsConst / static_cast<int64_t>(gcd); |
| 1341 | // If the divisor becomes 1, the updated LHS is the result. (The |
| 1342 | // divisor can't be negative since rhsConst is positive). |
| 1343 | if (divisor == 1) |
| 1344 | return; |
| 1345 | |
| 1346 | // If the divisor cannot be simplified to one, we will have to retain |
| 1347 | // the ceil/floor expr (simplified up until here). Add an existential |
| 1348 | // quantifier to express its result, i.e., expr1 div expr2 is replaced |
| 1349 | // by a new identifier, q. |
| 1350 | AffineExpr a = |
| 1351 | getAffineExprFromFlatForm(lhs, numDims, numSymbols, localExprs, context); |
| 1352 | AffineExpr b = getAffineConstantExpr(divisor, context); |
| 1353 | |
| 1354 | int loc; |
| 1355 | AffineExpr divExpr = isCeil ? a.ceilDiv(b) : a.floorDiv(b); |
| 1356 | if ((loc = findLocalId(divExpr)) == -1) { |
| 1357 | if (!isCeil) { |
| 1358 | SmallVector<int64_t, 8> dividend(lhs); |
| 1359 | addLocalFloorDivId(dividend, divisor, divExpr); |
| 1360 | } else { |
| 1361 | // lhs ceildiv c <=> (lhs + c - 1) floordiv c |
| 1362 | SmallVector<int64_t, 8> dividend(lhs); |
| 1363 | dividend.back() += divisor - 1; |
| 1364 | addLocalFloorDivId(dividend, divisor, divExpr); |
| 1365 | } |
| 1366 | } |
| 1367 | // Set the expression on stack to the local var introduced to capture the |
| 1368 | // result of the division (floor or ceil). |
| 1369 | std::fill(lhs.begin(), lhs.end(), 0); |
| 1370 | if (loc == -1) |
| 1371 | lhs[getLocalVarStartIndex() + numLocals - 1] = 1; |
| 1372 | else |
| 1373 | lhs[getLocalVarStartIndex() + loc] = 1; |
| 1374 | } |
| 1375 | |
| 1376 | // Add a local identifier (needed to flatten a mod, floordiv, ceildiv expr). |
| 1377 | // The local identifier added is always a floordiv of a pure add/mul affine |
| 1378 | // function of other identifiers, coefficients of which are specified in |
| 1379 | // dividend and with respect to a positive constant divisor. localExpr is the |
| 1380 | // simplified tree expression (AffineExpr) corresponding to the quantifier. |
| 1381 | void SimpleAffineExprFlattener::addLocalFloorDivId(ArrayRef<int64_t> dividend, |
| 1382 | int64_t divisor, |
| 1383 | AffineExpr localExpr) { |
| 1384 | assert(divisor > 0 && "positive constant divisor expected")(static_cast <bool> (divisor > 0 && "positive constant divisor expected" ) ? void (0) : __assert_fail ("divisor > 0 && \"positive constant divisor expected\"" , "mlir/lib/IR/AffineExpr.cpp", 1384, __extension__ __PRETTY_FUNCTION__ )); |
| 1385 | for (SmallVector<int64_t, 8> &subExpr : operandExprStack) |
| 1386 | subExpr.insert(subExpr.begin() + getLocalVarStartIndex() + numLocals, 0); |
| 1387 | localExprs.push_back(localExpr); |
| 1388 | numLocals++; |
| 1389 | // dividend and divisor are not used here; an override of this method uses it. |
| 1390 | } |
| 1391 | |
| 1392 | void SimpleAffineExprFlattener::addLocalIdSemiAffine(AffineExpr localExpr) { |
| 1393 | for (SmallVector<int64_t, 8> &subExpr : operandExprStack) |
| 1394 | subExpr.insert(subExpr.begin() + getLocalVarStartIndex() + numLocals, 0); |
| 1395 | localExprs.push_back(localExpr); |
| 1396 | ++numLocals; |
| 1397 | } |
| 1398 | |
| 1399 | int SimpleAffineExprFlattener::findLocalId(AffineExpr localExpr) { |
| 1400 | SmallVectorImpl<AffineExpr>::iterator it; |
| 1401 | if ((it = llvm::find(localExprs, localExpr)) == localExprs.end()) |
| 1402 | return -1; |
| 1403 | return it - localExprs.begin(); |
| 1404 | } |
| 1405 | |
| 1406 | /// Simplify the affine expression by flattening it and reconstructing it. |
| 1407 | AffineExpr mlir::simplifyAffineExpr(AffineExpr expr, unsigned numDims, |
| 1408 | unsigned numSymbols) { |
| 1409 | // Simplify semi-affine expressions separately. |
| 1410 | if (!expr.isPureAffine()) |
| 1411 | expr = simplifySemiAffine(expr); |
| 1412 | |
| 1413 | SimpleAffineExprFlattener flattener(numDims, numSymbols); |
| 1414 | flattener.walkPostOrder(expr); |
| 1415 | ArrayRef<int64_t> flattenedExpr = flattener.operandExprStack.back(); |
| 1416 | if (!expr.isPureAffine() && |
| 1417 | expr == getAffineExprFromFlatForm(flattenedExpr, numDims, numSymbols, |
| 1418 | flattener.localExprs, |
| 1419 | expr.getContext())) |
| 1420 | return expr; |
| 1421 | AffineExpr simplifiedExpr = |
| 1422 | expr.isPureAffine() |
| 1423 | ? getAffineExprFromFlatForm(flattenedExpr, numDims, numSymbols, |
| 1424 | flattener.localExprs, expr.getContext()) |
| 1425 | : getSemiAffineExprFromFlatForm(flattenedExpr, numDims, numSymbols, |
| 1426 | flattener.localExprs, |
| 1427 | expr.getContext()); |
| 1428 | |
| 1429 | flattener.operandExprStack.pop_back(); |
| 1430 | assert(flattener.operandExprStack.empty())(static_cast <bool> (flattener.operandExprStack.empty() ) ? void (0) : __assert_fail ("flattener.operandExprStack.empty()" , "mlir/lib/IR/AffineExpr.cpp", 1430, __extension__ __PRETTY_FUNCTION__ )); |
| 1431 | return simplifiedExpr; |
| 1432 | } |