| File: | build/source/mlir/lib/IR/AffineExpr.cpp |
| Warning: | line 25, column 54 Access to field 'context' results in a dereference of a null pointer (loaded from field 'expr') |
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
| |||
| 50 | return lhs + rhs; | |||
| 51 | if (kind
| |||
| 52 | return lhs * rhs; | |||
| 53 | if (kind
| |||
| 54 | return lhs.floorDiv(rhs); | |||
| 55 | if (kind
| |||
| 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(); | |||
| 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 | } |