From 218bf2aa82ca6a177e6aca66d9b6327bd0383f49 Mon Sep 17 00:00:00 2001 From: Matthias Springer Date: Wed, 25 Oct 2023 11:16:20 +0900 Subject: [PATCH] [mlir][NFC] Move `foldAttributesIntoMap` to `IR` build unit `foldAttributesIntoMap` is a helper function that folds constant `OpFoldResult` into an affine map. This commit moves the function from the affine dialect to `AffineMap.h`, so that it can be used without depending on the affine dialect. --- mlir/include/mlir/IR/AffineMap.h | 9 ++++++ mlir/lib/Dialect/Affine/IR/AffineOps.cpp | 38 +++--------------------- mlir/lib/IR/AffineMap.cpp | 29 ++++++++++++++++++ 3 files changed, 42 insertions(+), 34 deletions(-) diff --git a/mlir/include/mlir/IR/AffineMap.h b/mlir/include/mlir/IR/AffineMap.h index 3430db2b99c3f..5af7835258f6b 100644 --- a/mlir/include/mlir/IR/AffineMap.h +++ b/mlir/include/mlir/IR/AffineMap.h @@ -15,6 +15,7 @@ #define MLIR_IR_AFFINEMAP_H #include "mlir/IR/AffineExpr.h" +#include "mlir/IR/Value.h" #include "mlir/Support/LLVM.h" #include "llvm/ADT/ArrayRef.h" #include "llvm/ADT/DenseMapInfo.h" @@ -33,7 +34,9 @@ struct AffineMapStorage; } // namespace detail class Attribute; +class Builder; struct LogicalResult; +class OpFoldResult; class MLIRContext; /// A multi-dimensional affine map @@ -447,6 +450,12 @@ AffineMap compressUnusedSymbols(AffineMap map); /// dims and symbols. SmallVector compressUnusedSymbols(ArrayRef maps); +/// Fold all attributes among the given operands into the affine map. Return the +/// folded affine map. Return all remaining values via `remainingValues`. +AffineMap foldAttributesIntoMap(Builder &b, AffineMap map, + ArrayRef operands, + SmallVector &remainingValues); + /// Returns a map with the same dimension and symbol count as `map`, but whose /// results are the unique affine expressions of `map`. AffineMap removeDuplicateExprs(AffineMap map); diff --git a/mlir/lib/Dialect/Affine/IR/AffineOps.cpp b/mlir/lib/Dialect/Affine/IR/AffineOps.cpp index 85d16088c43fb..4d79c458889d2 100644 --- a/mlir/lib/Dialect/Affine/IR/AffineOps.cpp +++ b/mlir/lib/Dialect/Affine/IR/AffineOps.cpp @@ -1193,42 +1193,11 @@ void mlir::affine::fullyComposeAffineMapAndOperands( } } -/// Fold all attributes among the given operands into the affine map and return -/// all remaining values. The affine map is modified in-place. -static SmallVector -foldAttributesIntoMap(Builder &b, AffineMap *map, - ArrayRef operands) { - SmallVector dimReplacements, symReplacements; - SmallVector valueOperands; - int64_t numDims = 0; - for (int64_t i = 0; i < map->getNumDims(); ++i) { - if (auto attr = operands[i].dyn_cast()) { - dimReplacements.push_back( - b.getAffineConstantExpr(attr.cast().getInt())); - } else { - dimReplacements.push_back(b.getAffineDimExpr(numDims++)); - valueOperands.push_back(operands[i].get()); - } - } - int64_t numSymbols = 0; - for (int64_t i = 0; i < map->getNumSymbols(); ++i) { - if (auto attr = operands[i + map->getNumDims()].dyn_cast()) { - symReplacements.push_back( - b.getAffineConstantExpr(attr.cast().getInt())); - } else { - symReplacements.push_back(b.getAffineSymbolExpr(numSymbols++)); - valueOperands.push_back(operands[i + map->getNumDims()].get()); - } - } - *map = map->replaceDimsAndSymbols(dimReplacements, symReplacements, numDims, - numSymbols); - return valueOperands; -} - AffineApplyOp mlir::affine::makeComposedAffineApply(OpBuilder &b, Location loc, AffineMap map, ArrayRef operands) { - SmallVector valueOperands = foldAttributesIntoMap(b, &map, operands); + SmallVector valueOperands; + map = foldAttributesIntoMap(b, map, operands, valueOperands); composeAffineMapAndOperands(&map, &valueOperands); assert(map); return b.create(loc, map, valueOperands); @@ -1331,7 +1300,8 @@ mlir::affine::makeComposedFoldedMultiResultAffineApply( template static OpTy makeComposedMinMax(OpBuilder &b, Location loc, AffineMap map, ArrayRef operands) { - SmallVector valueOperands = foldAttributesIntoMap(b, &map, operands); + SmallVector valueOperands; + map = foldAttributesIntoMap(b, map, operands, valueOperands); composeMultiResultAffineMap(map, valueOperands); return b.create(loc, b.getIndexType(), map, valueOperands); } diff --git a/mlir/lib/IR/AffineMap.cpp b/mlir/lib/IR/AffineMap.cpp index 9cdac964710ca..3bd1181b6c7bb 100644 --- a/mlir/lib/IR/AffineMap.cpp +++ b/mlir/lib/IR/AffineMap.cpp @@ -9,6 +9,7 @@ #include "mlir/IR/AffineMap.h" #include "AffineMapDetail.h" #include "mlir/IR/AffineExpr.h" +#include "mlir/IR/Builders.h" #include "mlir/IR/BuiltinAttributes.h" #include "mlir/IR/BuiltinTypes.h" #include "mlir/Support/LogicalResult.h" @@ -658,6 +659,34 @@ SmallVector mlir::compressUnusedSymbols(ArrayRef maps) { maps, [](AffineMap m) { return compressUnusedSymbols(m); }); } +AffineMap mlir::foldAttributesIntoMap(Builder &b, AffineMap map, + ArrayRef operands, + SmallVector &remainingValues) { + SmallVector dimReplacements, symReplacements; + int64_t numDims = 0; + for (int64_t i = 0; i < map.getNumDims(); ++i) { + if (auto attr = operands[i].dyn_cast()) { + dimReplacements.push_back( + b.getAffineConstantExpr(attr.cast().getInt())); + } else { + dimReplacements.push_back(b.getAffineDimExpr(numDims++)); + remainingValues.push_back(operands[i].get()); + } + } + int64_t numSymbols = 0; + for (int64_t i = 0; i < map.getNumSymbols(); ++i) { + if (auto attr = operands[i + map.getNumDims()].dyn_cast()) { + symReplacements.push_back( + b.getAffineConstantExpr(attr.cast().getInt())); + } else { + symReplacements.push_back(b.getAffineSymbolExpr(numSymbols++)); + remainingValues.push_back(operands[i + map.getNumDims()].get()); + } + } + return map.replaceDimsAndSymbols(dimReplacements, symReplacements, numDims, + numSymbols); +} + AffineMap mlir::simplifyAffineMap(AffineMap map) { SmallVector exprs; for (auto e : map.getResults()) {