1 //===- AffineMap.cpp - MLIR Affine Map 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 "mlir/IR/AffineMap.h" 10 #include "AffineMapDetail.h" 11 #include "mlir/IR/BuiltinAttributes.h" 12 #include "mlir/IR/BuiltinTypes.h" 13 #include "mlir/Support/LogicalResult.h" 14 #include "mlir/Support/MathExtras.h" 15 #include "llvm/ADT/SmallSet.h" 16 #include "llvm/ADT/StringRef.h" 17 #include "llvm/Support/raw_ostream.h" 18 19 using namespace mlir; 20 21 namespace { 22 23 // AffineExprConstantFolder evaluates an affine expression using constant 24 // operands passed in 'operandConsts'. Returns an IntegerAttr attribute 25 // representing the constant value of the affine expression evaluated on 26 // constant 'operandConsts', or nullptr if it can't be folded. 27 class AffineExprConstantFolder { 28 public: 29 AffineExprConstantFolder(unsigned numDims, ArrayRef<Attribute> operandConsts) 30 : numDims(numDims), operandConsts(operandConsts) {} 31 32 /// Attempt to constant fold the specified affine expr, or return null on 33 /// failure. 34 IntegerAttr constantFold(AffineExpr expr) { 35 if (auto result = constantFoldImpl(expr)) 36 return IntegerAttr::get(IndexType::get(expr.getContext()), *result); 37 return nullptr; 38 } 39 40 private: 41 Optional<int64_t> constantFoldImpl(AffineExpr expr) { 42 switch (expr.getKind()) { 43 case AffineExprKind::Add: 44 return constantFoldBinExpr( 45 expr, [](int64_t lhs, int64_t rhs) { return lhs + rhs; }); 46 case AffineExprKind::Mul: 47 return constantFoldBinExpr( 48 expr, [](int64_t lhs, int64_t rhs) { return lhs * rhs; }); 49 case AffineExprKind::Mod: 50 return constantFoldBinExpr( 51 expr, [](int64_t lhs, int64_t rhs) { return mod(lhs, rhs); }); 52 case AffineExprKind::FloorDiv: 53 return constantFoldBinExpr( 54 expr, [](int64_t lhs, int64_t rhs) { return floorDiv(lhs, rhs); }); 55 case AffineExprKind::CeilDiv: 56 return constantFoldBinExpr( 57 expr, [](int64_t lhs, int64_t rhs) { return ceilDiv(lhs, rhs); }); 58 case AffineExprKind::Constant: 59 return expr.cast<AffineConstantExpr>().getValue(); 60 case AffineExprKind::DimId: 61 if (auto attr = operandConsts[expr.cast<AffineDimExpr>().getPosition()] 62 .dyn_cast_or_null<IntegerAttr>()) 63 return attr.getInt(); 64 return llvm::None; 65 case AffineExprKind::SymbolId: 66 if (auto attr = operandConsts[numDims + 67 expr.cast<AffineSymbolExpr>().getPosition()] 68 .dyn_cast_or_null<IntegerAttr>()) 69 return attr.getInt(); 70 return llvm::None; 71 } 72 llvm_unreachable("Unknown AffineExpr"); 73 } 74 75 // TODO: Change these to operate on APInts too. 76 Optional<int64_t> constantFoldBinExpr(AffineExpr expr, 77 int64_t (*op)(int64_t, int64_t)) { 78 auto binOpExpr = expr.cast<AffineBinaryOpExpr>(); 79 if (auto lhs = constantFoldImpl(binOpExpr.getLHS())) 80 if (auto rhs = constantFoldImpl(binOpExpr.getRHS())) 81 return op(*lhs, *rhs); 82 return llvm::None; 83 } 84 85 // The number of dimension operands in AffineMap containing this expression. 86 unsigned numDims; 87 // The constant valued operands used to evaluate this AffineExpr. 88 ArrayRef<Attribute> operandConsts; 89 }; 90 91 } // end anonymous namespace 92 93 /// Returns a single constant result affine map. 94 AffineMap AffineMap::getConstantMap(int64_t val, MLIRContext *context) { 95 return get(/*dimCount=*/0, /*symbolCount=*/0, 96 {getAffineConstantExpr(val, context)}); 97 } 98 99 /// Returns an identity affine map (d0, ..., dn) -> (dp, ..., dn) on the most 100 /// minor dimensions. 101 AffineMap AffineMap::getMinorIdentityMap(unsigned dims, unsigned results, 102 MLIRContext *context) { 103 assert(dims >= results && "Dimension mismatch"); 104 auto id = AffineMap::getMultiDimIdentityMap(dims, context); 105 return AffineMap::get(dims, 0, id.getResults().take_back(results), context); 106 } 107 108 bool AffineMap::isMinorIdentity() const { 109 return getNumDims() >= getNumResults() && 110 *this == 111 getMinorIdentityMap(getNumDims(), getNumResults(), getContext()); 112 } 113 114 /// Returns true if this affine map is a minor identity up to broadcasted 115 /// dimensions which are indicated by value 0 in the result. 116 bool AffineMap::isMinorIdentityWithBroadcasting( 117 SmallVectorImpl<unsigned> *broadcastedDims) const { 118 if (broadcastedDims) 119 broadcastedDims->clear(); 120 if (getNumDims() < getNumResults()) 121 return false; 122 unsigned suffixStart = getNumDims() - getNumResults(); 123 for (auto idxAndExpr : llvm::enumerate(getResults())) { 124 unsigned resIdx = idxAndExpr.index(); 125 AffineExpr expr = idxAndExpr.value(); 126 if (auto constExpr = expr.dyn_cast<AffineConstantExpr>()) { 127 // Each result may be either a constant 0 (broadcasted dimension). 128 if (constExpr.getValue() != 0) 129 return false; 130 if (broadcastedDims) 131 broadcastedDims->push_back(resIdx); 132 } else if (auto dimExpr = expr.dyn_cast<AffineDimExpr>()) { 133 // Or it may be the input dimension corresponding to this result position. 134 if (dimExpr.getPosition() != suffixStart + resIdx) 135 return false; 136 } else { 137 return false; 138 } 139 } 140 return true; 141 } 142 143 /// Returns an AffineMap representing a permutation. 144 AffineMap AffineMap::getPermutationMap(ArrayRef<unsigned> permutation, 145 MLIRContext *context) { 146 assert(!permutation.empty() && 147 "Cannot create permutation map from empty permutation vector"); 148 SmallVector<AffineExpr, 4> affExprs; 149 for (auto index : permutation) 150 affExprs.push_back(getAffineDimExpr(index, context)); 151 auto m = std::max_element(permutation.begin(), permutation.end()); 152 auto permutationMap = AffineMap::get(*m + 1, 0, affExprs, context); 153 assert(permutationMap.isPermutation() && "Invalid permutation vector"); 154 return permutationMap; 155 } 156 157 template <typename AffineExprContainer> 158 static void getMaxDimAndSymbol(ArrayRef<AffineExprContainer> exprsList, 159 int64_t &maxDim, int64_t &maxSym) { 160 for (const auto &exprs : exprsList) { 161 for (auto expr : exprs) { 162 expr.walk([&maxDim, &maxSym](AffineExpr e) { 163 if (auto d = e.dyn_cast<AffineDimExpr>()) 164 maxDim = std::max(maxDim, static_cast<int64_t>(d.getPosition())); 165 if (auto s = e.dyn_cast<AffineSymbolExpr>()) 166 maxSym = std::max(maxSym, static_cast<int64_t>(s.getPosition())); 167 }); 168 } 169 } 170 } 171 172 template <typename AffineExprContainer> 173 static SmallVector<AffineMap, 4> 174 inferFromExprList(ArrayRef<AffineExprContainer> exprsList) { 175 assert(!exprsList.empty()); 176 assert(!exprsList[0].empty()); 177 auto context = exprsList[0][0].getContext(); 178 int64_t maxDim = -1, maxSym = -1; 179 getMaxDimAndSymbol(exprsList, maxDim, maxSym); 180 SmallVector<AffineMap, 4> maps; 181 maps.reserve(exprsList.size()); 182 for (const auto &exprs : exprsList) 183 maps.push_back(AffineMap::get(/*dimCount=*/maxDim + 1, 184 /*symbolCount=*/maxSym + 1, exprs, context)); 185 return maps; 186 } 187 188 SmallVector<AffineMap, 4> 189 AffineMap::inferFromExprList(ArrayRef<ArrayRef<AffineExpr>> exprsList) { 190 return ::inferFromExprList(exprsList); 191 } 192 193 SmallVector<AffineMap, 4> 194 AffineMap::inferFromExprList(ArrayRef<SmallVector<AffineExpr, 4>> exprsList) { 195 return ::inferFromExprList(exprsList); 196 } 197 198 AffineMap AffineMap::getMultiDimIdentityMap(unsigned numDims, 199 MLIRContext *context) { 200 SmallVector<AffineExpr, 4> dimExprs; 201 dimExprs.reserve(numDims); 202 for (unsigned i = 0; i < numDims; ++i) 203 dimExprs.push_back(mlir::getAffineDimExpr(i, context)); 204 return get(/*dimCount=*/numDims, /*symbolCount=*/0, dimExprs, context); 205 } 206 207 MLIRContext *AffineMap::getContext() const { return map->context; } 208 209 bool AffineMap::isIdentity() const { 210 if (getNumDims() != getNumResults()) 211 return false; 212 ArrayRef<AffineExpr> results = getResults(); 213 for (unsigned i = 0, numDims = getNumDims(); i < numDims; ++i) { 214 auto expr = results[i].dyn_cast<AffineDimExpr>(); 215 if (!expr || expr.getPosition() != i) 216 return false; 217 } 218 return true; 219 } 220 221 bool AffineMap::isEmpty() const { 222 return getNumDims() == 0 && getNumSymbols() == 0 && getNumResults() == 0; 223 } 224 225 bool AffineMap::isSingleConstant() const { 226 return getNumResults() == 1 && getResult(0).isa<AffineConstantExpr>(); 227 } 228 229 int64_t AffineMap::getSingleConstantResult() const { 230 assert(isSingleConstant() && "map must have a single constant result"); 231 return getResult(0).cast<AffineConstantExpr>().getValue(); 232 } 233 234 unsigned AffineMap::getNumDims() const { 235 assert(map && "uninitialized map storage"); 236 return map->numDims; 237 } 238 unsigned AffineMap::getNumSymbols() const { 239 assert(map && "uninitialized map storage"); 240 return map->numSymbols; 241 } 242 unsigned AffineMap::getNumResults() const { 243 assert(map && "uninitialized map storage"); 244 return map->results.size(); 245 } 246 unsigned AffineMap::getNumInputs() const { 247 assert(map && "uninitialized map storage"); 248 return map->numDims + map->numSymbols; 249 } 250 251 ArrayRef<AffineExpr> AffineMap::getResults() const { 252 assert(map && "uninitialized map storage"); 253 return map->results; 254 } 255 AffineExpr AffineMap::getResult(unsigned idx) const { 256 assert(map && "uninitialized map storage"); 257 return map->results[idx]; 258 } 259 260 unsigned AffineMap::getDimPosition(unsigned idx) const { 261 return getResult(idx).cast<AffineDimExpr>().getPosition(); 262 } 263 264 /// Folds the results of the application of an affine map on the provided 265 /// operands to a constant if possible. Returns false if the folding happens, 266 /// true otherwise. 267 LogicalResult 268 AffineMap::constantFold(ArrayRef<Attribute> operandConstants, 269 SmallVectorImpl<Attribute> &results) const { 270 // Attempt partial folding. 271 SmallVector<int64_t, 2> integers; 272 partialConstantFold(operandConstants, &integers); 273 274 // If all expressions folded to a constant, populate results with attributes 275 // containing those constants. 276 if (integers.empty()) 277 return failure(); 278 279 auto range = llvm::map_range(integers, [this](int64_t i) { 280 return IntegerAttr::get(IndexType::get(getContext()), i); 281 }); 282 results.append(range.begin(), range.end()); 283 return success(); 284 } 285 286 AffineMap 287 AffineMap::partialConstantFold(ArrayRef<Attribute> operandConstants, 288 SmallVectorImpl<int64_t> *results) const { 289 assert(getNumInputs() == operandConstants.size()); 290 291 // Fold each of the result expressions. 292 AffineExprConstantFolder exprFolder(getNumDims(), operandConstants); 293 SmallVector<AffineExpr, 4> exprs; 294 exprs.reserve(getNumResults()); 295 296 for (auto expr : getResults()) { 297 auto folded = exprFolder.constantFold(expr); 298 // If did not fold to a constant, keep the original expression, and clear 299 // the integer results vector. 300 if (folded) { 301 exprs.push_back( 302 getAffineConstantExpr(folded.getInt(), folded.getContext())); 303 if (results) 304 results->push_back(folded.getInt()); 305 } else { 306 exprs.push_back(expr); 307 if (results) { 308 results->clear(); 309 results = nullptr; 310 } 311 } 312 } 313 314 return get(getNumDims(), getNumSymbols(), exprs, getContext()); 315 } 316 317 /// Walk all of the AffineExpr's in this mapping. Each node in an expression 318 /// tree is visited in postorder. 319 void AffineMap::walkExprs(std::function<void(AffineExpr)> callback) const { 320 for (auto expr : getResults()) 321 expr.walk(callback); 322 } 323 324 /// This method substitutes any uses of dimensions and symbols (e.g. 325 /// dim#0 with dimReplacements[0]) in subexpressions and returns the modified 326 /// expression mapping. Because this can be used to eliminate dims and 327 /// symbols, the client needs to specify the number of dims and symbols in 328 /// the result. The returned map always has the same number of results. 329 AffineMap AffineMap::replaceDimsAndSymbols(ArrayRef<AffineExpr> dimReplacements, 330 ArrayRef<AffineExpr> symReplacements, 331 unsigned numResultDims, 332 unsigned numResultSyms) const { 333 SmallVector<AffineExpr, 8> results; 334 results.reserve(getNumResults()); 335 for (auto expr : getResults()) 336 results.push_back( 337 expr.replaceDimsAndSymbols(dimReplacements, symReplacements)); 338 return get(numResultDims, numResultSyms, results, getContext()); 339 } 340 341 /// Sparse replace method. Apply AffineExpr::replace(`expr`, `replacement`) to 342 /// each of the results and return a new AffineMap with the new results and 343 /// with the specified number of dims and symbols. 344 AffineMap AffineMap::replace(AffineExpr expr, AffineExpr replacement, 345 unsigned numResultDims, 346 unsigned numResultSyms) const { 347 SmallVector<AffineExpr, 4> newResults; 348 newResults.reserve(getNumResults()); 349 for (AffineExpr e : getResults()) 350 newResults.push_back(e.replace(expr, replacement)); 351 return AffineMap::get(numResultDims, numResultSyms, newResults, getContext()); 352 } 353 354 /// Sparse replace method. Apply AffineExpr::replace(`map`) to each of the 355 /// results and return a new AffineMap with the new results and with the 356 /// specified number of dims and symbols. 357 AffineMap AffineMap::replace(const DenseMap<AffineExpr, AffineExpr> &map, 358 unsigned numResultDims, 359 unsigned numResultSyms) const { 360 SmallVector<AffineExpr, 4> newResults; 361 newResults.reserve(getNumResults()); 362 for (AffineExpr e : getResults()) 363 newResults.push_back(e.replace(map)); 364 return AffineMap::get(numResultDims, numResultSyms, newResults, getContext()); 365 } 366 367 AffineMap AffineMap::compose(AffineMap map) const { 368 assert(getNumDims() == map.getNumResults() && "Number of results mismatch"); 369 // Prepare `map` by concatenating the symbols and rewriting its exprs. 370 unsigned numDims = map.getNumDims(); 371 unsigned numSymbolsThisMap = getNumSymbols(); 372 unsigned numSymbols = numSymbolsThisMap + map.getNumSymbols(); 373 SmallVector<AffineExpr, 8> newDims(numDims); 374 for (unsigned idx = 0; idx < numDims; ++idx) { 375 newDims[idx] = getAffineDimExpr(idx, getContext()); 376 } 377 SmallVector<AffineExpr, 8> newSymbols(numSymbols - numSymbolsThisMap); 378 for (unsigned idx = numSymbolsThisMap; idx < numSymbols; ++idx) { 379 newSymbols[idx - numSymbolsThisMap] = 380 getAffineSymbolExpr(idx, getContext()); 381 } 382 auto newMap = 383 map.replaceDimsAndSymbols(newDims, newSymbols, numDims, numSymbols); 384 SmallVector<AffineExpr, 8> exprs; 385 exprs.reserve(getResults().size()); 386 for (auto expr : getResults()) 387 exprs.push_back(expr.compose(newMap)); 388 return AffineMap::get(numDims, numSymbols, exprs, map.getContext()); 389 } 390 391 SmallVector<int64_t, 4> AffineMap::compose(ArrayRef<int64_t> values) const { 392 assert(getNumSymbols() == 0 && "Expected symbol-less map"); 393 SmallVector<AffineExpr, 4> exprs; 394 exprs.reserve(values.size()); 395 MLIRContext *ctx = getContext(); 396 for (auto v : values) 397 exprs.push_back(getAffineConstantExpr(v, ctx)); 398 auto resMap = compose(AffineMap::get(0, 0, exprs, ctx)); 399 SmallVector<int64_t, 4> res; 400 res.reserve(resMap.getNumResults()); 401 for (auto e : resMap.getResults()) 402 res.push_back(e.cast<AffineConstantExpr>().getValue()); 403 return res; 404 } 405 406 bool AffineMap::isProjectedPermutation() const { 407 if (getNumSymbols() > 0) 408 return false; 409 SmallVector<bool, 8> seen(getNumInputs(), false); 410 for (auto expr : getResults()) { 411 if (auto dim = expr.dyn_cast<AffineDimExpr>()) { 412 if (seen[dim.getPosition()]) 413 return false; 414 seen[dim.getPosition()] = true; 415 continue; 416 } 417 return false; 418 } 419 return true; 420 } 421 422 bool AffineMap::isPermutation() const { 423 if (getNumDims() != getNumResults()) 424 return false; 425 return isProjectedPermutation(); 426 } 427 428 AffineMap AffineMap::getSubMap(ArrayRef<unsigned> resultPos) const { 429 SmallVector<AffineExpr, 4> exprs; 430 exprs.reserve(resultPos.size()); 431 for (auto idx : resultPos) 432 exprs.push_back(getResult(idx)); 433 return AffineMap::get(getNumDims(), getNumSymbols(), exprs, getContext()); 434 } 435 436 AffineMap AffineMap::getMajorSubMap(unsigned numResults) const { 437 if (numResults == 0) 438 return AffineMap(); 439 if (numResults > getNumResults()) 440 return *this; 441 return getSubMap(llvm::to_vector<4>(llvm::seq<unsigned>(0, numResults))); 442 } 443 444 AffineMap AffineMap::getMinorSubMap(unsigned numResults) const { 445 if (numResults == 0) 446 return AffineMap(); 447 if (numResults > getNumResults()) 448 return *this; 449 return getSubMap(llvm::to_vector<4>( 450 llvm::seq<unsigned>(getNumResults() - numResults, getNumResults()))); 451 } 452 453 AffineMap mlir::compressDims(AffineMap map, 454 const llvm::SmallDenseSet<unsigned> &unusedDims) { 455 unsigned numDims = 0; 456 SmallVector<AffineExpr> dimReplacements; 457 dimReplacements.reserve(map.getNumDims()); 458 MLIRContext *context = map.getContext(); 459 for (unsigned dim = 0, e = map.getNumDims(); dim < e; ++dim) { 460 if (unusedDims.contains(dim)) 461 dimReplacements.push_back(getAffineConstantExpr(0, context)); 462 else 463 dimReplacements.push_back(getAffineDimExpr(numDims++, context)); 464 } 465 SmallVector<AffineExpr> resultExprs; 466 resultExprs.reserve(map.getNumResults()); 467 for (auto e : map.getResults()) 468 resultExprs.push_back(e.replaceDims(dimReplacements)); 469 return AffineMap::get(numDims, map.getNumSymbols(), resultExprs, context); 470 } 471 472 AffineMap mlir::compressUnusedDims(AffineMap map) { 473 llvm::SmallDenseSet<unsigned> usedDims; 474 map.walkExprs([&](AffineExpr expr) { 475 if (auto dimExpr = expr.dyn_cast<AffineDimExpr>()) 476 usedDims.insert(dimExpr.getPosition()); 477 }); 478 llvm::SmallDenseSet<unsigned> unusedDims; 479 for (unsigned d = 0, e = map.getNumDims(); d != e; ++d) 480 if (!usedDims.contains(d)) 481 unusedDims.insert(d); 482 return compressDims(map, unusedDims); 483 } 484 485 AffineMap 486 mlir::compressSymbols(AffineMap map, 487 const llvm::SmallDenseSet<unsigned> &unusedSymbols) { 488 unsigned numSymbols = 0; 489 SmallVector<AffineExpr> symReplacements; 490 symReplacements.reserve(map.getNumSymbols()); 491 MLIRContext *context = map.getContext(); 492 for (unsigned sym = 0, e = map.getNumSymbols(); sym < e; ++sym) { 493 if (unusedSymbols.contains(sym)) 494 symReplacements.push_back(getAffineConstantExpr(0, context)); 495 else 496 symReplacements.push_back(getAffineSymbolExpr(numSymbols++, context)); 497 } 498 SmallVector<AffineExpr> resultExprs; 499 resultExprs.reserve(map.getNumResults()); 500 for (auto e : map.getResults()) 501 resultExprs.push_back(e.replaceSymbols(symReplacements)); 502 return AffineMap::get(map.getNumDims(), numSymbols, resultExprs, context); 503 } 504 505 AffineMap mlir::compressUnusedSymbols(AffineMap map) { 506 llvm::SmallDenseSet<unsigned> usedSymbols; 507 map.walkExprs([&](AffineExpr expr) { 508 if (auto symExpr = expr.dyn_cast<AffineSymbolExpr>()) 509 usedSymbols.insert(symExpr.getPosition()); 510 }); 511 llvm::SmallDenseSet<unsigned> unusedSymbols; 512 for (unsigned d = 0, e = map.getNumSymbols(); d != e; ++d) 513 if (!usedSymbols.contains(d)) 514 unusedSymbols.insert(d); 515 return compressSymbols(map, unusedSymbols); 516 } 517 518 AffineMap mlir::simplifyAffineMap(AffineMap map) { 519 SmallVector<AffineExpr, 8> exprs; 520 for (auto e : map.getResults()) { 521 exprs.push_back( 522 simplifyAffineExpr(e, map.getNumDims(), map.getNumSymbols())); 523 } 524 return AffineMap::get(map.getNumDims(), map.getNumSymbols(), exprs, 525 map.getContext()); 526 } 527 528 AffineMap mlir::removeDuplicateExprs(AffineMap map) { 529 auto results = map.getResults(); 530 SmallVector<AffineExpr, 4> uniqueExprs(results.begin(), results.end()); 531 uniqueExprs.erase(std::unique(uniqueExprs.begin(), uniqueExprs.end()), 532 uniqueExprs.end()); 533 return AffineMap::get(map.getNumDims(), map.getNumSymbols(), uniqueExprs, 534 map.getContext()); 535 } 536 537 AffineMap mlir::inversePermutation(AffineMap map) { 538 if (map.isEmpty()) 539 return map; 540 assert(map.getNumSymbols() == 0 && "expected map without symbols"); 541 SmallVector<AffineExpr, 4> exprs(map.getNumDims()); 542 for (auto en : llvm::enumerate(map.getResults())) { 543 auto expr = en.value(); 544 // Skip non-permutations. 545 if (auto d = expr.dyn_cast<AffineDimExpr>()) { 546 if (exprs[d.getPosition()]) 547 continue; 548 exprs[d.getPosition()] = getAffineDimExpr(en.index(), d.getContext()); 549 } 550 } 551 SmallVector<AffineExpr, 4> seenExprs; 552 seenExprs.reserve(map.getNumDims()); 553 for (auto expr : exprs) 554 if (expr) 555 seenExprs.push_back(expr); 556 if (seenExprs.size() != map.getNumInputs()) 557 return AffineMap(); 558 return AffineMap::get(map.getNumResults(), 0, seenExprs, map.getContext()); 559 } 560 561 AffineMap mlir::concatAffineMaps(ArrayRef<AffineMap> maps) { 562 unsigned numResults = 0, numDims = 0, numSymbols = 0; 563 for (auto m : maps) 564 numResults += m.getNumResults(); 565 SmallVector<AffineExpr, 8> results; 566 results.reserve(numResults); 567 for (auto m : maps) { 568 for (auto res : m.getResults()) 569 results.push_back(res.shiftSymbols(m.getNumSymbols(), numSymbols)); 570 571 numSymbols += m.getNumSymbols(); 572 numDims = std::max(m.getNumDims(), numDims); 573 } 574 return AffineMap::get(numDims, numSymbols, results, 575 maps.front().getContext()); 576 } 577 578 AffineMap 579 mlir::getProjectedMap(AffineMap map, 580 const llvm::SmallDenseSet<unsigned> &unusedDims) { 581 return compressUnusedSymbols(compressDims(map, unusedDims)); 582 } 583 584 //===----------------------------------------------------------------------===// 585 // MutableAffineMap. 586 //===----------------------------------------------------------------------===// 587 588 MutableAffineMap::MutableAffineMap(AffineMap map) 589 : numDims(map.getNumDims()), numSymbols(map.getNumSymbols()), 590 context(map.getContext()) { 591 for (auto result : map.getResults()) 592 results.push_back(result); 593 } 594 595 void MutableAffineMap::reset(AffineMap map) { 596 results.clear(); 597 numDims = map.getNumDims(); 598 numSymbols = map.getNumSymbols(); 599 context = map.getContext(); 600 for (auto result : map.getResults()) 601 results.push_back(result); 602 } 603 604 bool MutableAffineMap::isMultipleOf(unsigned idx, int64_t factor) const { 605 if (results[idx].isMultipleOf(factor)) 606 return true; 607 608 // TODO: use simplifyAffineExpr and FlatAffineConstraints to 609 // complete this (for a more powerful analysis). 610 return false; 611 } 612 613 // Simplifies the result affine expressions of this map. The expressions have to 614 // be pure for the simplification implemented. 615 void MutableAffineMap::simplify() { 616 // Simplify each of the results if possible. 617 // TODO: functional-style map 618 for (unsigned i = 0, e = getNumResults(); i < e; i++) { 619 results[i] = simplifyAffineExpr(getResult(i), numDims, numSymbols); 620 } 621 } 622 623 AffineMap MutableAffineMap::getAffineMap() const { 624 return AffineMap::get(numDims, numSymbols, results, context); 625 } 626