forked from OSchip/llvm-project
412 lines
14 KiB
C++
412 lines
14 KiB
C++
//===- AffineMap.cpp - MLIR Affine Map Classes ----------------------------===//
|
|
//
|
|
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
|
|
// See https://llvm.org/LICENSE.txt for license information.
|
|
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
|
|
//
|
|
//===----------------------------------------------------------------------===//
|
|
|
|
#include "mlir/IR/AffineMap.h"
|
|
#include "AffineMapDetail.h"
|
|
#include "mlir/IR/AffineExpr.h"
|
|
#include "mlir/IR/Attributes.h"
|
|
#include "mlir/IR/StandardTypes.h"
|
|
#include "mlir/Support/Functional.h"
|
|
#include "mlir/Support/LogicalResult.h"
|
|
#include "mlir/Support/MathExtras.h"
|
|
#include "llvm/ADT/StringRef.h"
|
|
#include "llvm/Support/raw_ostream.h"
|
|
|
|
using namespace mlir;
|
|
|
|
namespace {
|
|
|
|
// AffineExprConstantFolder evaluates an affine expression using constant
|
|
// operands passed in 'operandConsts'. Returns an IntegerAttr attribute
|
|
// representing the constant value of the affine expression evaluated on
|
|
// constant 'operandConsts', or nullptr if it can't be folded.
|
|
class AffineExprConstantFolder {
|
|
public:
|
|
AffineExprConstantFolder(unsigned numDims, ArrayRef<Attribute> operandConsts)
|
|
: numDims(numDims), operandConsts(operandConsts) {}
|
|
|
|
/// Attempt to constant fold the specified affine expr, or return null on
|
|
/// failure.
|
|
IntegerAttr constantFold(AffineExpr expr) {
|
|
if (auto result = constantFoldImpl(expr))
|
|
return IntegerAttr::get(IndexType::get(expr.getContext()), *result);
|
|
return nullptr;
|
|
}
|
|
|
|
private:
|
|
Optional<int64_t> constantFoldImpl(AffineExpr expr) {
|
|
switch (expr.getKind()) {
|
|
case AffineExprKind::Add:
|
|
return constantFoldBinExpr(
|
|
expr, [](int64_t lhs, int64_t rhs) { return lhs + rhs; });
|
|
case AffineExprKind::Mul:
|
|
return constantFoldBinExpr(
|
|
expr, [](int64_t lhs, int64_t rhs) { return lhs * rhs; });
|
|
case AffineExprKind::Mod:
|
|
return constantFoldBinExpr(
|
|
expr, [](int64_t lhs, int64_t rhs) { return mod(lhs, rhs); });
|
|
case AffineExprKind::FloorDiv:
|
|
return constantFoldBinExpr(
|
|
expr, [](int64_t lhs, int64_t rhs) { return floorDiv(lhs, rhs); });
|
|
case AffineExprKind::CeilDiv:
|
|
return constantFoldBinExpr(
|
|
expr, [](int64_t lhs, int64_t rhs) { return ceilDiv(lhs, rhs); });
|
|
case AffineExprKind::Constant:
|
|
return expr.cast<AffineConstantExpr>().getValue();
|
|
case AffineExprKind::DimId:
|
|
if (auto attr = operandConsts[expr.cast<AffineDimExpr>().getPosition()]
|
|
.dyn_cast_or_null<IntegerAttr>())
|
|
return attr.getInt();
|
|
return llvm::None;
|
|
case AffineExprKind::SymbolId:
|
|
if (auto attr = operandConsts[numDims +
|
|
expr.cast<AffineSymbolExpr>().getPosition()]
|
|
.dyn_cast_or_null<IntegerAttr>())
|
|
return attr.getInt();
|
|
return llvm::None;
|
|
}
|
|
llvm_unreachable("Unknown AffineExpr");
|
|
}
|
|
|
|
// TODO: Change these to operate on APInts too.
|
|
Optional<int64_t> constantFoldBinExpr(AffineExpr expr,
|
|
int64_t (*op)(int64_t, int64_t)) {
|
|
auto binOpExpr = expr.cast<AffineBinaryOpExpr>();
|
|
if (auto lhs = constantFoldImpl(binOpExpr.getLHS()))
|
|
if (auto rhs = constantFoldImpl(binOpExpr.getRHS()))
|
|
return op(*lhs, *rhs);
|
|
return llvm::None;
|
|
}
|
|
|
|
// The number of dimension operands in AffineMap containing this expression.
|
|
unsigned numDims;
|
|
// The constant valued operands used to evaluate this AffineExpr.
|
|
ArrayRef<Attribute> operandConsts;
|
|
};
|
|
|
|
} // end anonymous namespace
|
|
|
|
/// Returns a single constant result affine map.
|
|
AffineMap AffineMap::getConstantMap(int64_t val, MLIRContext *context) {
|
|
return get(/*dimCount=*/0, /*symbolCount=*/0,
|
|
{getAffineConstantExpr(val, context)});
|
|
}
|
|
|
|
/// Returns an AffineMap representing a permutation.
|
|
AffineMap AffineMap::getPermutationMap(ArrayRef<unsigned> permutation,
|
|
MLIRContext *context) {
|
|
assert(!permutation.empty() &&
|
|
"Cannot create permutation map from empty permutation vector");
|
|
SmallVector<AffineExpr, 4> affExprs;
|
|
for (auto index : permutation)
|
|
affExprs.push_back(getAffineDimExpr(index, context));
|
|
auto m = std::max_element(permutation.begin(), permutation.end());
|
|
auto permutationMap = AffineMap::get(*m + 1, 0, affExprs);
|
|
assert(permutationMap.isPermutation() && "Invalid permutation vector");
|
|
return permutationMap;
|
|
}
|
|
|
|
template <typename AffineExprContainer>
|
|
static void getMaxDimAndSymbol(ArrayRef<AffineExprContainer> exprsList,
|
|
int64_t &maxDim, int64_t &maxSym) {
|
|
for (const auto &exprs : exprsList) {
|
|
for (auto expr : exprs) {
|
|
expr.walk([&maxDim, &maxSym](AffineExpr e) {
|
|
if (auto d = e.dyn_cast<AffineDimExpr>())
|
|
maxDim = std::max(maxDim, static_cast<int64_t>(d.getPosition()));
|
|
if (auto s = e.dyn_cast<AffineSymbolExpr>())
|
|
maxSym = std::max(maxSym, static_cast<int64_t>(s.getPosition()));
|
|
});
|
|
}
|
|
}
|
|
}
|
|
|
|
template <typename AffineExprContainer>
|
|
static SmallVector<AffineMap, 4>
|
|
inferFromExprList(ArrayRef<AffineExprContainer> exprsList) {
|
|
int64_t maxDim = -1, maxSym = -1;
|
|
getMaxDimAndSymbol(exprsList, maxDim, maxSym);
|
|
SmallVector<AffineMap, 4> maps;
|
|
maps.reserve(exprsList.size());
|
|
for (const auto &exprs : exprsList)
|
|
maps.push_back(AffineMap::get(/*dimCount=*/maxDim + 1,
|
|
/*symbolCount=*/maxSym + 1, exprs));
|
|
return maps;
|
|
}
|
|
|
|
SmallVector<AffineMap, 4>
|
|
AffineMap::inferFromExprList(ArrayRef<ArrayRef<AffineExpr>> exprsList) {
|
|
return ::inferFromExprList(exprsList);
|
|
}
|
|
|
|
SmallVector<AffineMap, 4>
|
|
AffineMap::inferFromExprList(ArrayRef<SmallVector<AffineExpr, 4>> exprsList) {
|
|
return ::inferFromExprList(exprsList);
|
|
}
|
|
|
|
AffineMap AffineMap::getMultiDimIdentityMap(unsigned numDims,
|
|
MLIRContext *context) {
|
|
SmallVector<AffineExpr, 4> dimExprs;
|
|
dimExprs.reserve(numDims);
|
|
for (unsigned i = 0; i < numDims; ++i)
|
|
dimExprs.push_back(mlir::getAffineDimExpr(i, context));
|
|
return get(/*dimCount=*/numDims, /*symbolCount=*/0, dimExprs);
|
|
}
|
|
|
|
MLIRContext *AffineMap::getContext() const { return map->context; }
|
|
|
|
bool AffineMap::isIdentity() const {
|
|
if (getNumDims() != getNumResults())
|
|
return false;
|
|
ArrayRef<AffineExpr> results = getResults();
|
|
for (unsigned i = 0, numDims = getNumDims(); i < numDims; ++i) {
|
|
auto expr = results[i].dyn_cast<AffineDimExpr>();
|
|
if (!expr || expr.getPosition() != i)
|
|
return false;
|
|
}
|
|
return true;
|
|
}
|
|
|
|
bool AffineMap::isEmpty() const {
|
|
return getNumDims() == 0 && getNumSymbols() == 0 && getNumResults() == 0;
|
|
}
|
|
|
|
bool AffineMap::isSingleConstant() const {
|
|
return getNumResults() == 1 && getResult(0).isa<AffineConstantExpr>();
|
|
}
|
|
|
|
int64_t AffineMap::getSingleConstantResult() const {
|
|
assert(isSingleConstant() && "map must have a single constant result");
|
|
return getResult(0).cast<AffineConstantExpr>().getValue();
|
|
}
|
|
|
|
unsigned AffineMap::getNumDims() const {
|
|
assert(map && "uninitialized map storage");
|
|
return map->numDims;
|
|
}
|
|
unsigned AffineMap::getNumSymbols() const {
|
|
assert(map && "uninitialized map storage");
|
|
return map->numSymbols;
|
|
}
|
|
unsigned AffineMap::getNumResults() const {
|
|
assert(map && "uninitialized map storage");
|
|
return map->results.size();
|
|
}
|
|
unsigned AffineMap::getNumInputs() const {
|
|
assert(map && "uninitialized map storage");
|
|
return map->numDims + map->numSymbols;
|
|
}
|
|
|
|
ArrayRef<AffineExpr> AffineMap::getResults() const {
|
|
assert(map && "uninitialized map storage");
|
|
return map->results;
|
|
}
|
|
AffineExpr AffineMap::getResult(unsigned idx) const {
|
|
assert(map && "uninitialized map storage");
|
|
return map->results[idx];
|
|
}
|
|
|
|
/// Folds the results of the application of an affine map on the provided
|
|
/// operands to a constant if possible. Returns false if the folding happens,
|
|
/// true otherwise.
|
|
LogicalResult
|
|
AffineMap::constantFold(ArrayRef<Attribute> operandConstants,
|
|
SmallVectorImpl<Attribute> &results) const {
|
|
assert(getNumInputs() == operandConstants.size());
|
|
|
|
// Fold each of the result expressions.
|
|
AffineExprConstantFolder exprFolder(getNumDims(), operandConstants);
|
|
// Constant fold each AffineExpr in AffineMap and add to 'results'.
|
|
for (auto expr : getResults()) {
|
|
auto folded = exprFolder.constantFold(expr);
|
|
// If we didn't fold to a constant, then folding fails.
|
|
if (!folded)
|
|
return failure();
|
|
|
|
results.push_back(folded);
|
|
}
|
|
assert(results.size() == getNumResults() &&
|
|
"constant folding produced the wrong number of results");
|
|
return success();
|
|
}
|
|
|
|
/// Walk all of the AffineExpr's in this mapping. Each node in an expression
|
|
/// tree is visited in postorder.
|
|
void AffineMap::walkExprs(std::function<void(AffineExpr)> callback) const {
|
|
for (auto expr : getResults())
|
|
expr.walk(callback);
|
|
}
|
|
|
|
/// This method substitutes any uses of dimensions and symbols (e.g.
|
|
/// dim#0 with dimReplacements[0]) in subexpressions and returns the modified
|
|
/// expression mapping. Because this can be used to eliminate dims and
|
|
/// symbols, the client needs to specify the number of dims and symbols in
|
|
/// the result. The returned map always has the same number of results.
|
|
AffineMap AffineMap::replaceDimsAndSymbols(ArrayRef<AffineExpr> dimReplacements,
|
|
ArrayRef<AffineExpr> symReplacements,
|
|
unsigned numResultDims,
|
|
unsigned numResultSyms) {
|
|
SmallVector<AffineExpr, 8> results;
|
|
results.reserve(getNumResults());
|
|
for (auto expr : getResults())
|
|
results.push_back(
|
|
expr.replaceDimsAndSymbols(dimReplacements, symReplacements));
|
|
|
|
return get(numResultDims, numResultSyms, results);
|
|
}
|
|
|
|
AffineMap AffineMap::compose(AffineMap map) {
|
|
assert(getNumDims() == map.getNumResults() && "Number of results mismatch");
|
|
// Prepare `map` by concatenating the symbols and rewriting its exprs.
|
|
unsigned numDims = map.getNumDims();
|
|
unsigned numSymbolsThisMap = getNumSymbols();
|
|
unsigned numSymbols = numSymbolsThisMap + map.getNumSymbols();
|
|
SmallVector<AffineExpr, 8> newDims(numDims);
|
|
for (unsigned idx = 0; idx < numDims; ++idx) {
|
|
newDims[idx] = getAffineDimExpr(idx, getContext());
|
|
}
|
|
SmallVector<AffineExpr, 8> newSymbols(numSymbols);
|
|
for (unsigned idx = numSymbolsThisMap; idx < numSymbols; ++idx) {
|
|
newSymbols[idx - numSymbolsThisMap] =
|
|
getAffineSymbolExpr(idx, getContext());
|
|
}
|
|
auto newMap =
|
|
map.replaceDimsAndSymbols(newDims, newSymbols, numDims, numSymbols);
|
|
SmallVector<AffineExpr, 8> exprs;
|
|
exprs.reserve(getResults().size());
|
|
for (auto expr : getResults())
|
|
exprs.push_back(expr.compose(newMap));
|
|
return AffineMap::get(numDims, numSymbols, exprs);
|
|
}
|
|
|
|
bool AffineMap::isProjectedPermutation() {
|
|
if (getNumSymbols() > 0)
|
|
return false;
|
|
SmallVector<bool, 8> seen(getNumInputs(), false);
|
|
for (auto expr : getResults()) {
|
|
if (auto dim = expr.dyn_cast<AffineDimExpr>()) {
|
|
if (seen[dim.getPosition()])
|
|
return false;
|
|
seen[dim.getPosition()] = true;
|
|
continue;
|
|
}
|
|
return false;
|
|
}
|
|
return true;
|
|
}
|
|
|
|
bool AffineMap::isPermutation() {
|
|
if (getNumDims() != getNumResults())
|
|
return false;
|
|
return isProjectedPermutation();
|
|
}
|
|
|
|
AffineMap AffineMap::getSubMap(ArrayRef<unsigned> resultPos) {
|
|
SmallVector<AffineExpr, 4> exprs;
|
|
exprs.reserve(resultPos.size());
|
|
for (auto idx : resultPos) {
|
|
exprs.push_back(getResult(idx));
|
|
}
|
|
return AffineMap::get(getNumDims(), getNumSymbols(), exprs);
|
|
}
|
|
|
|
AffineMap mlir::simplifyAffineMap(AffineMap map) {
|
|
SmallVector<AffineExpr, 8> exprs;
|
|
for (auto e : map.getResults()) {
|
|
exprs.push_back(
|
|
simplifyAffineExpr(e, map.getNumDims(), map.getNumSymbols()));
|
|
}
|
|
return AffineMap::get(map.getNumDims(), map.getNumSymbols(), exprs);
|
|
}
|
|
|
|
AffineMap mlir::inversePermutation(AffineMap map) {
|
|
if (!map)
|
|
return map;
|
|
assert(map.getNumSymbols() == 0 && "expected map without symbols");
|
|
SmallVector<AffineExpr, 4> exprs(map.getNumDims());
|
|
for (auto en : llvm::enumerate(map.getResults())) {
|
|
auto expr = en.value();
|
|
// Skip non-permutations.
|
|
if (auto d = expr.dyn_cast<AffineDimExpr>()) {
|
|
if (exprs[d.getPosition()])
|
|
continue;
|
|
exprs[d.getPosition()] = getAffineDimExpr(en.index(), d.getContext());
|
|
}
|
|
}
|
|
SmallVector<AffineExpr, 4> seenExprs;
|
|
seenExprs.reserve(map.getNumDims());
|
|
for (auto expr : exprs)
|
|
if (expr)
|
|
seenExprs.push_back(expr);
|
|
if (seenExprs.size() != map.getNumInputs())
|
|
return AffineMap();
|
|
return AffineMap::get(map.getNumResults(), 0, seenExprs);
|
|
}
|
|
|
|
AffineMap mlir::concatAffineMaps(ArrayRef<AffineMap> maps) {
|
|
unsigned numResults = 0;
|
|
for (auto m : maps)
|
|
numResults += (m && !m.isSingleConstant()) ? m.getNumResults() : 0;
|
|
unsigned numDims = 0;
|
|
SmallVector<AffineExpr, 8> results;
|
|
results.reserve(numResults);
|
|
for (auto m : maps) {
|
|
if (!m || m.isSingleConstant())
|
|
continue;
|
|
assert(m.getNumSymbols() == 0 && "expected map without symbols");
|
|
results.append(m.getResults().begin(), m.getResults().end());
|
|
numDims = std::max(m.getNumDims(), numDims);
|
|
}
|
|
return numDims == 0 ? AffineMap() : AffineMap::get(numDims, 0, results);
|
|
}
|
|
|
|
//===----------------------------------------------------------------------===//
|
|
// MutableAffineMap.
|
|
//===----------------------------------------------------------------------===//
|
|
|
|
MutableAffineMap::MutableAffineMap(AffineMap map)
|
|
: numDims(map.getNumDims()), numSymbols(map.getNumSymbols()),
|
|
// A map always has at least 1 result by construction
|
|
context(map.getResult(0).getContext()) {
|
|
for (auto result : map.getResults())
|
|
results.push_back(result);
|
|
}
|
|
|
|
void MutableAffineMap::reset(AffineMap map) {
|
|
results.clear();
|
|
numDims = map.getNumDims();
|
|
numSymbols = map.getNumSymbols();
|
|
// A map always has at least 1 result by construction
|
|
context = map.getResult(0).getContext();
|
|
for (auto result : map.getResults())
|
|
results.push_back(result);
|
|
}
|
|
|
|
bool MutableAffineMap::isMultipleOf(unsigned idx, int64_t factor) const {
|
|
if (results[idx].isMultipleOf(factor))
|
|
return true;
|
|
|
|
// TODO(bondhugula): use simplifyAffineExpr and FlatAffineConstraints to
|
|
// complete this (for a more powerful analysis).
|
|
return false;
|
|
}
|
|
|
|
// Simplifies the result affine expressions of this map. The expressions have to
|
|
// be pure for the simplification implemented.
|
|
void MutableAffineMap::simplify() {
|
|
// Simplify each of the results if possible.
|
|
// TODO(ntv): functional-style map
|
|
for (unsigned i = 0, e = getNumResults(); i < e; i++) {
|
|
results[i] = simplifyAffineExpr(getResult(i), numDims, numSymbols);
|
|
}
|
|
}
|
|
|
|
AffineMap MutableAffineMap::getAffineMap() const {
|
|
return AffineMap::get(numDims, numSymbols, results);
|
|
}
|