2018-06-28 02:03:08 +08:00
|
|
|
//===- AffineMap.cpp - MLIR Affine Map Classes ----------------------------===//
|
|
|
|
//
|
|
|
|
// Copyright 2019 The MLIR Authors.
|
|
|
|
//
|
|
|
|
// Licensed under the Apache License, Version 2.0 (the "License");
|
|
|
|
// you may not use this file except in compliance with the License.
|
|
|
|
// You may obtain a copy of the License at
|
|
|
|
//
|
|
|
|
// http://www.apache.org/licenses/LICENSE-2.0
|
|
|
|
//
|
|
|
|
// Unless required by applicable law or agreed to in writing, software
|
|
|
|
// distributed under the License is distributed on an "AS IS" BASIS,
|
|
|
|
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
|
|
// See the License for the specific language governing permissions and
|
|
|
|
// limitations under the License.
|
|
|
|
// =============================================================================
|
|
|
|
|
|
|
|
#include "mlir/IR/AffineMap.h"
|
2018-10-10 07:39:24 +08:00
|
|
|
#include "AffineMapDetail.h"
|
2018-07-10 00:00:25 +08:00
|
|
|
#include "mlir/IR/AffineExpr.h"
|
2018-10-06 09:24:18 +08:00
|
|
|
#include "mlir/IR/Attributes.h"
|
2019-02-13 03:08:04 +08:00
|
|
|
#include "mlir/IR/StandardTypes.h"
|
2019-01-08 12:05:14 +08:00
|
|
|
#include "mlir/Support/Functional.h"
|
2019-03-11 08:44:25 +08:00
|
|
|
#include "mlir/Support/LogicalResult.h"
|
2018-10-04 01:07:54 +08:00
|
|
|
#include "mlir/Support/MathExtras.h"
|
2018-06-28 02:03:08 +08:00
|
|
|
#include "llvm/ADT/StringRef.h"
|
2019-06-22 09:31:52 +08:00
|
|
|
#include "llvm/Support/raw_ostream.h"
|
2018-06-28 02:03:08 +08:00
|
|
|
|
|
|
|
using namespace mlir;
|
|
|
|
|
2018-10-06 09:24:18 +08:00
|
|
|
namespace {
|
|
|
|
|
|
|
|
// AffineExprConstantFolder evaluates an affine expression using constant
|
2019-03-15 02:34:12 +08:00
|
|
|
// 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.
|
2018-10-06 09:24:18 +08:00
|
|
|
class AffineExprConstantFolder {
|
|
|
|
public:
|
2018-10-26 06:46:10 +08:00
|
|
|
AffineExprConstantFolder(unsigned numDims, ArrayRef<Attribute> operandConsts)
|
2018-10-06 09:24:18 +08:00
|
|
|
: numDims(numDims), operandConsts(operandConsts) {}
|
|
|
|
|
|
|
|
/// Attempt to constant fold the specified affine expr, or return null on
|
|
|
|
/// failure.
|
2018-10-26 06:46:10 +08:00
|
|
|
IntegerAttr constantFold(AffineExpr expr) {
|
2019-03-15 02:34:12 +08:00
|
|
|
if (auto result = constantFoldImpl(expr))
|
|
|
|
return IntegerAttr::get(IndexType::get(expr.getContext()), *result);
|
|
|
|
return nullptr;
|
|
|
|
}
|
|
|
|
|
|
|
|
private:
|
|
|
|
llvm::Optional<int64_t> constantFoldImpl(AffineExpr expr) {
|
2018-10-10 01:59:27 +08:00
|
|
|
switch (expr.getKind()) {
|
2018-10-09 01:20:25 +08:00
|
|
|
case AffineExprKind::Add:
|
2018-10-06 09:24:18 +08:00
|
|
|
return constantFoldBinExpr(
|
|
|
|
expr, [](int64_t lhs, int64_t rhs) { return lhs + rhs; });
|
2018-10-09 01:20:25 +08:00
|
|
|
case AffineExprKind::Mul:
|
2018-10-06 09:24:18 +08:00
|
|
|
return constantFoldBinExpr(
|
|
|
|
expr, [](int64_t lhs, int64_t rhs) { return lhs * rhs; });
|
2018-10-09 01:20:25 +08:00
|
|
|
case AffineExprKind::Mod:
|
2018-10-06 09:24:18 +08:00
|
|
|
return constantFoldBinExpr(
|
2019-03-15 02:34:12 +08:00
|
|
|
expr, [](int64_t lhs, int64_t rhs) { return mod(lhs, rhs); });
|
2018-10-09 01:20:25 +08:00
|
|
|
case AffineExprKind::FloorDiv:
|
2018-10-06 09:24:18 +08:00
|
|
|
return constantFoldBinExpr(
|
2019-03-15 02:34:12 +08:00
|
|
|
expr, [](int64_t lhs, int64_t rhs) { return floorDiv(lhs, rhs); });
|
2018-10-09 01:20:25 +08:00
|
|
|
case AffineExprKind::CeilDiv:
|
2018-10-06 09:24:18 +08:00
|
|
|
return constantFoldBinExpr(
|
2019-03-15 02:34:12 +08:00
|
|
|
expr, [](int64_t lhs, int64_t rhs) { return ceilDiv(lhs, rhs); });
|
2018-10-09 01:20:25 +08:00
|
|
|
case AffineExprKind::Constant:
|
2019-03-15 02:34:12 +08:00
|
|
|
return expr.cast<AffineConstantExpr>().getValue();
|
2018-10-09 01:20:25 +08:00
|
|
|
case AffineExprKind::DimId:
|
2019-03-15 02:34:12 +08:00
|
|
|
if (auto attr = operandConsts[expr.cast<AffineDimExpr>().getPosition()]
|
|
|
|
.dyn_cast_or_null<IntegerAttr>())
|
|
|
|
return attr.getInt();
|
|
|
|
return llvm::None;
|
2018-10-09 01:20:25 +08:00
|
|
|
case AffineExprKind::SymbolId:
|
2019-03-15 02:34:12 +08:00
|
|
|
if (auto attr = operandConsts[numDims +
|
|
|
|
expr.cast<AffineSymbolExpr>().getPosition()]
|
|
|
|
.dyn_cast_or_null<IntegerAttr>())
|
|
|
|
return attr.getInt();
|
|
|
|
return llvm::None;
|
2018-10-06 09:24:18 +08:00
|
|
|
}
|
2019-05-11 05:06:10 +08:00
|
|
|
llvm_unreachable("Unknown AffineExpr");
|
2018-10-06 09:24:18 +08:00
|
|
|
}
|
|
|
|
|
2018-11-12 22:33:22 +08:00
|
|
|
// TODO: Change these to operate on APInts too.
|
2019-03-15 02:34:12 +08:00
|
|
|
llvm::Optional<int64_t> constantFoldBinExpr(AffineExpr expr,
|
|
|
|
int64_t (*op)(int64_t, int64_t)) {
|
2018-10-09 04:47:18 +08:00
|
|
|
auto binOpExpr = expr.cast<AffineBinaryOpExpr>();
|
2019-03-15 02:34:12 +08:00
|
|
|
if (auto lhs = constantFoldImpl(binOpExpr.getLHS()))
|
|
|
|
if (auto rhs = constantFoldImpl(binOpExpr.getRHS()))
|
|
|
|
return op(*lhs, *rhs);
|
|
|
|
return llvm::None;
|
2018-10-06 09:24:18 +08:00
|
|
|
}
|
|
|
|
|
|
|
|
// The number of dimension operands in AffineMap containing this expression.
|
|
|
|
unsigned numDims;
|
|
|
|
// The constant valued operands used to evaluate this AffineExpr.
|
2018-10-26 06:46:10 +08:00
|
|
|
ArrayRef<Attribute> operandConsts;
|
2018-10-06 09:24:18 +08:00
|
|
|
};
|
|
|
|
|
|
|
|
} // end anonymous namespace
|
|
|
|
|
2018-10-10 07:39:24 +08:00
|
|
|
/// Returns a single constant result affine map.
|
|
|
|
AffineMap AffineMap::getConstantMap(int64_t val, MLIRContext *context) {
|
|
|
|
return get(/*dimCount=*/0, /*symbolCount=*/0,
|
2019-05-30 05:56:41 +08:00
|
|
|
{getAffineConstantExpr(val, context)});
|
2018-10-10 07:39:24 +08:00
|
|
|
}
|
2018-07-10 00:00:25 +08:00
|
|
|
|
2018-11-01 22:26:00 +08:00
|
|
|
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));
|
2019-05-30 05:56:41 +08:00
|
|
|
return get(/*dimCount=*/numDims, /*symbolCount=*/0, dimExprs);
|
2018-11-01 22:26:00 +08:00
|
|
|
}
|
|
|
|
|
2018-10-23 02:47:10 +08:00
|
|
|
MLIRContext *AffineMap::getContext() const { return getResult(0).getContext(); }
|
|
|
|
|
2018-10-10 07:39:24 +08:00
|
|
|
bool AffineMap::isIdentity() const {
|
2018-08-15 03:43:51 +08:00
|
|
|
if (getNumDims() != getNumResults())
|
|
|
|
return false;
|
2018-10-09 04:47:18 +08:00
|
|
|
ArrayRef<AffineExpr> results = getResults();
|
2018-08-16 06:14:45 +08:00
|
|
|
for (unsigned i = 0, numDims = getNumDims(); i < numDims; ++i) {
|
2018-10-09 04:47:18 +08:00
|
|
|
auto expr = results[i].dyn_cast<AffineDimExpr>();
|
2018-10-10 01:59:27 +08:00
|
|
|
if (!expr || expr.getPosition() != i)
|
2018-08-15 03:43:51 +08:00
|
|
|
return false;
|
|
|
|
}
|
|
|
|
return true;
|
|
|
|
}
|
|
|
|
|
2018-10-10 07:39:24 +08:00
|
|
|
bool AffineMap::isSingleConstant() const {
|
2018-10-09 04:47:18 +08:00
|
|
|
return getNumResults() == 1 && getResult(0).isa<AffineConstantExpr>();
|
2018-08-25 14:38:14 +08:00
|
|
|
}
|
|
|
|
|
2018-10-10 07:39:24 +08:00
|
|
|
int64_t AffineMap::getSingleConstantResult() const {
|
2018-08-25 14:38:14 +08:00
|
|
|
assert(isSingleConstant() && "map must have a single constant result");
|
2018-10-10 01:59:27 +08:00
|
|
|
return getResult(0).cast<AffineConstantExpr>().getValue();
|
2018-08-25 14:38:14 +08:00
|
|
|
}
|
|
|
|
|
2019-01-12 08:08:16 +08:00
|
|
|
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();
|
|
|
|
}
|
2018-10-10 07:39:24 +08:00
|
|
|
unsigned AffineMap::getNumInputs() const {
|
2019-01-12 08:08:16 +08:00
|
|
|
assert(map && "uninitialized map storage");
|
2018-10-10 07:39:24 +08:00
|
|
|
return map->numDims + map->numSymbols;
|
|
|
|
}
|
|
|
|
|
2019-01-12 08:08:16 +08:00
|
|
|
ArrayRef<AffineExpr> AffineMap::getResults() const {
|
|
|
|
assert(map && "uninitialized map storage");
|
|
|
|
return map->results;
|
|
|
|
}
|
2018-10-10 07:39:24 +08:00
|
|
|
AffineExpr AffineMap::getResult(unsigned idx) const {
|
2019-01-12 08:08:16 +08:00
|
|
|
assert(map && "uninitialized map storage");
|
2018-10-10 07:39:24 +08:00
|
|
|
return map->results[idx];
|
|
|
|
}
|
[RFC][MLIR] Use AffineExprRef in place of AffineExpr* in IR
This CL starts by replacing AffineExpr* with value-type AffineExprRef in a few
places in the IR. By a domino effect that is pretty telling of the
inconsistencies in the codebase, const is removed where it makes sense.
The rationale is that the decision was concisously made that unique'd types
have pointer semantics without const specifier. This is fine but we should be
consistent. In the end, the only logical invariant is that there should never
be such a thing as a const AffineExpr*, const AffineMap* or const IntegerSet*
in our codebase.
This CL takes a number of shortcuts to killing const with fire, in particular
forcing const AffineExprRef to return the underlying non-const
AffineExpr*. This will be removed once AffineExpr* has disappeared in
containers but for now such shortcuts allow a bit of sanity in this long quest
for cleanups.
The **only** places where const AffineExpr*, const AffineMap* or const
IntegerSet* may still appear is by transitive needs from containers,
comparison operators etc.
There is still one major thing remaining here: figure out why cast/dyn_cast
return me a const AffineXXX*, which in turn requires a bunch of ugly
const_casts. I suspect this is due to the classof
taking const AffineXXXExpr*. I wonder whether this is a side effect of 1., if
it is coming from llvm itself (I'd doubt it) or something else (clattner@?)
In light of this, the whole discussion about const makes total sense to me now
and I would systematically apply the rule that in the end, we should never
have any const XXX in our codebase for unique'd types (assuming we can remove
them all in containers and no additional constness constraint is added on us
from the outside world).
PiperOrigin-RevId: 215811554
2018-10-05 06:10:33 +08:00
|
|
|
|
2018-10-06 09:24:18 +08:00
|
|
|
/// 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.
|
2019-03-11 08:44:25 +08:00
|
|
|
LogicalResult
|
|
|
|
AffineMap::constantFold(ArrayRef<Attribute> operandConstants,
|
|
|
|
SmallVectorImpl<Attribute> &results) const {
|
2018-10-06 09:24:18 +08:00
|
|
|
assert(getNumInputs() == operandConstants.size());
|
|
|
|
|
|
|
|
// Fold each of the result expressions.
|
|
|
|
AffineExprConstantFolder exprFolder(getNumDims(), operandConstants);
|
2018-10-10 01:59:27 +08:00
|
|
|
// Constant fold each AffineExpr in AffineMap and add to 'results'.
|
2018-10-06 09:24:18 +08:00
|
|
|
for (auto expr : getResults()) {
|
2018-10-26 06:46:10 +08:00
|
|
|
auto folded = exprFolder.constantFold(expr);
|
2018-10-06 09:24:18 +08:00
|
|
|
// If we didn't fold to a constant, then folding fails.
|
|
|
|
if (!folded)
|
2019-03-11 08:44:25 +08:00
|
|
|
return failure();
|
2018-10-06 09:24:18 +08:00
|
|
|
|
|
|
|
results.push_back(folded);
|
|
|
|
}
|
|
|
|
assert(results.size() == getNumResults() &&
|
|
|
|
"constant folding produced the wrong number of results");
|
2019-03-11 08:44:25 +08:00
|
|
|
return success();
|
2018-10-06 09:24:18 +08:00
|
|
|
}
|
2019-01-04 23:23:28 +08:00
|
|
|
|
2019-05-30 05:56:41 +08:00
|
|
|
/// Walk all of the AffineExpr's in this mapping. Each node in an expression
|
2019-01-04 23:23:28 +08:00
|
|
|
/// 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));
|
|
|
|
|
2019-05-30 05:56:41 +08:00
|
|
|
return get(numResultDims, numResultSyms, results);
|
2019-01-04 23:23:28 +08:00
|
|
|
}
|
2019-01-08 12:05:14 +08:00
|
|
|
|
|
|
|
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));
|
2019-05-30 05:56:41 +08:00
|
|
|
return AffineMap::get(numDims, numSymbols, exprs);
|
2019-01-08 12:05:14 +08:00
|
|
|
}
|
2019-02-06 09:00:13 +08:00
|
|
|
|
2019-06-15 01:18:08 +08:00
|
|
|
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();
|
|
|
|
}
|
|
|
|
|
2019-02-06 09:00:13 +08:00
|
|
|
AffineMap mlir::simplifyAffineMap(AffineMap map) {
|
2019-05-30 05:56:41 +08:00
|
|
|
SmallVector<AffineExpr, 8> exprs;
|
2019-02-06 09:00:13 +08:00
|
|
|
for (auto e : map.getResults()) {
|
|
|
|
exprs.push_back(
|
|
|
|
simplifyAffineExpr(e, map.getNumDims(), map.getNumSymbols()));
|
|
|
|
}
|
2019-05-30 05:56:41 +08:00
|
|
|
return AffineMap::get(map.getNumDims(), map.getNumSymbols(), exprs);
|
2019-02-06 09:00:13 +08:00
|
|
|
}
|
2019-05-01 21:47:32 +08:00
|
|
|
|
|
|
|
AffineMap mlir::inversePermutation(AffineMap map) {
|
2019-06-21 20:34:58 +08:00
|
|
|
if (!map)
|
|
|
|
return map;
|
2019-05-01 21:47:32 +08:00
|
|
|
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();
|
2019-06-22 09:31:52 +08:00
|
|
|
// Skip non-permutations.
|
|
|
|
if (auto d = expr.dyn_cast<AffineDimExpr>()) {
|
|
|
|
if (exprs[d.getPosition()])
|
|
|
|
continue;
|
|
|
|
exprs[d.getPosition()] = getAffineDimExpr(en.index(), d.getContext());
|
|
|
|
}
|
2019-05-01 21:47:32 +08:00
|
|
|
}
|
|
|
|
SmallVector<AffineExpr, 4> seenExprs;
|
|
|
|
seenExprs.reserve(map.getNumDims());
|
|
|
|
for (auto expr : exprs)
|
|
|
|
if (expr)
|
|
|
|
seenExprs.push_back(expr);
|
2019-06-22 09:31:52 +08:00
|
|
|
assert(seenExprs.size() == map.getNumInputs() &&
|
|
|
|
"map does not include a full rank permutation");
|
2019-05-30 05:56:41 +08:00
|
|
|
return AffineMap::get(map.getNumResults(), 0, seenExprs);
|
2019-05-01 21:47:32 +08:00
|
|
|
}
|
|
|
|
|
|
|
|
AffineMap mlir::concatAffineMaps(ArrayRef<AffineMap> maps) {
|
|
|
|
unsigned numResults = 0;
|
|
|
|
for (auto m : maps)
|
|
|
|
numResults += m ? m.getNumResults() : 0;
|
|
|
|
unsigned numDims = 0;
|
|
|
|
llvm::SmallVector<AffineExpr, 8> results;
|
|
|
|
results.reserve(numResults);
|
|
|
|
for (auto m : maps) {
|
|
|
|
if (!m)
|
|
|
|
continue;
|
|
|
|
assert(m.getNumSymbols() == 0 && "expected map without symbols");
|
|
|
|
results.append(m.getResults().begin(), m.getResults().end());
|
|
|
|
numDims = std::max(m.getNumDims(), numDims);
|
|
|
|
}
|
2019-06-21 20:34:58 +08:00
|
|
|
return numDims == 0 ? AffineMap() : AffineMap::get(numDims, 0, results);
|
2019-05-01 21:47:32 +08:00
|
|
|
}
|