mirror of
https://github.com/llvm/llvm-project.git
synced 2025-04-25 10:26:06 +00:00
60b44d31af
Whenever `symbolicDivide` returns nullptr when called from inside `simplifySemiAffine` we substitute the result with the original expression (`expr`). nullptr simply indicates that the floordiv expression cannot be simplified further. Fixes: https://github.com/llvm/llvm-project/issues/122231
1682 lines
67 KiB
C++
1682 lines
67 KiB
C++
//===- AffineExpr.cpp - MLIR Affine Expr 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 <cmath>
|
|
#include <cstdint>
|
|
#include <limits>
|
|
#include <utility>
|
|
|
|
#include "AffineExprDetail.h"
|
|
#include "mlir/IR/AffineExpr.h"
|
|
#include "mlir/IR/AffineExprVisitor.h"
|
|
#include "mlir/IR/AffineMap.h"
|
|
#include "mlir/IR/IntegerSet.h"
|
|
#include "mlir/Support/TypeID.h"
|
|
#include "llvm/ADT/STLExtras.h"
|
|
#include "llvm/Support/MathExtras.h"
|
|
#include <numeric>
|
|
#include <optional>
|
|
|
|
using namespace mlir;
|
|
using namespace mlir::detail;
|
|
|
|
using llvm::divideCeilSigned;
|
|
using llvm::divideFloorSigned;
|
|
using llvm::divideSignedWouldOverflow;
|
|
using llvm::mod;
|
|
|
|
MLIRContext *AffineExpr::getContext() const { return expr->context; }
|
|
|
|
AffineExprKind AffineExpr::getKind() const { return expr->kind; }
|
|
|
|
/// Walk all of the AffineExprs in `e` in postorder. This is a private factory
|
|
/// method to help handle lambda walk functions. Users should use the regular
|
|
/// (non-static) `walk` method.
|
|
template <typename WalkRetTy>
|
|
WalkRetTy mlir::AffineExpr::walk(AffineExpr e,
|
|
function_ref<WalkRetTy(AffineExpr)> callback) {
|
|
struct AffineExprWalker
|
|
: public AffineExprVisitor<AffineExprWalker, WalkRetTy> {
|
|
function_ref<WalkRetTy(AffineExpr)> callback;
|
|
|
|
AffineExprWalker(function_ref<WalkRetTy(AffineExpr)> callback)
|
|
: callback(callback) {}
|
|
|
|
WalkRetTy visitAffineBinaryOpExpr(AffineBinaryOpExpr expr) {
|
|
return callback(expr);
|
|
}
|
|
WalkRetTy visitConstantExpr(AffineConstantExpr expr) {
|
|
return callback(expr);
|
|
}
|
|
WalkRetTy visitDimExpr(AffineDimExpr expr) { return callback(expr); }
|
|
WalkRetTy visitSymbolExpr(AffineSymbolExpr expr) { return callback(expr); }
|
|
};
|
|
|
|
return AffineExprWalker(callback).walkPostOrder(e);
|
|
}
|
|
// Explicitly instantiate for the two supported return types.
|
|
template void mlir::AffineExpr::walk(AffineExpr e,
|
|
function_ref<void(AffineExpr)> callback);
|
|
template WalkResult
|
|
mlir::AffineExpr::walk(AffineExpr e,
|
|
function_ref<WalkResult(AffineExpr)> callback);
|
|
|
|
// Dispatch affine expression construction based on kind.
|
|
AffineExpr mlir::getAffineBinaryOpExpr(AffineExprKind kind, AffineExpr lhs,
|
|
AffineExpr rhs) {
|
|
if (kind == AffineExprKind::Add)
|
|
return lhs + rhs;
|
|
if (kind == AffineExprKind::Mul)
|
|
return lhs * rhs;
|
|
if (kind == AffineExprKind::FloorDiv)
|
|
return lhs.floorDiv(rhs);
|
|
if (kind == AffineExprKind::CeilDiv)
|
|
return lhs.ceilDiv(rhs);
|
|
if (kind == AffineExprKind::Mod)
|
|
return lhs % rhs;
|
|
|
|
llvm_unreachable("unknown binary operation on affine expressions");
|
|
}
|
|
|
|
/// This method substitutes any uses of dimensions and symbols (e.g.
|
|
/// dim#0 with dimReplacements[0]) and returns the modified expression tree.
|
|
AffineExpr
|
|
AffineExpr::replaceDimsAndSymbols(ArrayRef<AffineExpr> dimReplacements,
|
|
ArrayRef<AffineExpr> symReplacements) const {
|
|
switch (getKind()) {
|
|
case AffineExprKind::Constant:
|
|
return *this;
|
|
case AffineExprKind::DimId: {
|
|
unsigned dimId = llvm::cast<AffineDimExpr>(*this).getPosition();
|
|
if (dimId >= dimReplacements.size())
|
|
return *this;
|
|
return dimReplacements[dimId];
|
|
}
|
|
case AffineExprKind::SymbolId: {
|
|
unsigned symId = llvm::cast<AffineSymbolExpr>(*this).getPosition();
|
|
if (symId >= symReplacements.size())
|
|
return *this;
|
|
return symReplacements[symId];
|
|
}
|
|
case AffineExprKind::Add:
|
|
case AffineExprKind::Mul:
|
|
case AffineExprKind::FloorDiv:
|
|
case AffineExprKind::CeilDiv:
|
|
case AffineExprKind::Mod:
|
|
auto binOp = llvm::cast<AffineBinaryOpExpr>(*this);
|
|
auto lhs = binOp.getLHS(), rhs = binOp.getRHS();
|
|
auto newLHS = lhs.replaceDimsAndSymbols(dimReplacements, symReplacements);
|
|
auto newRHS = rhs.replaceDimsAndSymbols(dimReplacements, symReplacements);
|
|
if (newLHS == lhs && newRHS == rhs)
|
|
return *this;
|
|
return getAffineBinaryOpExpr(getKind(), newLHS, newRHS);
|
|
}
|
|
llvm_unreachable("Unknown AffineExpr");
|
|
}
|
|
|
|
AffineExpr AffineExpr::replaceDims(ArrayRef<AffineExpr> dimReplacements) const {
|
|
return replaceDimsAndSymbols(dimReplacements, {});
|
|
}
|
|
|
|
AffineExpr
|
|
AffineExpr::replaceSymbols(ArrayRef<AffineExpr> symReplacements) const {
|
|
return replaceDimsAndSymbols({}, symReplacements);
|
|
}
|
|
|
|
/// Replace dims[offset ... numDims)
|
|
/// by dims[offset + shift ... shift + numDims).
|
|
AffineExpr AffineExpr::shiftDims(unsigned numDims, unsigned shift,
|
|
unsigned offset) const {
|
|
SmallVector<AffineExpr, 4> dims;
|
|
for (unsigned idx = 0; idx < offset; ++idx)
|
|
dims.push_back(getAffineDimExpr(idx, getContext()));
|
|
for (unsigned idx = offset; idx < numDims; ++idx)
|
|
dims.push_back(getAffineDimExpr(idx + shift, getContext()));
|
|
return replaceDimsAndSymbols(dims, {});
|
|
}
|
|
|
|
/// Replace symbols[offset ... numSymbols)
|
|
/// by symbols[offset + shift ... shift + numSymbols).
|
|
AffineExpr AffineExpr::shiftSymbols(unsigned numSymbols, unsigned shift,
|
|
unsigned offset) const {
|
|
SmallVector<AffineExpr, 4> symbols;
|
|
for (unsigned idx = 0; idx < offset; ++idx)
|
|
symbols.push_back(getAffineSymbolExpr(idx, getContext()));
|
|
for (unsigned idx = offset; idx < numSymbols; ++idx)
|
|
symbols.push_back(getAffineSymbolExpr(idx + shift, getContext()));
|
|
return replaceDimsAndSymbols({}, symbols);
|
|
}
|
|
|
|
/// Sparse replace method. Return the modified expression tree.
|
|
AffineExpr
|
|
AffineExpr::replace(const DenseMap<AffineExpr, AffineExpr> &map) const {
|
|
auto it = map.find(*this);
|
|
if (it != map.end())
|
|
return it->second;
|
|
switch (getKind()) {
|
|
default:
|
|
return *this;
|
|
case AffineExprKind::Add:
|
|
case AffineExprKind::Mul:
|
|
case AffineExprKind::FloorDiv:
|
|
case AffineExprKind::CeilDiv:
|
|
case AffineExprKind::Mod:
|
|
auto binOp = llvm::cast<AffineBinaryOpExpr>(*this);
|
|
auto lhs = binOp.getLHS(), rhs = binOp.getRHS();
|
|
auto newLHS = lhs.replace(map);
|
|
auto newRHS = rhs.replace(map);
|
|
if (newLHS == lhs && newRHS == rhs)
|
|
return *this;
|
|
return getAffineBinaryOpExpr(getKind(), newLHS, newRHS);
|
|
}
|
|
llvm_unreachable("Unknown AffineExpr");
|
|
}
|
|
|
|
/// Sparse replace method. Return the modified expression tree.
|
|
AffineExpr AffineExpr::replace(AffineExpr expr, AffineExpr replacement) const {
|
|
DenseMap<AffineExpr, AffineExpr> map;
|
|
map.insert(std::make_pair(expr, replacement));
|
|
return replace(map);
|
|
}
|
|
/// Returns true if this expression is made out of only symbols and
|
|
/// constants (no dimensional identifiers).
|
|
bool AffineExpr::isSymbolicOrConstant() const {
|
|
switch (getKind()) {
|
|
case AffineExprKind::Constant:
|
|
return true;
|
|
case AffineExprKind::DimId:
|
|
return false;
|
|
case AffineExprKind::SymbolId:
|
|
return true;
|
|
|
|
case AffineExprKind::Add:
|
|
case AffineExprKind::Mul:
|
|
case AffineExprKind::FloorDiv:
|
|
case AffineExprKind::CeilDiv:
|
|
case AffineExprKind::Mod: {
|
|
auto expr = llvm::cast<AffineBinaryOpExpr>(*this);
|
|
return expr.getLHS().isSymbolicOrConstant() &&
|
|
expr.getRHS().isSymbolicOrConstant();
|
|
}
|
|
}
|
|
llvm_unreachable("Unknown AffineExpr");
|
|
}
|
|
|
|
/// Returns true if this is a pure affine expression, i.e., multiplication,
|
|
/// floordiv, ceildiv, and mod is only allowed w.r.t constants.
|
|
bool AffineExpr::isPureAffine() const {
|
|
switch (getKind()) {
|
|
case AffineExprKind::SymbolId:
|
|
case AffineExprKind::DimId:
|
|
case AffineExprKind::Constant:
|
|
return true;
|
|
case AffineExprKind::Add: {
|
|
auto op = llvm::cast<AffineBinaryOpExpr>(*this);
|
|
return op.getLHS().isPureAffine() && op.getRHS().isPureAffine();
|
|
}
|
|
|
|
case AffineExprKind::Mul: {
|
|
// TODO: Canonicalize the constants in binary operators to the RHS when
|
|
// possible, allowing this to merge into the next case.
|
|
auto op = llvm::cast<AffineBinaryOpExpr>(*this);
|
|
return op.getLHS().isPureAffine() && op.getRHS().isPureAffine() &&
|
|
(llvm::isa<AffineConstantExpr>(op.getLHS()) ||
|
|
llvm::isa<AffineConstantExpr>(op.getRHS()));
|
|
}
|
|
case AffineExprKind::FloorDiv:
|
|
case AffineExprKind::CeilDiv:
|
|
case AffineExprKind::Mod: {
|
|
auto op = llvm::cast<AffineBinaryOpExpr>(*this);
|
|
return op.getLHS().isPureAffine() &&
|
|
llvm::isa<AffineConstantExpr>(op.getRHS());
|
|
}
|
|
}
|
|
llvm_unreachable("Unknown AffineExpr");
|
|
}
|
|
|
|
// Returns the greatest known integral divisor of this affine expression.
|
|
int64_t AffineExpr::getLargestKnownDivisor() const {
|
|
AffineBinaryOpExpr binExpr(nullptr);
|
|
switch (getKind()) {
|
|
case AffineExprKind::DimId:
|
|
[[fallthrough]];
|
|
case AffineExprKind::SymbolId:
|
|
return 1;
|
|
case AffineExprKind::CeilDiv:
|
|
[[fallthrough]];
|
|
case AffineExprKind::FloorDiv: {
|
|
// If the RHS is a constant and divides the known divisor on the LHS, the
|
|
// quotient is a known divisor of the expression.
|
|
binExpr = llvm::cast<AffineBinaryOpExpr>(*this);
|
|
auto rhs = llvm::dyn_cast<AffineConstantExpr>(binExpr.getRHS());
|
|
// Leave alone undefined expressions.
|
|
if (rhs && rhs.getValue() != 0) {
|
|
int64_t lhsDiv = binExpr.getLHS().getLargestKnownDivisor();
|
|
if (lhsDiv % rhs.getValue() == 0)
|
|
return std::abs(lhsDiv / rhs.getValue());
|
|
}
|
|
return 1;
|
|
}
|
|
case AffineExprKind::Constant:
|
|
return std::abs(llvm::cast<AffineConstantExpr>(*this).getValue());
|
|
case AffineExprKind::Mul: {
|
|
binExpr = llvm::cast<AffineBinaryOpExpr>(*this);
|
|
return binExpr.getLHS().getLargestKnownDivisor() *
|
|
binExpr.getRHS().getLargestKnownDivisor();
|
|
}
|
|
case AffineExprKind::Add:
|
|
[[fallthrough]];
|
|
case AffineExprKind::Mod: {
|
|
binExpr = llvm::cast<AffineBinaryOpExpr>(*this);
|
|
return std::gcd((uint64_t)binExpr.getLHS().getLargestKnownDivisor(),
|
|
(uint64_t)binExpr.getRHS().getLargestKnownDivisor());
|
|
}
|
|
}
|
|
llvm_unreachable("Unknown AffineExpr");
|
|
}
|
|
|
|
bool AffineExpr::isMultipleOf(int64_t factor) const {
|
|
AffineBinaryOpExpr binExpr(nullptr);
|
|
uint64_t l, u;
|
|
switch (getKind()) {
|
|
case AffineExprKind::SymbolId:
|
|
[[fallthrough]];
|
|
case AffineExprKind::DimId:
|
|
return factor * factor == 1;
|
|
case AffineExprKind::Constant:
|
|
return llvm::cast<AffineConstantExpr>(*this).getValue() % factor == 0;
|
|
case AffineExprKind::Mul: {
|
|
binExpr = llvm::cast<AffineBinaryOpExpr>(*this);
|
|
// It's probably not worth optimizing this further (to not traverse the
|
|
// whole sub-tree under - it that would require a version of isMultipleOf
|
|
// that on a 'false' return also returns the largest known divisor).
|
|
return (l = binExpr.getLHS().getLargestKnownDivisor()) % factor == 0 ||
|
|
(u = binExpr.getRHS().getLargestKnownDivisor()) % factor == 0 ||
|
|
(l * u) % factor == 0;
|
|
}
|
|
case AffineExprKind::Add:
|
|
case AffineExprKind::FloorDiv:
|
|
case AffineExprKind::CeilDiv:
|
|
case AffineExprKind::Mod: {
|
|
binExpr = llvm::cast<AffineBinaryOpExpr>(*this);
|
|
return std::gcd((uint64_t)binExpr.getLHS().getLargestKnownDivisor(),
|
|
(uint64_t)binExpr.getRHS().getLargestKnownDivisor()) %
|
|
factor ==
|
|
0;
|
|
}
|
|
}
|
|
llvm_unreachable("Unknown AffineExpr");
|
|
}
|
|
|
|
bool AffineExpr::isFunctionOfDim(unsigned position) const {
|
|
if (getKind() == AffineExprKind::DimId) {
|
|
return *this == mlir::getAffineDimExpr(position, getContext());
|
|
}
|
|
if (auto expr = llvm::dyn_cast<AffineBinaryOpExpr>(*this)) {
|
|
return expr.getLHS().isFunctionOfDim(position) ||
|
|
expr.getRHS().isFunctionOfDim(position);
|
|
}
|
|
return false;
|
|
}
|
|
|
|
bool AffineExpr::isFunctionOfSymbol(unsigned position) const {
|
|
if (getKind() == AffineExprKind::SymbolId) {
|
|
return *this == mlir::getAffineSymbolExpr(position, getContext());
|
|
}
|
|
if (auto expr = llvm::dyn_cast<AffineBinaryOpExpr>(*this)) {
|
|
return expr.getLHS().isFunctionOfSymbol(position) ||
|
|
expr.getRHS().isFunctionOfSymbol(position);
|
|
}
|
|
return false;
|
|
}
|
|
|
|
AffineBinaryOpExpr::AffineBinaryOpExpr(AffineExpr::ImplType *ptr)
|
|
: AffineExpr(ptr) {}
|
|
AffineExpr AffineBinaryOpExpr::getLHS() const {
|
|
return static_cast<ImplType *>(expr)->lhs;
|
|
}
|
|
AffineExpr AffineBinaryOpExpr::getRHS() const {
|
|
return static_cast<ImplType *>(expr)->rhs;
|
|
}
|
|
|
|
AffineDimExpr::AffineDimExpr(AffineExpr::ImplType *ptr) : AffineExpr(ptr) {}
|
|
unsigned AffineDimExpr::getPosition() const {
|
|
return static_cast<ImplType *>(expr)->position;
|
|
}
|
|
|
|
/// Returns true if the expression is divisible by the given symbol with
|
|
/// position `symbolPos`. The argument `opKind` specifies here what kind of
|
|
/// division or mod operation called this division. It helps in implementing the
|
|
/// commutative property of the floordiv and ceildiv operations. If the argument
|
|
///`exprKind` is floordiv and `expr` is also a binary expression of a floordiv
|
|
/// operation, then the commutative property can be used otherwise, the floordiv
|
|
/// operation is not divisible. The same argument holds for ceildiv operation.
|
|
static bool canSimplifyDivisionBySymbol(AffineExpr expr, unsigned symbolPos,
|
|
AffineExprKind opKind,
|
|
bool fromMul = false) {
|
|
// The argument `opKind` can either be Modulo, Floordiv or Ceildiv only.
|
|
assert((opKind == AffineExprKind::Mod || opKind == AffineExprKind::FloorDiv ||
|
|
opKind == AffineExprKind::CeilDiv) &&
|
|
"unexpected opKind");
|
|
switch (expr.getKind()) {
|
|
case AffineExprKind::Constant:
|
|
return cast<AffineConstantExpr>(expr).getValue() == 0;
|
|
case AffineExprKind::DimId:
|
|
return false;
|
|
case AffineExprKind::SymbolId:
|
|
return (cast<AffineSymbolExpr>(expr).getPosition() == symbolPos);
|
|
// Checks divisibility by the given symbol for both operands.
|
|
case AffineExprKind::Add: {
|
|
AffineBinaryOpExpr binaryExpr = cast<AffineBinaryOpExpr>(expr);
|
|
return canSimplifyDivisionBySymbol(binaryExpr.getLHS(), symbolPos,
|
|
opKind) &&
|
|
canSimplifyDivisionBySymbol(binaryExpr.getRHS(), symbolPos, opKind);
|
|
}
|
|
// Checks divisibility by the given symbol for both operands. Consider the
|
|
// expression `(((s1*s0) floordiv w) mod ((s1 * s2) floordiv p)) floordiv s1`,
|
|
// this is a division by s1 and both the operands of modulo are divisible by
|
|
// s1 but it is not divisible by s1 always. The third argument is
|
|
// `AffineExprKind::Mod` for this reason.
|
|
case AffineExprKind::Mod: {
|
|
AffineBinaryOpExpr binaryExpr = cast<AffineBinaryOpExpr>(expr);
|
|
return canSimplifyDivisionBySymbol(binaryExpr.getLHS(), symbolPos,
|
|
AffineExprKind::Mod) &&
|
|
canSimplifyDivisionBySymbol(binaryExpr.getRHS(), symbolPos,
|
|
AffineExprKind::Mod);
|
|
}
|
|
// Checks if any of the operand divisible by the given symbol.
|
|
case AffineExprKind::Mul: {
|
|
AffineBinaryOpExpr binaryExpr = cast<AffineBinaryOpExpr>(expr);
|
|
return canSimplifyDivisionBySymbol(binaryExpr.getLHS(), symbolPos, opKind,
|
|
true) ||
|
|
canSimplifyDivisionBySymbol(binaryExpr.getRHS(), symbolPos, opKind,
|
|
true);
|
|
}
|
|
// Floordiv and ceildiv are divisible by the given symbol when the first
|
|
// operand is divisible, and the affine expression kind of the argument expr
|
|
// is same as the argument `opKind`. This can be inferred from commutative
|
|
// property of floordiv and ceildiv operations and are as follow:
|
|
// (exp1 floordiv exp2) floordiv exp3 = (exp1 floordiv exp3) floordiv exp2
|
|
// (exp1 ceildiv exp2) ceildiv exp3 = (exp1 ceildiv exp3) ceildiv expr2
|
|
// It will fail 1.if operations are not same. For example:
|
|
// (exps1 ceildiv exp2) floordiv exp3 can not be simplified. 2.if there is a
|
|
// multiplication operation in the expression. For example:
|
|
// (exps1 ceildiv exp2) mul exp3 ceildiv exp4 can not be simplified.
|
|
case AffineExprKind::FloorDiv:
|
|
case AffineExprKind::CeilDiv: {
|
|
AffineBinaryOpExpr binaryExpr = cast<AffineBinaryOpExpr>(expr);
|
|
if (opKind != expr.getKind())
|
|
return false;
|
|
if (fromMul)
|
|
return false;
|
|
return canSimplifyDivisionBySymbol(binaryExpr.getLHS(), symbolPos,
|
|
expr.getKind());
|
|
}
|
|
}
|
|
llvm_unreachable("Unknown AffineExpr");
|
|
}
|
|
|
|
/// Divides the given expression by the given symbol at position `symbolPos`. It
|
|
/// considers the divisibility condition is checked before calling itself. A
|
|
/// null expression is returned whenever the divisibility condition fails.
|
|
static AffineExpr symbolicDivide(AffineExpr expr, unsigned symbolPos,
|
|
AffineExprKind opKind) {
|
|
// THe argument `opKind` can either be Modulo, Floordiv or Ceildiv only.
|
|
assert((opKind == AffineExprKind::Mod || opKind == AffineExprKind::FloorDiv ||
|
|
opKind == AffineExprKind::CeilDiv) &&
|
|
"unexpected opKind");
|
|
switch (expr.getKind()) {
|
|
case AffineExprKind::Constant:
|
|
if (cast<AffineConstantExpr>(expr).getValue() != 0)
|
|
return nullptr;
|
|
return getAffineConstantExpr(0, expr.getContext());
|
|
case AffineExprKind::DimId:
|
|
return nullptr;
|
|
case AffineExprKind::SymbolId:
|
|
return getAffineConstantExpr(1, expr.getContext());
|
|
// Dividing both operands by the given symbol.
|
|
case AffineExprKind::Add: {
|
|
AffineBinaryOpExpr binaryExpr = cast<AffineBinaryOpExpr>(expr);
|
|
return getAffineBinaryOpExpr(
|
|
expr.getKind(), symbolicDivide(binaryExpr.getLHS(), symbolPos, opKind),
|
|
symbolicDivide(binaryExpr.getRHS(), symbolPos, opKind));
|
|
}
|
|
// Dividing both operands by the given symbol.
|
|
case AffineExprKind::Mod: {
|
|
AffineBinaryOpExpr binaryExpr = cast<AffineBinaryOpExpr>(expr);
|
|
return getAffineBinaryOpExpr(
|
|
expr.getKind(),
|
|
symbolicDivide(binaryExpr.getLHS(), symbolPos, expr.getKind()),
|
|
symbolicDivide(binaryExpr.getRHS(), symbolPos, expr.getKind()));
|
|
}
|
|
// Dividing any of the operand by the given symbol.
|
|
case AffineExprKind::Mul: {
|
|
AffineBinaryOpExpr binaryExpr = cast<AffineBinaryOpExpr>(expr);
|
|
if (!canSimplifyDivisionBySymbol(binaryExpr.getLHS(), symbolPos, opKind))
|
|
return binaryExpr.getLHS() *
|
|
symbolicDivide(binaryExpr.getRHS(), symbolPos, opKind);
|
|
return symbolicDivide(binaryExpr.getLHS(), symbolPos, opKind) *
|
|
binaryExpr.getRHS();
|
|
}
|
|
// Dividing first operand only by the given symbol.
|
|
case AffineExprKind::FloorDiv:
|
|
case AffineExprKind::CeilDiv: {
|
|
AffineBinaryOpExpr binaryExpr = cast<AffineBinaryOpExpr>(expr);
|
|
return getAffineBinaryOpExpr(
|
|
expr.getKind(),
|
|
symbolicDivide(binaryExpr.getLHS(), symbolPos, expr.getKind()),
|
|
binaryExpr.getRHS());
|
|
}
|
|
}
|
|
llvm_unreachable("Unknown AffineExpr");
|
|
}
|
|
|
|
/// Populate `result` with all summand operands of given (potentially nested)
|
|
/// addition. If the given expression is not an addition, just populate the
|
|
/// expression itself.
|
|
/// Example: Add(Add(7, 8), Mul(9, 10)) will return [7, 8, Mul(9, 10)].
|
|
static void getSummandExprs(AffineExpr expr, SmallVector<AffineExpr> &result) {
|
|
auto addExpr = dyn_cast<AffineBinaryOpExpr>(expr);
|
|
if (!addExpr || addExpr.getKind() != AffineExprKind::Add) {
|
|
result.push_back(expr);
|
|
return;
|
|
}
|
|
getSummandExprs(addExpr.getLHS(), result);
|
|
getSummandExprs(addExpr.getRHS(), result);
|
|
}
|
|
|
|
/// Return "true" if `candidate` is a negated expression, i.e., Mul(-1, expr).
|
|
/// If so, also return the non-negated expression via `expr`.
|
|
static bool isNegatedAffineExpr(AffineExpr candidate, AffineExpr &expr) {
|
|
auto mulExpr = dyn_cast<AffineBinaryOpExpr>(candidate);
|
|
if (!mulExpr || mulExpr.getKind() != AffineExprKind::Mul)
|
|
return false;
|
|
if (auto lhs = dyn_cast<AffineConstantExpr>(mulExpr.getLHS())) {
|
|
if (lhs.getValue() == -1) {
|
|
expr = mulExpr.getRHS();
|
|
return true;
|
|
}
|
|
}
|
|
if (auto rhs = dyn_cast<AffineConstantExpr>(mulExpr.getRHS())) {
|
|
if (rhs.getValue() == -1) {
|
|
expr = mulExpr.getLHS();
|
|
return true;
|
|
}
|
|
}
|
|
return false;
|
|
}
|
|
|
|
/// Return "true" if `lhs` % `rhs` is guaranteed to evaluate to zero based on
|
|
/// the fact that `lhs` contains another modulo expression that ensures that
|
|
/// `lhs` is divisible by `rhs`. This is a common pattern in the resulting IR
|
|
/// after loop peeling.
|
|
///
|
|
/// Example: lhs = ub - ub % step
|
|
/// rhs = step
|
|
/// => (ub - ub % step) % step is guaranteed to evaluate to 0.
|
|
static bool isModOfModSubtraction(AffineExpr lhs, AffineExpr rhs,
|
|
unsigned numDims, unsigned numSymbols) {
|
|
// TODO: Try to unify this function with `getBoundForAffineExpr`.
|
|
// Collect all summands in lhs.
|
|
SmallVector<AffineExpr> summands;
|
|
getSummandExprs(lhs, summands);
|
|
// Look for Mul(-1, Mod(x, rhs)) among the summands. If x matches the
|
|
// remaining summands, then lhs % rhs is guaranteed to evaluate to 0.
|
|
for (int64_t i = 0, e = summands.size(); i < e; ++i) {
|
|
AffineExpr current = summands[i];
|
|
AffineExpr beforeNegation;
|
|
if (!isNegatedAffineExpr(current, beforeNegation))
|
|
continue;
|
|
AffineBinaryOpExpr innerMod = dyn_cast<AffineBinaryOpExpr>(beforeNegation);
|
|
if (!innerMod || innerMod.getKind() != AffineExprKind::Mod)
|
|
continue;
|
|
if (innerMod.getRHS() != rhs)
|
|
continue;
|
|
// Sum all remaining summands and subtract x. If that expression can be
|
|
// simplified to zero, then the remaining summands and x are equal.
|
|
AffineExpr diff = getAffineConstantExpr(0, lhs.getContext());
|
|
for (int64_t j = 0; j < e; ++j)
|
|
if (i != j)
|
|
diff = diff + summands[j];
|
|
diff = diff - innerMod.getLHS();
|
|
diff = simplifyAffineExpr(diff, numDims, numSymbols);
|
|
auto constExpr = dyn_cast<AffineConstantExpr>(diff);
|
|
if (constExpr && constExpr.getValue() == 0)
|
|
return true;
|
|
}
|
|
return false;
|
|
}
|
|
|
|
/// Simplify a semi-affine expression by handling modulo, floordiv, or ceildiv
|
|
/// operations when the second operand simplifies to a symbol and the first
|
|
/// operand is divisible by that symbol. It can be applied to any semi-affine
|
|
/// expression. Returned expression can either be a semi-affine or pure affine
|
|
/// expression.
|
|
static AffineExpr simplifySemiAffine(AffineExpr expr, unsigned numDims,
|
|
unsigned numSymbols) {
|
|
switch (expr.getKind()) {
|
|
case AffineExprKind::Constant:
|
|
case AffineExprKind::DimId:
|
|
case AffineExprKind::SymbolId:
|
|
return expr;
|
|
case AffineExprKind::Add:
|
|
case AffineExprKind::Mul: {
|
|
AffineBinaryOpExpr binaryExpr = cast<AffineBinaryOpExpr>(expr);
|
|
return getAffineBinaryOpExpr(
|
|
expr.getKind(),
|
|
simplifySemiAffine(binaryExpr.getLHS(), numDims, numSymbols),
|
|
simplifySemiAffine(binaryExpr.getRHS(), numDims, numSymbols));
|
|
}
|
|
// Check if the simplification of the second operand is a symbol, and the
|
|
// first operand is divisible by it. If the operation is a modulo, a constant
|
|
// zero expression is returned. In the case of floordiv and ceildiv, the
|
|
// symbol from the simplification of the second operand divides the first
|
|
// operand. Otherwise, simplification is not possible.
|
|
case AffineExprKind::FloorDiv:
|
|
case AffineExprKind::CeilDiv:
|
|
case AffineExprKind::Mod: {
|
|
AffineBinaryOpExpr binaryExpr = cast<AffineBinaryOpExpr>(expr);
|
|
AffineExpr sLHS =
|
|
simplifySemiAffine(binaryExpr.getLHS(), numDims, numSymbols);
|
|
AffineExpr sRHS =
|
|
simplifySemiAffine(binaryExpr.getRHS(), numDims, numSymbols);
|
|
if (isModOfModSubtraction(sLHS, sRHS, numDims, numSymbols))
|
|
return getAffineConstantExpr(0, expr.getContext());
|
|
AffineSymbolExpr symbolExpr = dyn_cast<AffineSymbolExpr>(
|
|
simplifySemiAffine(binaryExpr.getRHS(), numDims, numSymbols));
|
|
if (!symbolExpr)
|
|
return getAffineBinaryOpExpr(expr.getKind(), sLHS, sRHS);
|
|
unsigned symbolPos = symbolExpr.getPosition();
|
|
if (!canSimplifyDivisionBySymbol(binaryExpr.getLHS(), symbolPos,
|
|
expr.getKind()))
|
|
return getAffineBinaryOpExpr(expr.getKind(), sLHS, sRHS);
|
|
if (expr.getKind() == AffineExprKind::Mod)
|
|
return getAffineConstantExpr(0, expr.getContext());
|
|
AffineExpr simplifiedQuotient =
|
|
symbolicDivide(sLHS, symbolPos, expr.getKind());
|
|
return simplifiedQuotient
|
|
? simplifiedQuotient
|
|
: getAffineBinaryOpExpr(expr.getKind(), sLHS, sRHS);
|
|
}
|
|
}
|
|
llvm_unreachable("Unknown AffineExpr");
|
|
}
|
|
|
|
static AffineExpr getAffineDimOrSymbol(AffineExprKind kind, unsigned position,
|
|
MLIRContext *context) {
|
|
auto assignCtx = [context](AffineDimExprStorage *storage) {
|
|
storage->context = context;
|
|
};
|
|
|
|
StorageUniquer &uniquer = context->getAffineUniquer();
|
|
return uniquer.get<AffineDimExprStorage>(
|
|
assignCtx, static_cast<unsigned>(kind), position);
|
|
}
|
|
|
|
AffineExpr mlir::getAffineDimExpr(unsigned position, MLIRContext *context) {
|
|
return getAffineDimOrSymbol(AffineExprKind::DimId, position, context);
|
|
}
|
|
|
|
AffineSymbolExpr::AffineSymbolExpr(AffineExpr::ImplType *ptr)
|
|
: AffineExpr(ptr) {}
|
|
unsigned AffineSymbolExpr::getPosition() const {
|
|
return static_cast<ImplType *>(expr)->position;
|
|
}
|
|
|
|
AffineExpr mlir::getAffineSymbolExpr(unsigned position, MLIRContext *context) {
|
|
return getAffineDimOrSymbol(AffineExprKind::SymbolId, position, context);
|
|
}
|
|
|
|
AffineConstantExpr::AffineConstantExpr(AffineExpr::ImplType *ptr)
|
|
: AffineExpr(ptr) {}
|
|
int64_t AffineConstantExpr::getValue() const {
|
|
return static_cast<ImplType *>(expr)->constant;
|
|
}
|
|
|
|
bool AffineExpr::operator==(int64_t v) const {
|
|
return *this == getAffineConstantExpr(v, getContext());
|
|
}
|
|
|
|
AffineExpr mlir::getAffineConstantExpr(int64_t constant, MLIRContext *context) {
|
|
auto assignCtx = [context](AffineConstantExprStorage *storage) {
|
|
storage->context = context;
|
|
};
|
|
|
|
StorageUniquer &uniquer = context->getAffineUniquer();
|
|
return uniquer.get<AffineConstantExprStorage>(assignCtx, constant);
|
|
}
|
|
|
|
SmallVector<AffineExpr>
|
|
mlir::getAffineConstantExprs(ArrayRef<int64_t> constants,
|
|
MLIRContext *context) {
|
|
return llvm::to_vector(llvm::map_range(constants, [&](int64_t constant) {
|
|
return getAffineConstantExpr(constant, context);
|
|
}));
|
|
}
|
|
|
|
/// Simplify add expression. Return nullptr if it can't be simplified.
|
|
static AffineExpr simplifyAdd(AffineExpr lhs, AffineExpr rhs) {
|
|
auto lhsConst = dyn_cast<AffineConstantExpr>(lhs);
|
|
auto rhsConst = dyn_cast<AffineConstantExpr>(rhs);
|
|
// Fold if both LHS, RHS are a constant and the sum does not overflow.
|
|
if (lhsConst && rhsConst) {
|
|
int64_t sum;
|
|
if (llvm::AddOverflow(lhsConst.getValue(), rhsConst.getValue(), sum)) {
|
|
return nullptr;
|
|
}
|
|
return getAffineConstantExpr(sum, lhs.getContext());
|
|
}
|
|
|
|
// Canonicalize so that only the RHS is a constant. (4 + d0 becomes d0 + 4).
|
|
// If only one of them is a symbolic expressions, make it the RHS.
|
|
if (isa<AffineConstantExpr>(lhs) ||
|
|
(lhs.isSymbolicOrConstant() && !rhs.isSymbolicOrConstant())) {
|
|
return rhs + lhs;
|
|
}
|
|
|
|
// At this point, if there was a constant, it would be on the right.
|
|
|
|
// Addition with a zero is a noop, return the other input.
|
|
if (rhsConst) {
|
|
if (rhsConst.getValue() == 0)
|
|
return lhs;
|
|
}
|
|
// Fold successive additions like (d0 + 2) + 3 into d0 + 5.
|
|
auto lBin = dyn_cast<AffineBinaryOpExpr>(lhs);
|
|
if (lBin && rhsConst && lBin.getKind() == AffineExprKind::Add) {
|
|
if (auto lrhs = dyn_cast<AffineConstantExpr>(lBin.getRHS()))
|
|
return lBin.getLHS() + (lrhs.getValue() + rhsConst.getValue());
|
|
}
|
|
|
|
// Detect "c1 * expr + c_2 * expr" as "(c1 + c2) * expr".
|
|
// c1 is rRhsConst, c2 is rLhsConst; firstExpr, secondExpr are their
|
|
// respective multiplicands.
|
|
std::optional<int64_t> rLhsConst, rRhsConst;
|
|
AffineExpr firstExpr, secondExpr;
|
|
AffineConstantExpr rLhsConstExpr;
|
|
auto lBinOpExpr = dyn_cast<AffineBinaryOpExpr>(lhs);
|
|
if (lBinOpExpr && lBinOpExpr.getKind() == AffineExprKind::Mul &&
|
|
(rLhsConstExpr = dyn_cast<AffineConstantExpr>(lBinOpExpr.getRHS()))) {
|
|
rLhsConst = rLhsConstExpr.getValue();
|
|
firstExpr = lBinOpExpr.getLHS();
|
|
} else {
|
|
rLhsConst = 1;
|
|
firstExpr = lhs;
|
|
}
|
|
|
|
auto rBinOpExpr = dyn_cast<AffineBinaryOpExpr>(rhs);
|
|
AffineConstantExpr rRhsConstExpr;
|
|
if (rBinOpExpr && rBinOpExpr.getKind() == AffineExprKind::Mul &&
|
|
(rRhsConstExpr = dyn_cast<AffineConstantExpr>(rBinOpExpr.getRHS()))) {
|
|
rRhsConst = rRhsConstExpr.getValue();
|
|
secondExpr = rBinOpExpr.getLHS();
|
|
} else {
|
|
rRhsConst = 1;
|
|
secondExpr = rhs;
|
|
}
|
|
|
|
if (rLhsConst && rRhsConst && firstExpr == secondExpr)
|
|
return getAffineBinaryOpExpr(
|
|
AffineExprKind::Mul, firstExpr,
|
|
getAffineConstantExpr(*rLhsConst + *rRhsConst, lhs.getContext()));
|
|
|
|
// When doing successive additions, bring constant to the right: turn (d0 + 2)
|
|
// + d1 into (d0 + d1) + 2.
|
|
if (lBin && lBin.getKind() == AffineExprKind::Add) {
|
|
if (auto lrhs = dyn_cast<AffineConstantExpr>(lBin.getRHS())) {
|
|
return lBin.getLHS() + rhs + lrhs;
|
|
}
|
|
}
|
|
|
|
// Detect and transform "expr - q * (expr floordiv q)" to "expr mod q", where
|
|
// q may be a constant or symbolic expression. This leads to a much more
|
|
// efficient form when 'c' is a power of two, and in general a more compact
|
|
// and readable form.
|
|
|
|
// Process '(expr floordiv c) * (-c)'.
|
|
if (!rBinOpExpr)
|
|
return nullptr;
|
|
|
|
auto lrhs = rBinOpExpr.getLHS();
|
|
auto rrhs = rBinOpExpr.getRHS();
|
|
|
|
AffineExpr llrhs, rlrhs;
|
|
|
|
// Check if lrhsBinOpExpr is of the form (expr floordiv q) * q, where q is a
|
|
// symbolic expression.
|
|
auto lrhsBinOpExpr = dyn_cast<AffineBinaryOpExpr>(lrhs);
|
|
// Check rrhsConstOpExpr = -1.
|
|
auto rrhsConstOpExpr = dyn_cast<AffineConstantExpr>(rrhs);
|
|
if (rrhsConstOpExpr && rrhsConstOpExpr.getValue() == -1 && lrhsBinOpExpr &&
|
|
lrhsBinOpExpr.getKind() == AffineExprKind::Mul) {
|
|
// Check llrhs = expr floordiv q.
|
|
llrhs = lrhsBinOpExpr.getLHS();
|
|
// Check rlrhs = q.
|
|
rlrhs = lrhsBinOpExpr.getRHS();
|
|
auto llrhsBinOpExpr = dyn_cast<AffineBinaryOpExpr>(llrhs);
|
|
if (!llrhsBinOpExpr || llrhsBinOpExpr.getKind() != AffineExprKind::FloorDiv)
|
|
return nullptr;
|
|
if (llrhsBinOpExpr.getRHS() == rlrhs && lhs == llrhsBinOpExpr.getLHS())
|
|
return lhs % rlrhs;
|
|
}
|
|
|
|
// Process lrhs, which is 'expr floordiv c'.
|
|
// expr + (expr // c * -c) = expr % c
|
|
AffineBinaryOpExpr lrBinOpExpr = dyn_cast<AffineBinaryOpExpr>(lrhs);
|
|
if (!lrBinOpExpr || rhs.getKind() != AffineExprKind::Mul ||
|
|
lrBinOpExpr.getKind() != AffineExprKind::FloorDiv)
|
|
return nullptr;
|
|
|
|
llrhs = lrBinOpExpr.getLHS();
|
|
rlrhs = lrBinOpExpr.getRHS();
|
|
auto rlrhsConstOpExpr = dyn_cast<AffineConstantExpr>(rlrhs);
|
|
// We don't support modulo with a negative RHS.
|
|
bool isPositiveRhs = rlrhsConstOpExpr && rlrhsConstOpExpr.getValue() > 0;
|
|
|
|
if (isPositiveRhs && lhs == llrhs && rlrhs == -rrhs) {
|
|
return lhs % rlrhs;
|
|
}
|
|
return nullptr;
|
|
}
|
|
|
|
AffineExpr AffineExpr::operator+(int64_t v) const {
|
|
return *this + getAffineConstantExpr(v, getContext());
|
|
}
|
|
AffineExpr AffineExpr::operator+(AffineExpr other) const {
|
|
if (auto simplified = simplifyAdd(*this, other))
|
|
return simplified;
|
|
|
|
StorageUniquer &uniquer = getContext()->getAffineUniquer();
|
|
return uniquer.get<AffineBinaryOpExprStorage>(
|
|
/*initFn=*/{}, static_cast<unsigned>(AffineExprKind::Add), *this, other);
|
|
}
|
|
|
|
/// Simplify a multiply expression. Return nullptr if it can't be simplified.
|
|
static AffineExpr simplifyMul(AffineExpr lhs, AffineExpr rhs) {
|
|
auto lhsConst = dyn_cast<AffineConstantExpr>(lhs);
|
|
auto rhsConst = dyn_cast<AffineConstantExpr>(rhs);
|
|
|
|
if (lhsConst && rhsConst) {
|
|
int64_t product;
|
|
if (llvm::MulOverflow(lhsConst.getValue(), rhsConst.getValue(), product)) {
|
|
return nullptr;
|
|
}
|
|
return getAffineConstantExpr(product, lhs.getContext());
|
|
}
|
|
|
|
if (!lhs.isSymbolicOrConstant() && !rhs.isSymbolicOrConstant())
|
|
return nullptr;
|
|
|
|
// Canonicalize the mul expression so that the constant/symbolic term is the
|
|
// RHS. If both the lhs and rhs are symbolic, swap them if the lhs is a
|
|
// constant. (Note that a constant is trivially symbolic).
|
|
if (!rhs.isSymbolicOrConstant() || isa<AffineConstantExpr>(lhs)) {
|
|
// At least one of them has to be symbolic.
|
|
return rhs * lhs;
|
|
}
|
|
|
|
// At this point, if there was a constant, it would be on the right.
|
|
|
|
// Multiplication with a one is a noop, return the other input.
|
|
if (rhsConst) {
|
|
if (rhsConst.getValue() == 1)
|
|
return lhs;
|
|
// Multiplication with zero.
|
|
if (rhsConst.getValue() == 0)
|
|
return rhsConst;
|
|
}
|
|
|
|
// Fold successive multiplications: eg: (d0 * 2) * 3 into d0 * 6.
|
|
auto lBin = dyn_cast<AffineBinaryOpExpr>(lhs);
|
|
if (lBin && rhsConst && lBin.getKind() == AffineExprKind::Mul) {
|
|
if (auto lrhs = dyn_cast<AffineConstantExpr>(lBin.getRHS()))
|
|
return lBin.getLHS() * (lrhs.getValue() * rhsConst.getValue());
|
|
}
|
|
|
|
// When doing successive multiplication, bring constant to the right: turn (d0
|
|
// * 2) * d1 into (d0 * d1) * 2.
|
|
if (lBin && lBin.getKind() == AffineExprKind::Mul) {
|
|
if (auto lrhs = dyn_cast<AffineConstantExpr>(lBin.getRHS())) {
|
|
return (lBin.getLHS() * rhs) * lrhs;
|
|
}
|
|
}
|
|
|
|
return nullptr;
|
|
}
|
|
|
|
AffineExpr AffineExpr::operator*(int64_t v) const {
|
|
return *this * getAffineConstantExpr(v, getContext());
|
|
}
|
|
AffineExpr AffineExpr::operator*(AffineExpr other) const {
|
|
if (auto simplified = simplifyMul(*this, other))
|
|
return simplified;
|
|
|
|
StorageUniquer &uniquer = getContext()->getAffineUniquer();
|
|
return uniquer.get<AffineBinaryOpExprStorage>(
|
|
/*initFn=*/{}, static_cast<unsigned>(AffineExprKind::Mul), *this, other);
|
|
}
|
|
|
|
// Unary minus, delegate to operator*.
|
|
AffineExpr AffineExpr::operator-() const {
|
|
return *this * getAffineConstantExpr(-1, getContext());
|
|
}
|
|
|
|
// Delegate to operator+.
|
|
AffineExpr AffineExpr::operator-(int64_t v) const { return *this + (-v); }
|
|
AffineExpr AffineExpr::operator-(AffineExpr other) const {
|
|
return *this + (-other);
|
|
}
|
|
|
|
static AffineExpr simplifyFloorDiv(AffineExpr lhs, AffineExpr rhs) {
|
|
auto lhsConst = dyn_cast<AffineConstantExpr>(lhs);
|
|
auto rhsConst = dyn_cast<AffineConstantExpr>(rhs);
|
|
|
|
if (!rhsConst || rhsConst.getValue() == 0)
|
|
return nullptr;
|
|
|
|
if (lhsConst) {
|
|
if (divideSignedWouldOverflow(lhsConst.getValue(), rhsConst.getValue()))
|
|
return nullptr;
|
|
return getAffineConstantExpr(
|
|
divideFloorSigned(lhsConst.getValue(), rhsConst.getValue()),
|
|
lhs.getContext());
|
|
}
|
|
|
|
// Fold floordiv of a multiply with a constant that is a multiple of the
|
|
// divisor. Eg: (i * 128) floordiv 64 = i * 2.
|
|
if (rhsConst == 1)
|
|
return lhs;
|
|
|
|
// Simplify `(expr * lrhs) floordiv rhsConst` when `lrhs` is known to be a
|
|
// multiple of `rhsConst`.
|
|
auto lBin = dyn_cast<AffineBinaryOpExpr>(lhs);
|
|
if (lBin && lBin.getKind() == AffineExprKind::Mul) {
|
|
if (auto lrhs = dyn_cast<AffineConstantExpr>(lBin.getRHS())) {
|
|
// `rhsConst` is known to be a nonzero constant.
|
|
if (lrhs.getValue() % rhsConst.getValue() == 0)
|
|
return lBin.getLHS() * (lrhs.getValue() / rhsConst.getValue());
|
|
}
|
|
}
|
|
|
|
// Simplify (expr1 + expr2) floordiv divConst when either expr1 or expr2 is
|
|
// known to be a multiple of divConst.
|
|
if (lBin && lBin.getKind() == AffineExprKind::Add) {
|
|
int64_t llhsDiv = lBin.getLHS().getLargestKnownDivisor();
|
|
int64_t lrhsDiv = lBin.getRHS().getLargestKnownDivisor();
|
|
// rhsConst is known to be a nonzero constant.
|
|
if (llhsDiv % rhsConst.getValue() == 0 ||
|
|
lrhsDiv % rhsConst.getValue() == 0)
|
|
return lBin.getLHS().floorDiv(rhsConst.getValue()) +
|
|
lBin.getRHS().floorDiv(rhsConst.getValue());
|
|
}
|
|
|
|
return nullptr;
|
|
}
|
|
|
|
AffineExpr AffineExpr::floorDiv(uint64_t v) const {
|
|
return floorDiv(getAffineConstantExpr(v, getContext()));
|
|
}
|
|
AffineExpr AffineExpr::floorDiv(AffineExpr other) const {
|
|
if (auto simplified = simplifyFloorDiv(*this, other))
|
|
return simplified;
|
|
|
|
StorageUniquer &uniquer = getContext()->getAffineUniquer();
|
|
return uniquer.get<AffineBinaryOpExprStorage>(
|
|
/*initFn=*/{}, static_cast<unsigned>(AffineExprKind::FloorDiv), *this,
|
|
other);
|
|
}
|
|
|
|
static AffineExpr simplifyCeilDiv(AffineExpr lhs, AffineExpr rhs) {
|
|
auto lhsConst = dyn_cast<AffineConstantExpr>(lhs);
|
|
auto rhsConst = dyn_cast<AffineConstantExpr>(rhs);
|
|
|
|
if (!rhsConst || rhsConst.getValue() == 0)
|
|
return nullptr;
|
|
|
|
if (lhsConst) {
|
|
if (divideSignedWouldOverflow(lhsConst.getValue(), rhsConst.getValue()))
|
|
return nullptr;
|
|
return getAffineConstantExpr(
|
|
divideCeilSigned(lhsConst.getValue(), rhsConst.getValue()),
|
|
lhs.getContext());
|
|
}
|
|
|
|
// Fold ceildiv of a multiply with a constant that is a multiple of the
|
|
// divisor. Eg: (i * 128) ceildiv 64 = i * 2.
|
|
if (rhsConst.getValue() == 1)
|
|
return lhs;
|
|
|
|
// Simplify `(expr * lrhs) ceildiv rhsConst` when `lrhs` is known to be a
|
|
// multiple of `rhsConst`.
|
|
auto lBin = dyn_cast<AffineBinaryOpExpr>(lhs);
|
|
if (lBin && lBin.getKind() == AffineExprKind::Mul) {
|
|
if (auto lrhs = dyn_cast<AffineConstantExpr>(lBin.getRHS())) {
|
|
// `rhsConst` is known to be a nonzero constant.
|
|
if (lrhs.getValue() % rhsConst.getValue() == 0)
|
|
return lBin.getLHS() * (lrhs.getValue() / rhsConst.getValue());
|
|
}
|
|
}
|
|
|
|
return nullptr;
|
|
}
|
|
|
|
AffineExpr AffineExpr::ceilDiv(uint64_t v) const {
|
|
return ceilDiv(getAffineConstantExpr(v, getContext()));
|
|
}
|
|
AffineExpr AffineExpr::ceilDiv(AffineExpr other) const {
|
|
if (auto simplified = simplifyCeilDiv(*this, other))
|
|
return simplified;
|
|
|
|
StorageUniquer &uniquer = getContext()->getAffineUniquer();
|
|
return uniquer.get<AffineBinaryOpExprStorage>(
|
|
/*initFn=*/{}, static_cast<unsigned>(AffineExprKind::CeilDiv), *this,
|
|
other);
|
|
}
|
|
|
|
static AffineExpr simplifyMod(AffineExpr lhs, AffineExpr rhs) {
|
|
auto lhsConst = dyn_cast<AffineConstantExpr>(lhs);
|
|
auto rhsConst = dyn_cast<AffineConstantExpr>(rhs);
|
|
|
|
// mod w.r.t zero or negative numbers is undefined and preserved as is.
|
|
if (!rhsConst || rhsConst.getValue() < 1)
|
|
return nullptr;
|
|
|
|
if (lhsConst) {
|
|
// mod never overflows.
|
|
return getAffineConstantExpr(mod(lhsConst.getValue(), rhsConst.getValue()),
|
|
lhs.getContext());
|
|
}
|
|
|
|
// Fold modulo of an expression that is known to be a multiple of a constant
|
|
// to zero if that constant is a multiple of the modulo factor. Eg: (i * 128)
|
|
// mod 64 is folded to 0, and less trivially, (i*(j*4*(k*32))) mod 128 = 0.
|
|
if (lhs.getLargestKnownDivisor() % rhsConst.getValue() == 0)
|
|
return getAffineConstantExpr(0, lhs.getContext());
|
|
|
|
// Simplify (expr1 + expr2) mod divConst when either expr1 or expr2 is
|
|
// known to be a multiple of divConst.
|
|
auto lBin = dyn_cast<AffineBinaryOpExpr>(lhs);
|
|
if (lBin && lBin.getKind() == AffineExprKind::Add) {
|
|
int64_t llhsDiv = lBin.getLHS().getLargestKnownDivisor();
|
|
int64_t lrhsDiv = lBin.getRHS().getLargestKnownDivisor();
|
|
// rhsConst is known to be a positive constant.
|
|
if (llhsDiv % rhsConst.getValue() == 0)
|
|
return lBin.getRHS() % rhsConst.getValue();
|
|
if (lrhsDiv % rhsConst.getValue() == 0)
|
|
return lBin.getLHS() % rhsConst.getValue();
|
|
}
|
|
|
|
// Simplify (e % a) % b to e % b when b evenly divides a
|
|
if (lBin && lBin.getKind() == AffineExprKind::Mod) {
|
|
auto intermediate = dyn_cast<AffineConstantExpr>(lBin.getRHS());
|
|
if (intermediate && intermediate.getValue() >= 1 &&
|
|
mod(intermediate.getValue(), rhsConst.getValue()) == 0) {
|
|
return lBin.getLHS() % rhsConst.getValue();
|
|
}
|
|
}
|
|
|
|
return nullptr;
|
|
}
|
|
|
|
AffineExpr AffineExpr::operator%(uint64_t v) const {
|
|
return *this % getAffineConstantExpr(v, getContext());
|
|
}
|
|
AffineExpr AffineExpr::operator%(AffineExpr other) const {
|
|
if (auto simplified = simplifyMod(*this, other))
|
|
return simplified;
|
|
|
|
StorageUniquer &uniquer = getContext()->getAffineUniquer();
|
|
return uniquer.get<AffineBinaryOpExprStorage>(
|
|
/*initFn=*/{}, static_cast<unsigned>(AffineExprKind::Mod), *this, other);
|
|
}
|
|
|
|
AffineExpr AffineExpr::compose(AffineMap map) const {
|
|
SmallVector<AffineExpr, 8> dimReplacements(map.getResults());
|
|
return replaceDimsAndSymbols(dimReplacements, {});
|
|
}
|
|
raw_ostream &mlir::operator<<(raw_ostream &os, AffineExpr expr) {
|
|
expr.print(os);
|
|
return os;
|
|
}
|
|
|
|
/// Constructs an affine expression from a flat ArrayRef. If there are local
|
|
/// identifiers (neither dimensional nor symbolic) that appear in the sum of
|
|
/// products expression, `localExprs` is expected to have the AffineExpr
|
|
/// for it, and is substituted into. The ArrayRef `flatExprs` is expected to be
|
|
/// in the format [dims, symbols, locals, constant term].
|
|
AffineExpr mlir::getAffineExprFromFlatForm(ArrayRef<int64_t> flatExprs,
|
|
unsigned numDims,
|
|
unsigned numSymbols,
|
|
ArrayRef<AffineExpr> localExprs,
|
|
MLIRContext *context) {
|
|
// Assert expected numLocals = flatExprs.size() - numDims - numSymbols - 1.
|
|
assert(flatExprs.size() - numDims - numSymbols - 1 == localExprs.size() &&
|
|
"unexpected number of local expressions");
|
|
|
|
auto expr = getAffineConstantExpr(0, context);
|
|
// Dimensions and symbols.
|
|
for (unsigned j = 0; j < numDims + numSymbols; j++) {
|
|
if (flatExprs[j] == 0)
|
|
continue;
|
|
auto id = j < numDims ? getAffineDimExpr(j, context)
|
|
: getAffineSymbolExpr(j - numDims, context);
|
|
expr = expr + id * flatExprs[j];
|
|
}
|
|
|
|
// Local identifiers.
|
|
for (unsigned j = numDims + numSymbols, e = flatExprs.size() - 1; j < e;
|
|
j++) {
|
|
if (flatExprs[j] == 0)
|
|
continue;
|
|
auto term = localExprs[j - numDims - numSymbols] * flatExprs[j];
|
|
expr = expr + term;
|
|
}
|
|
|
|
// Constant term.
|
|
int64_t constTerm = flatExprs[flatExprs.size() - 1];
|
|
if (constTerm != 0)
|
|
expr = expr + constTerm;
|
|
return expr;
|
|
}
|
|
|
|
/// Constructs a semi-affine expression from a flat ArrayRef. If there are
|
|
/// local identifiers (neither dimensional nor symbolic) that appear in the sum
|
|
/// of products expression, `localExprs` is expected to have the AffineExprs for
|
|
/// it, and is substituted into. The ArrayRef `flatExprs` is expected to be in
|
|
/// the format [dims, symbols, locals, constant term]. The semi-affine
|
|
/// expression is constructed in the sorted order of dimension and symbol
|
|
/// position numbers. Note: local expressions/ids are used for mod, div as well
|
|
/// as symbolic RHS terms for terms that are not pure affine.
|
|
static AffineExpr getSemiAffineExprFromFlatForm(ArrayRef<int64_t> flatExprs,
|
|
unsigned numDims,
|
|
unsigned numSymbols,
|
|
ArrayRef<AffineExpr> localExprs,
|
|
MLIRContext *context) {
|
|
assert(!flatExprs.empty() && "flatExprs cannot be empty");
|
|
|
|
// Assert expected numLocals = flatExprs.size() - numDims - numSymbols - 1.
|
|
assert(flatExprs.size() - numDims - numSymbols - 1 == localExprs.size() &&
|
|
"unexpected number of local expressions");
|
|
|
|
AffineExpr expr = getAffineConstantExpr(0, context);
|
|
|
|
// We design indices as a pair which help us present the semi-affine map as
|
|
// sum of product where terms are sorted based on dimension or symbol
|
|
// position: <keyA, keyB> for expressions of the form dimension * symbol,
|
|
// where keyA is the position number of the dimension and keyB is the
|
|
// position number of the symbol. For dimensional expressions we set the index
|
|
// as (position number of the dimension, -1), as we want dimensional
|
|
// expressions to appear before symbolic and product of dimensional and
|
|
// symbolic expressions having the dimension with the same position number.
|
|
// For symbolic expression set the index as (position number of the symbol,
|
|
// maximum of last dimension and symbol position) number. For example, we want
|
|
// the expression we are constructing to look something like: d0 + d0 * s0 +
|
|
// s0 + d1*s1 + s1.
|
|
|
|
// Stores the affine expression corresponding to a given index.
|
|
DenseMap<std::pair<unsigned, signed>, AffineExpr> indexToExprMap;
|
|
// Stores the constant coefficient value corresponding to a given
|
|
// dimension, symbol or a non-pure affine expression stored in `localExprs`.
|
|
DenseMap<std::pair<unsigned, signed>, int64_t> coefficients;
|
|
// Stores the indices as defined above, and later sorted to produce
|
|
// the semi-affine expression in the desired form.
|
|
SmallVector<std::pair<unsigned, signed>, 8> indices;
|
|
|
|
// Example: expression = d0 + d0 * s0 + 2 * s0.
|
|
// indices = [{0,-1}, {0, 0}, {0, 1}]
|
|
// coefficients = [{{0, -1}, 1}, {{0, 0}, 1}, {{0, 1}, 2}]
|
|
// indexToExprMap = [{{0, -1}, d0}, {{0, 0}, d0 * s0}, {{0, 1}, s0}]
|
|
|
|
// Adds entries to `indexToExprMap`, `coefficients` and `indices`.
|
|
auto addEntry = [&](std::pair<unsigned, signed> index, int64_t coefficient,
|
|
AffineExpr expr) {
|
|
assert(!llvm::is_contained(indices, index) &&
|
|
"Key is already present in indices vector and overwriting will "
|
|
"happen in `indexToExprMap` and `coefficients`!");
|
|
|
|
indices.push_back(index);
|
|
coefficients.insert({index, coefficient});
|
|
indexToExprMap.insert({index, expr});
|
|
};
|
|
|
|
// Design indices for dimensional or symbolic terms, and store the indices,
|
|
// constant coefficient corresponding to the indices in `coefficients` map,
|
|
// and affine expression corresponding to indices in `indexToExprMap` map.
|
|
|
|
// Ensure we do not have duplicate keys in `indexToExpr` map.
|
|
unsigned offsetSym = 0;
|
|
signed offsetDim = -1;
|
|
for (unsigned j = numDims; j < numDims + numSymbols; ++j) {
|
|
if (flatExprs[j] == 0)
|
|
continue;
|
|
// For symbolic expression set the index as <position number
|
|
// of the symbol, max(dimCount, symCount)> number,
|
|
// as we want symbolic expressions with the same positional number to
|
|
// appear after dimensional expressions having the same positional number.
|
|
std::pair<unsigned, signed> indexEntry(
|
|
j - numDims, std::max(numDims, numSymbols) + offsetSym++);
|
|
addEntry(indexEntry, flatExprs[j],
|
|
getAffineSymbolExpr(j - numDims, context));
|
|
}
|
|
|
|
// Denotes semi-affine product, modulo or division terms, which has been added
|
|
// to the `indexToExpr` map.
|
|
SmallVector<bool, 4> addedToMap(flatExprs.size() - numDims - numSymbols - 1,
|
|
false);
|
|
unsigned lhsPos, rhsPos;
|
|
// Construct indices for product terms involving dimension, symbol or constant
|
|
// as lhs/rhs, and store the indices, constant coefficient corresponding to
|
|
// the indices in `coefficients` map, and affine expression corresponding to
|
|
// in indices in `indexToExprMap` map.
|
|
for (const auto &it : llvm::enumerate(localExprs)) {
|
|
AffineExpr expr = it.value();
|
|
if (flatExprs[numDims + numSymbols + it.index()] == 0)
|
|
continue;
|
|
AffineExpr lhs = cast<AffineBinaryOpExpr>(expr).getLHS();
|
|
AffineExpr rhs = cast<AffineBinaryOpExpr>(expr).getRHS();
|
|
if (!((isa<AffineDimExpr>(lhs) || isa<AffineSymbolExpr>(lhs)) &&
|
|
(isa<AffineDimExpr>(rhs) || isa<AffineSymbolExpr>(rhs) ||
|
|
isa<AffineConstantExpr>(rhs)))) {
|
|
continue;
|
|
}
|
|
if (isa<AffineConstantExpr>(rhs)) {
|
|
// For product/modulo/division expressions, when rhs of modulo/division
|
|
// expression is constant, we put 0 in place of keyB, because we want
|
|
// them to appear earlier in the semi-affine expression we are
|
|
// constructing. When rhs is constant, we place 0 in place of keyB.
|
|
if (isa<AffineDimExpr>(lhs)) {
|
|
lhsPos = cast<AffineDimExpr>(lhs).getPosition();
|
|
std::pair<unsigned, signed> indexEntry(lhsPos, offsetDim--);
|
|
addEntry(indexEntry, flatExprs[numDims + numSymbols + it.index()],
|
|
expr);
|
|
} else {
|
|
lhsPos = cast<AffineSymbolExpr>(lhs).getPosition();
|
|
std::pair<unsigned, signed> indexEntry(
|
|
lhsPos, std::max(numDims, numSymbols) + offsetSym++);
|
|
addEntry(indexEntry, flatExprs[numDims + numSymbols + it.index()],
|
|
expr);
|
|
}
|
|
} else if (isa<AffineDimExpr>(lhs)) {
|
|
// For product/modulo/division expressions having lhs as dimension and rhs
|
|
// as symbol, we order the terms in the semi-affine expression based on
|
|
// the pair: <keyA, keyB> for expressions of the form dimension * symbol,
|
|
// where keyA is the position number of the dimension and keyB is the
|
|
// position number of the symbol.
|
|
lhsPos = cast<AffineDimExpr>(lhs).getPosition();
|
|
rhsPos = cast<AffineSymbolExpr>(rhs).getPosition();
|
|
std::pair<unsigned, signed> indexEntry(lhsPos, rhsPos);
|
|
addEntry(indexEntry, flatExprs[numDims + numSymbols + it.index()], expr);
|
|
} else {
|
|
// For product/modulo/division expressions having both lhs and rhs as
|
|
// symbol, we design indices as a pair: <keyA, keyB> for expressions
|
|
// of the form dimension * symbol, where keyA is the position number of
|
|
// the dimension and keyB is the position number of the symbol.
|
|
lhsPos = cast<AffineSymbolExpr>(lhs).getPosition();
|
|
rhsPos = cast<AffineSymbolExpr>(rhs).getPosition();
|
|
std::pair<unsigned, signed> indexEntry(
|
|
lhsPos, std::max(numDims, numSymbols) + offsetSym++);
|
|
addEntry(indexEntry, flatExprs[numDims + numSymbols + it.index()], expr);
|
|
}
|
|
addedToMap[it.index()] = true;
|
|
}
|
|
|
|
for (unsigned j = 0; j < numDims; ++j) {
|
|
if (flatExprs[j] == 0)
|
|
continue;
|
|
// For dimensional expressions we set the index as <position number of the
|
|
// dimension, 0>, as we want dimensional expressions to appear before
|
|
// symbolic ones and products of dimensional and symbolic expressions
|
|
// having the dimension with the same position number.
|
|
std::pair<unsigned, signed> indexEntry(j, offsetDim--);
|
|
addEntry(indexEntry, flatExprs[j], getAffineDimExpr(j, context));
|
|
}
|
|
|
|
// Constructing the simplified semi-affine sum of product/division/mod
|
|
// expression from the flattened form in the desired sorted order of indices
|
|
// of the various individual product/division/mod expressions.
|
|
llvm::sort(indices);
|
|
for (const std::pair<unsigned, unsigned> index : indices) {
|
|
assert(indexToExprMap.lookup(index) &&
|
|
"cannot find key in `indexToExprMap` map");
|
|
expr = expr + indexToExprMap.lookup(index) * coefficients.lookup(index);
|
|
}
|
|
|
|
// Local identifiers.
|
|
for (unsigned j = numDims + numSymbols, e = flatExprs.size() - 1; j < e;
|
|
j++) {
|
|
// If the coefficient of the local expression is 0, continue as we need not
|
|
// add it in out final expression.
|
|
if (flatExprs[j] == 0 || addedToMap[j - numDims - numSymbols])
|
|
continue;
|
|
auto term = localExprs[j - numDims - numSymbols] * flatExprs[j];
|
|
expr = expr + term;
|
|
}
|
|
|
|
// Constant term.
|
|
int64_t constTerm = flatExprs.back();
|
|
if (constTerm != 0)
|
|
expr = expr + constTerm;
|
|
return expr;
|
|
}
|
|
|
|
SimpleAffineExprFlattener::SimpleAffineExprFlattener(unsigned numDims,
|
|
unsigned numSymbols)
|
|
: numDims(numDims), numSymbols(numSymbols), numLocals(0) {
|
|
operandExprStack.reserve(8);
|
|
}
|
|
|
|
// In pure affine t = expr * c, we multiply each coefficient of lhs with c.
|
|
//
|
|
// In case of semi affine multiplication expressions, t = expr * symbolic_expr,
|
|
// introduce a local variable p (= expr * symbolic_expr), and the affine
|
|
// expression expr * symbolic_expr is added to `localExprs`.
|
|
LogicalResult SimpleAffineExprFlattener::visitMulExpr(AffineBinaryOpExpr expr) {
|
|
assert(operandExprStack.size() >= 2);
|
|
SmallVector<int64_t, 8> rhs = operandExprStack.back();
|
|
operandExprStack.pop_back();
|
|
SmallVector<int64_t, 8> &lhs = operandExprStack.back();
|
|
|
|
// Flatten semi-affine multiplication expressions by introducing a local
|
|
// variable in place of the product; the affine expression
|
|
// corresponding to the quantifier is added to `localExprs`.
|
|
if (!isa<AffineConstantExpr>(expr.getRHS())) {
|
|
SmallVector<int64_t, 8> mulLhs(lhs);
|
|
MLIRContext *context = expr.getContext();
|
|
AffineExpr a = getAffineExprFromFlatForm(lhs, numDims, numSymbols,
|
|
localExprs, context);
|
|
AffineExpr b = getAffineExprFromFlatForm(rhs, numDims, numSymbols,
|
|
localExprs, context);
|
|
return addLocalVariableSemiAffine(mulLhs, rhs, a * b, lhs, lhs.size());
|
|
}
|
|
|
|
// Get the RHS constant.
|
|
int64_t rhsConst = rhs[getConstantIndex()];
|
|
for (int64_t &lhsElt : lhs)
|
|
lhsElt *= rhsConst;
|
|
|
|
return success();
|
|
}
|
|
|
|
LogicalResult SimpleAffineExprFlattener::visitAddExpr(AffineBinaryOpExpr expr) {
|
|
assert(operandExprStack.size() >= 2);
|
|
const auto &rhs = operandExprStack.back();
|
|
auto &lhs = operandExprStack[operandExprStack.size() - 2];
|
|
assert(lhs.size() == rhs.size());
|
|
// Update the LHS in place.
|
|
for (unsigned i = 0, e = rhs.size(); i < e; i++) {
|
|
lhs[i] += rhs[i];
|
|
}
|
|
// Pop off the RHS.
|
|
operandExprStack.pop_back();
|
|
return success();
|
|
}
|
|
|
|
//
|
|
// t = expr mod c <=> t = expr - c*q and c*q <= expr <= c*q + c - 1
|
|
//
|
|
// A mod expression "expr mod c" is thus flattened by introducing a new local
|
|
// variable q (= expr floordiv c), such that expr mod c is replaced with
|
|
// 'expr - c * q' and c * q <= expr <= c * q + c - 1 are added to localVarCst.
|
|
//
|
|
// In case of semi-affine modulo expressions, t = expr mod symbolic_expr,
|
|
// introduce a local variable m (= expr mod symbolic_expr), and the affine
|
|
// expression expr mod symbolic_expr is added to `localExprs`.
|
|
LogicalResult SimpleAffineExprFlattener::visitModExpr(AffineBinaryOpExpr expr) {
|
|
assert(operandExprStack.size() >= 2);
|
|
|
|
SmallVector<int64_t, 8> rhs = operandExprStack.back();
|
|
operandExprStack.pop_back();
|
|
SmallVector<int64_t, 8> &lhs = operandExprStack.back();
|
|
MLIRContext *context = expr.getContext();
|
|
|
|
// Flatten semi affine modulo expressions by introducing a local
|
|
// variable in place of the modulo value, and the affine expression
|
|
// corresponding to the quantifier is added to `localExprs`.
|
|
if (!isa<AffineConstantExpr>(expr.getRHS())) {
|
|
SmallVector<int64_t, 8> modLhs(lhs);
|
|
AffineExpr dividendExpr = getAffineExprFromFlatForm(
|
|
lhs, numDims, numSymbols, localExprs, context);
|
|
AffineExpr divisorExpr = getAffineExprFromFlatForm(rhs, numDims, numSymbols,
|
|
localExprs, context);
|
|
AffineExpr modExpr = dividendExpr % divisorExpr;
|
|
return addLocalVariableSemiAffine(modLhs, rhs, modExpr, lhs, lhs.size());
|
|
}
|
|
|
|
int64_t rhsConst = rhs[getConstantIndex()];
|
|
if (rhsConst <= 0)
|
|
return failure();
|
|
|
|
// Check if the LHS expression is a multiple of modulo factor.
|
|
unsigned i, e;
|
|
for (i = 0, e = lhs.size(); i < e; i++)
|
|
if (lhs[i] % rhsConst != 0)
|
|
break;
|
|
// If yes, modulo expression here simplifies to zero.
|
|
if (i == lhs.size()) {
|
|
std::fill(lhs.begin(), lhs.end(), 0);
|
|
return success();
|
|
}
|
|
|
|
// Add a local variable for the quotient, i.e., expr % c is replaced by
|
|
// (expr - q * c) where q = expr floordiv c. Do this while canceling out
|
|
// the GCD of expr and c.
|
|
SmallVector<int64_t, 8> floorDividend(lhs);
|
|
uint64_t gcd = rhsConst;
|
|
for (int64_t lhsElt : lhs)
|
|
gcd = std::gcd(gcd, (uint64_t)std::abs(lhsElt));
|
|
// Simplify the numerator and the denominator.
|
|
if (gcd != 1) {
|
|
for (int64_t &floorDividendElt : floorDividend)
|
|
floorDividendElt = floorDividendElt / static_cast<int64_t>(gcd);
|
|
}
|
|
int64_t floorDivisor = rhsConst / static_cast<int64_t>(gcd);
|
|
|
|
// Construct the AffineExpr form of the floordiv to store in localExprs.
|
|
|
|
AffineExpr dividendExpr = getAffineExprFromFlatForm(
|
|
floorDividend, numDims, numSymbols, localExprs, context);
|
|
AffineExpr divisorExpr = getAffineConstantExpr(floorDivisor, context);
|
|
AffineExpr floorDivExpr = dividendExpr.floorDiv(divisorExpr);
|
|
int loc;
|
|
if ((loc = findLocalId(floorDivExpr)) == -1) {
|
|
addLocalFloorDivId(floorDividend, floorDivisor, floorDivExpr);
|
|
// Set result at top of stack to "lhs - rhsConst * q".
|
|
lhs[getLocalVarStartIndex() + numLocals - 1] = -rhsConst;
|
|
} else {
|
|
// Reuse the existing local id.
|
|
lhs[getLocalVarStartIndex() + loc] -= rhsConst;
|
|
}
|
|
return success();
|
|
}
|
|
|
|
LogicalResult
|
|
SimpleAffineExprFlattener::visitCeilDivExpr(AffineBinaryOpExpr expr) {
|
|
return visitDivExpr(expr, /*isCeil=*/true);
|
|
}
|
|
LogicalResult
|
|
SimpleAffineExprFlattener::visitFloorDivExpr(AffineBinaryOpExpr expr) {
|
|
return visitDivExpr(expr, /*isCeil=*/false);
|
|
}
|
|
|
|
LogicalResult SimpleAffineExprFlattener::visitDimExpr(AffineDimExpr expr) {
|
|
operandExprStack.emplace_back(SmallVector<int64_t, 32>(getNumCols(), 0));
|
|
auto &eq = operandExprStack.back();
|
|
assert(expr.getPosition() < numDims && "Inconsistent number of dims");
|
|
eq[getDimStartIndex() + expr.getPosition()] = 1;
|
|
return success();
|
|
}
|
|
|
|
LogicalResult
|
|
SimpleAffineExprFlattener::visitSymbolExpr(AffineSymbolExpr expr) {
|
|
operandExprStack.emplace_back(SmallVector<int64_t, 32>(getNumCols(), 0));
|
|
auto &eq = operandExprStack.back();
|
|
assert(expr.getPosition() < numSymbols && "inconsistent number of symbols");
|
|
eq[getSymbolStartIndex() + expr.getPosition()] = 1;
|
|
return success();
|
|
}
|
|
|
|
LogicalResult
|
|
SimpleAffineExprFlattener::visitConstantExpr(AffineConstantExpr expr) {
|
|
operandExprStack.emplace_back(SmallVector<int64_t, 32>(getNumCols(), 0));
|
|
auto &eq = operandExprStack.back();
|
|
eq[getConstantIndex()] = expr.getValue();
|
|
return success();
|
|
}
|
|
|
|
LogicalResult SimpleAffineExprFlattener::addLocalVariableSemiAffine(
|
|
ArrayRef<int64_t> lhs, ArrayRef<int64_t> rhs, AffineExpr localExpr,
|
|
SmallVectorImpl<int64_t> &result, unsigned long resultSize) {
|
|
assert(result.size() == resultSize &&
|
|
"`result` vector passed is not of correct size");
|
|
int loc;
|
|
if ((loc = findLocalId(localExpr)) == -1) {
|
|
if (failed(addLocalIdSemiAffine(lhs, rhs, localExpr)))
|
|
return failure();
|
|
}
|
|
std::fill(result.begin(), result.end(), 0);
|
|
if (loc == -1)
|
|
result[getLocalVarStartIndex() + numLocals - 1] = 1;
|
|
else
|
|
result[getLocalVarStartIndex() + loc] = 1;
|
|
return success();
|
|
}
|
|
|
|
// t = expr floordiv c <=> t = q, c * q <= expr <= c * q + c - 1
|
|
// A floordiv is thus flattened by introducing a new local variable q, and
|
|
// replacing that expression with 'q' while adding the constraints
|
|
// c * q <= expr <= c * q + c - 1 to localVarCst (done by
|
|
// IntegerRelation::addLocalFloorDiv).
|
|
//
|
|
// A ceildiv is similarly flattened:
|
|
// t = expr ceildiv c <=> t = (expr + c - 1) floordiv c
|
|
//
|
|
// In case of semi affine division expressions, t = expr floordiv symbolic_expr
|
|
// or t = expr ceildiv symbolic_expr, introduce a local variable q (= expr
|
|
// floordiv/ceildiv symbolic_expr), and the affine floordiv/ceildiv is added to
|
|
// `localExprs`.
|
|
LogicalResult SimpleAffineExprFlattener::visitDivExpr(AffineBinaryOpExpr expr,
|
|
bool isCeil) {
|
|
assert(operandExprStack.size() >= 2);
|
|
|
|
MLIRContext *context = expr.getContext();
|
|
SmallVector<int64_t, 8> rhs = operandExprStack.back();
|
|
operandExprStack.pop_back();
|
|
SmallVector<int64_t, 8> &lhs = operandExprStack.back();
|
|
|
|
// Flatten semi affine division expressions by introducing a local
|
|
// variable in place of the quotient, and the affine expression corresponding
|
|
// to the quantifier is added to `localExprs`.
|
|
if (!isa<AffineConstantExpr>(expr.getRHS())) {
|
|
SmallVector<int64_t, 8> divLhs(lhs);
|
|
AffineExpr a = getAffineExprFromFlatForm(lhs, numDims, numSymbols,
|
|
localExprs, context);
|
|
AffineExpr b = getAffineExprFromFlatForm(rhs, numDims, numSymbols,
|
|
localExprs, context);
|
|
AffineExpr divExpr = isCeil ? a.ceilDiv(b) : a.floorDiv(b);
|
|
return addLocalVariableSemiAffine(divLhs, rhs, divExpr, lhs, lhs.size());
|
|
}
|
|
|
|
// This is a pure affine expr; the RHS is a positive constant.
|
|
int64_t rhsConst = rhs[getConstantIndex()];
|
|
if (rhsConst <= 0)
|
|
return failure();
|
|
|
|
// Simplify the floordiv, ceildiv if possible by canceling out the greatest
|
|
// common divisors of the numerator and denominator.
|
|
uint64_t gcd = std::abs(rhsConst);
|
|
for (int64_t lhsElt : lhs)
|
|
gcd = std::gcd(gcd, (uint64_t)std::abs(lhsElt));
|
|
// Simplify the numerator and the denominator.
|
|
if (gcd != 1) {
|
|
for (int64_t &lhsElt : lhs)
|
|
lhsElt = lhsElt / static_cast<int64_t>(gcd);
|
|
}
|
|
int64_t divisor = rhsConst / static_cast<int64_t>(gcd);
|
|
// If the divisor becomes 1, the updated LHS is the result. (The
|
|
// divisor can't be negative since rhsConst is positive).
|
|
if (divisor == 1)
|
|
return success();
|
|
|
|
// If the divisor cannot be simplified to one, we will have to retain
|
|
// the ceil/floor expr (simplified up until here). Add an existential
|
|
// quantifier to express its result, i.e., expr1 div expr2 is replaced
|
|
// by a new identifier, q.
|
|
AffineExpr a =
|
|
getAffineExprFromFlatForm(lhs, numDims, numSymbols, localExprs, context);
|
|
AffineExpr b = getAffineConstantExpr(divisor, context);
|
|
|
|
int loc;
|
|
AffineExpr divExpr = isCeil ? a.ceilDiv(b) : a.floorDiv(b);
|
|
if ((loc = findLocalId(divExpr)) == -1) {
|
|
if (!isCeil) {
|
|
SmallVector<int64_t, 8> dividend(lhs);
|
|
addLocalFloorDivId(dividend, divisor, divExpr);
|
|
} else {
|
|
// lhs ceildiv c <=> (lhs + c - 1) floordiv c
|
|
SmallVector<int64_t, 8> dividend(lhs);
|
|
dividend.back() += divisor - 1;
|
|
addLocalFloorDivId(dividend, divisor, divExpr);
|
|
}
|
|
}
|
|
// Set the expression on stack to the local var introduced to capture the
|
|
// result of the division (floor or ceil).
|
|
std::fill(lhs.begin(), lhs.end(), 0);
|
|
if (loc == -1)
|
|
lhs[getLocalVarStartIndex() + numLocals - 1] = 1;
|
|
else
|
|
lhs[getLocalVarStartIndex() + loc] = 1;
|
|
return success();
|
|
}
|
|
|
|
// Add a local identifier (needed to flatten a mod, floordiv, ceildiv expr).
|
|
// The local identifier added is always a floordiv of a pure add/mul affine
|
|
// function of other identifiers, coefficients of which are specified in
|
|
// dividend and with respect to a positive constant divisor. localExpr is the
|
|
// simplified tree expression (AffineExpr) corresponding to the quantifier.
|
|
void SimpleAffineExprFlattener::addLocalFloorDivId(ArrayRef<int64_t> dividend,
|
|
int64_t divisor,
|
|
AffineExpr localExpr) {
|
|
assert(divisor > 0 && "positive constant divisor expected");
|
|
for (SmallVector<int64_t, 8> &subExpr : operandExprStack)
|
|
subExpr.insert(subExpr.begin() + getLocalVarStartIndex() + numLocals, 0);
|
|
localExprs.push_back(localExpr);
|
|
numLocals++;
|
|
// dividend and divisor are not used here; an override of this method uses it.
|
|
}
|
|
|
|
LogicalResult SimpleAffineExprFlattener::addLocalIdSemiAffine(
|
|
ArrayRef<int64_t> lhs, ArrayRef<int64_t> rhs, AffineExpr localExpr) {
|
|
for (SmallVector<int64_t, 8> &subExpr : operandExprStack)
|
|
subExpr.insert(subExpr.begin() + getLocalVarStartIndex() + numLocals, 0);
|
|
localExprs.push_back(localExpr);
|
|
++numLocals;
|
|
// lhs and rhs are not used here; an override of this method uses them.
|
|
return success();
|
|
}
|
|
|
|
int SimpleAffineExprFlattener::findLocalId(AffineExpr localExpr) {
|
|
SmallVectorImpl<AffineExpr>::iterator it;
|
|
if ((it = llvm::find(localExprs, localExpr)) == localExprs.end())
|
|
return -1;
|
|
return it - localExprs.begin();
|
|
}
|
|
|
|
/// Simplify the affine expression by flattening it and reconstructing it.
|
|
AffineExpr mlir::simplifyAffineExpr(AffineExpr expr, unsigned numDims,
|
|
unsigned numSymbols) {
|
|
// Simplify semi-affine expressions separately.
|
|
if (!expr.isPureAffine())
|
|
expr = simplifySemiAffine(expr, numDims, numSymbols);
|
|
|
|
SimpleAffineExprFlattener flattener(numDims, numSymbols);
|
|
// has poison expression
|
|
if (failed(flattener.walkPostOrder(expr)))
|
|
return expr;
|
|
ArrayRef<int64_t> flattenedExpr = flattener.operandExprStack.back();
|
|
if (!expr.isPureAffine() &&
|
|
expr == getAffineExprFromFlatForm(flattenedExpr, numDims, numSymbols,
|
|
flattener.localExprs,
|
|
expr.getContext()))
|
|
return expr;
|
|
AffineExpr simplifiedExpr =
|
|
expr.isPureAffine()
|
|
? getAffineExprFromFlatForm(flattenedExpr, numDims, numSymbols,
|
|
flattener.localExprs, expr.getContext())
|
|
: getSemiAffineExprFromFlatForm(flattenedExpr, numDims, numSymbols,
|
|
flattener.localExprs,
|
|
expr.getContext());
|
|
|
|
flattener.operandExprStack.pop_back();
|
|
assert(flattener.operandExprStack.empty());
|
|
return simplifiedExpr;
|
|
}
|
|
|
|
std::optional<int64_t> mlir::getBoundForAffineExpr(
|
|
AffineExpr expr, unsigned numDims, unsigned numSymbols,
|
|
ArrayRef<std::optional<int64_t>> constLowerBounds,
|
|
ArrayRef<std::optional<int64_t>> constUpperBounds, bool isUpper) {
|
|
// Handle divs and mods.
|
|
if (auto binOpExpr = dyn_cast<AffineBinaryOpExpr>(expr)) {
|
|
// If the LHS of a floor or ceil is bounded and the RHS is a constant, we
|
|
// can compute an upper bound.
|
|
if (binOpExpr.getKind() == AffineExprKind::FloorDiv) {
|
|
auto rhsConst = dyn_cast<AffineConstantExpr>(binOpExpr.getRHS());
|
|
if (!rhsConst || rhsConst.getValue() < 1)
|
|
return std::nullopt;
|
|
auto bound =
|
|
getBoundForAffineExpr(binOpExpr.getLHS(), numDims, numSymbols,
|
|
constLowerBounds, constUpperBounds, isUpper);
|
|
if (!bound)
|
|
return std::nullopt;
|
|
return divideFloorSigned(*bound, rhsConst.getValue());
|
|
}
|
|
if (binOpExpr.getKind() == AffineExprKind::CeilDiv) {
|
|
auto rhsConst = dyn_cast<AffineConstantExpr>(binOpExpr.getRHS());
|
|
if (rhsConst && rhsConst.getValue() >= 1) {
|
|
auto bound =
|
|
getBoundForAffineExpr(binOpExpr.getLHS(), numDims, numSymbols,
|
|
constLowerBounds, constUpperBounds, isUpper);
|
|
if (!bound)
|
|
return std::nullopt;
|
|
return divideCeilSigned(*bound, rhsConst.getValue());
|
|
}
|
|
return std::nullopt;
|
|
}
|
|
if (binOpExpr.getKind() == AffineExprKind::Mod) {
|
|
// lhs mod c is always <= c - 1 and non-negative. In addition, if `lhs` is
|
|
// bounded such that lb <= lhs <= ub and lb floordiv c == ub floordiv c
|
|
// (same "interval"), then lb mod c <= lhs mod c <= ub mod c.
|
|
auto rhsConst = dyn_cast<AffineConstantExpr>(binOpExpr.getRHS());
|
|
if (rhsConst && rhsConst.getValue() >= 1) {
|
|
int64_t rhsConstVal = rhsConst.getValue();
|
|
auto lb = getBoundForAffineExpr(binOpExpr.getLHS(), numDims, numSymbols,
|
|
constLowerBounds, constUpperBounds,
|
|
/*isUpper=*/false);
|
|
auto ub =
|
|
getBoundForAffineExpr(binOpExpr.getLHS(), numDims, numSymbols,
|
|
constLowerBounds, constUpperBounds, isUpper);
|
|
if (ub && lb &&
|
|
divideFloorSigned(*lb, rhsConstVal) ==
|
|
divideFloorSigned(*ub, rhsConstVal))
|
|
return isUpper ? mod(*ub, rhsConstVal) : mod(*lb, rhsConstVal);
|
|
return isUpper ? rhsConstVal - 1 : 0;
|
|
}
|
|
}
|
|
}
|
|
// Flatten the expression.
|
|
SimpleAffineExprFlattener flattener(numDims, numSymbols);
|
|
auto simpleResult = flattener.walkPostOrder(expr);
|
|
// has poison expression
|
|
if (failed(simpleResult))
|
|
return std::nullopt;
|
|
ArrayRef<int64_t> flattenedExpr = flattener.operandExprStack.back();
|
|
// TODO: Handle local variables. We can get hold of flattener.localExprs and
|
|
// get bound on the local expr recursively.
|
|
if (flattener.numLocals > 0)
|
|
return std::nullopt;
|
|
int64_t bound = 0;
|
|
// Substitute the constant lower or upper bound for the dimensional or
|
|
// symbolic input depending on `isUpper` to determine the bound.
|
|
for (unsigned i = 0, e = numDims + numSymbols; i < e; ++i) {
|
|
if (flattenedExpr[i] > 0) {
|
|
auto &constBound = isUpper ? constUpperBounds[i] : constLowerBounds[i];
|
|
if (!constBound)
|
|
return std::nullopt;
|
|
bound += *constBound * flattenedExpr[i];
|
|
} else if (flattenedExpr[i] < 0) {
|
|
auto &constBound = isUpper ? constLowerBounds[i] : constUpperBounds[i];
|
|
if (!constBound)
|
|
return std::nullopt;
|
|
bound += *constBound * flattenedExpr[i];
|
|
}
|
|
}
|
|
// Constant term.
|
|
bound += flattenedExpr.back();
|
|
return bound;
|
|
}
|