457 lines
20 KiB
C++
457 lines
20 KiB
C++
//===- AffineExprVisitor.h - MLIR AffineExpr Visitor Class ------*- C++ -*-===//
|
|
//
|
|
// 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
|
|
//
|
|
//===----------------------------------------------------------------------===//
|
|
//
|
|
// This file defines the AffineExpr visitor class.
|
|
//
|
|
//===----------------------------------------------------------------------===//
|
|
|
|
#ifndef MLIR_IR_AFFINEEXPRVISITOR_H
|
|
#define MLIR_IR_AFFINEEXPRVISITOR_H
|
|
|
|
#include "mlir/IR/AffineExpr.h"
|
|
#include "mlir/Support/LLVM.h"
|
|
#include "llvm/ADT/ArrayRef.h"
|
|
|
|
namespace mlir {
|
|
|
|
/// Base class for AffineExpr visitors/walkers.
|
|
///
|
|
/// AffineExpr visitors are used when you want to perform different actions
|
|
/// for different kinds of AffineExprs without having to use lots of casts
|
|
/// and a big switch instruction.
|
|
///
|
|
/// To define your own visitor, inherit from this class, specifying your
|
|
/// new type for the 'SubClass' template parameter, and "override" visitXXX
|
|
/// functions in your class. This class is defined in terms of statically
|
|
/// resolved overloading, not virtual functions.
|
|
///
|
|
/// The visitor is templated on its return type (`RetTy`). With a WalkResult
|
|
/// return type, the visitor supports interrupting walks.
|
|
///
|
|
/// For example, here is a visitor that counts the number of for AffineDimExprs
|
|
/// in an AffineExpr.
|
|
///
|
|
/// /// Declare the class. Note that we derive from AffineExprVisitor
|
|
/// /// instantiated with our new subclasses_ type.
|
|
///
|
|
/// struct DimExprCounter : public AffineExprVisitor<DimExprCounter> {
|
|
/// unsigned numDimExprs;
|
|
/// DimExprCounter() : numDimExprs(0) {}
|
|
/// void visitDimExpr(AffineDimExpr expr) { ++numDimExprs; }
|
|
/// };
|
|
///
|
|
/// And this class would be used like this:
|
|
/// DimExprCounter dec;
|
|
/// dec.visit(affineExpr);
|
|
/// numDimExprs = dec.numDimExprs;
|
|
///
|
|
/// AffineExprVisitor provides visit methods for the following binary affine
|
|
/// op expressions:
|
|
/// AffineBinaryAddOpExpr, AffineBinaryMulOpExpr,
|
|
/// AffineBinaryModOpExpr, AffineBinaryFloorDivOpExpr,
|
|
/// AffineBinaryCeilDivOpExpr. Note that default implementations of these
|
|
/// methods will call the general AffineBinaryOpExpr method.
|
|
///
|
|
/// In addition, visit methods are provided for the following affine
|
|
// expressions: AffineConstantExpr, AffineDimExpr, and
|
|
// AffineSymbolExpr.
|
|
///
|
|
/// Note that if you don't implement visitXXX for some affine expression type,
|
|
/// the visitXXX method for Instruction superclass will be invoked.
|
|
///
|
|
/// Note that this class is specifically designed as a template to avoid
|
|
/// virtual function call overhead. Defining and using a AffineExprVisitor is
|
|
/// just as efficient as having your own switch instruction over the instruction
|
|
/// opcode.
|
|
template <typename SubClass, typename RetTy>
|
|
class AffineExprVisitorBase {
|
|
public:
|
|
// Function to visit an AffineExpr.
|
|
RetTy visit(AffineExpr expr) {
|
|
static_assert(std::is_base_of<AffineExprVisitorBase, SubClass>::value,
|
|
"Must instantiate with a derived type of AffineExprVisitor");
|
|
auto self = static_cast<SubClass *>(this);
|
|
switch (expr.getKind()) {
|
|
case AffineExprKind::Add: {
|
|
auto binOpExpr = cast<AffineBinaryOpExpr>(expr);
|
|
return self->visitAddExpr(binOpExpr);
|
|
}
|
|
case AffineExprKind::Mul: {
|
|
auto binOpExpr = cast<AffineBinaryOpExpr>(expr);
|
|
return self->visitMulExpr(binOpExpr);
|
|
}
|
|
case AffineExprKind::Mod: {
|
|
auto binOpExpr = cast<AffineBinaryOpExpr>(expr);
|
|
return self->visitModExpr(binOpExpr);
|
|
}
|
|
case AffineExprKind::FloorDiv: {
|
|
auto binOpExpr = cast<AffineBinaryOpExpr>(expr);
|
|
return self->visitFloorDivExpr(binOpExpr);
|
|
}
|
|
case AffineExprKind::CeilDiv: {
|
|
auto binOpExpr = cast<AffineBinaryOpExpr>(expr);
|
|
return self->visitCeilDivExpr(binOpExpr);
|
|
}
|
|
case AffineExprKind::Constant:
|
|
return self->visitConstantExpr(cast<AffineConstantExpr>(expr));
|
|
case AffineExprKind::DimId:
|
|
return self->visitDimExpr(cast<AffineDimExpr>(expr));
|
|
case AffineExprKind::SymbolId:
|
|
return self->visitSymbolExpr(cast<AffineSymbolExpr>(expr));
|
|
}
|
|
llvm_unreachable("Unknown AffineExpr");
|
|
}
|
|
|
|
//===--------------------------------------------------------------------===//
|
|
// Visitation functions... these functions provide default fallbacks in case
|
|
// the user does not specify what to do for a particular instruction type.
|
|
// The default behavior is to generalize the instruction type to its subtype
|
|
// and try visiting the subtype. All of this should be inlined perfectly,
|
|
// because there are no virtual functions to get in the way.
|
|
//
|
|
|
|
// Default visit methods. Note that the default op-specific binary op visit
|
|
// methods call the general visitAffineBinaryOpExpr visit method.
|
|
RetTy visitAffineBinaryOpExpr(AffineBinaryOpExpr expr) { return RetTy(); }
|
|
RetTy visitAddExpr(AffineBinaryOpExpr expr) {
|
|
return static_cast<SubClass *>(this)->visitAffineBinaryOpExpr(expr);
|
|
}
|
|
RetTy visitMulExpr(AffineBinaryOpExpr expr) {
|
|
return static_cast<SubClass *>(this)->visitAffineBinaryOpExpr(expr);
|
|
}
|
|
RetTy visitModExpr(AffineBinaryOpExpr expr) {
|
|
return static_cast<SubClass *>(this)->visitAffineBinaryOpExpr(expr);
|
|
}
|
|
RetTy visitFloorDivExpr(AffineBinaryOpExpr expr) {
|
|
return static_cast<SubClass *>(this)->visitAffineBinaryOpExpr(expr);
|
|
}
|
|
RetTy visitCeilDivExpr(AffineBinaryOpExpr expr) {
|
|
return static_cast<SubClass *>(this)->visitAffineBinaryOpExpr(expr);
|
|
}
|
|
RetTy visitConstantExpr(AffineConstantExpr expr) { return RetTy(); }
|
|
RetTy visitDimExpr(AffineDimExpr expr) { return RetTy(); }
|
|
RetTy visitSymbolExpr(AffineSymbolExpr expr) { return RetTy(); }
|
|
};
|
|
|
|
/// See documentation for AffineExprVisitorBase. This visitor supports
|
|
/// interrupting walks when a `WalkResult` is used for `RetTy`.
|
|
template <typename SubClass, typename RetTy = void>
|
|
class AffineExprVisitor : public AffineExprVisitorBase<SubClass, RetTy> {
|
|
//===--------------------------------------------------------------------===//
|
|
// Interface code - This is the public interface of the AffineExprVisitor
|
|
// that you use to visit affine expressions...
|
|
public:
|
|
// Function to walk an AffineExpr (in post order).
|
|
RetTy walkPostOrder(AffineExpr expr) {
|
|
static_assert(std::is_base_of<AffineExprVisitor, SubClass>::value,
|
|
"Must instantiate with a derived type of AffineExprVisitor");
|
|
auto self = static_cast<SubClass *>(this);
|
|
switch (expr.getKind()) {
|
|
case AffineExprKind::Add: {
|
|
auto binOpExpr = cast<AffineBinaryOpExpr>(expr);
|
|
if constexpr (std::is_same<RetTy, WalkResult>::value) {
|
|
if (walkOperandsPostOrder(binOpExpr).wasInterrupted())
|
|
return WalkResult::interrupt();
|
|
} else {
|
|
walkOperandsPostOrder(binOpExpr);
|
|
}
|
|
return self->visitAddExpr(binOpExpr);
|
|
}
|
|
case AffineExprKind::Mul: {
|
|
auto binOpExpr = cast<AffineBinaryOpExpr>(expr);
|
|
if constexpr (std::is_same<RetTy, WalkResult>::value) {
|
|
if (walkOperandsPostOrder(binOpExpr).wasInterrupted())
|
|
return WalkResult::interrupt();
|
|
} else {
|
|
walkOperandsPostOrder(binOpExpr);
|
|
}
|
|
return self->visitMulExpr(binOpExpr);
|
|
}
|
|
case AffineExprKind::Mod: {
|
|
auto binOpExpr = cast<AffineBinaryOpExpr>(expr);
|
|
if constexpr (std::is_same<RetTy, WalkResult>::value) {
|
|
if (walkOperandsPostOrder(binOpExpr).wasInterrupted())
|
|
return WalkResult::interrupt();
|
|
} else {
|
|
walkOperandsPostOrder(binOpExpr);
|
|
}
|
|
return self->visitModExpr(binOpExpr);
|
|
}
|
|
case AffineExprKind::FloorDiv: {
|
|
auto binOpExpr = cast<AffineBinaryOpExpr>(expr);
|
|
if constexpr (std::is_same<RetTy, WalkResult>::value) {
|
|
if (walkOperandsPostOrder(binOpExpr).wasInterrupted())
|
|
return WalkResult::interrupt();
|
|
} else {
|
|
walkOperandsPostOrder(binOpExpr);
|
|
}
|
|
return self->visitFloorDivExpr(binOpExpr);
|
|
}
|
|
case AffineExprKind::CeilDiv: {
|
|
auto binOpExpr = cast<AffineBinaryOpExpr>(expr);
|
|
if constexpr (std::is_same<RetTy, WalkResult>::value) {
|
|
if (walkOperandsPostOrder(binOpExpr).wasInterrupted())
|
|
return WalkResult::interrupt();
|
|
} else {
|
|
walkOperandsPostOrder(binOpExpr);
|
|
}
|
|
return self->visitCeilDivExpr(binOpExpr);
|
|
}
|
|
case AffineExprKind::Constant:
|
|
return self->visitConstantExpr(cast<AffineConstantExpr>(expr));
|
|
case AffineExprKind::DimId:
|
|
return self->visitDimExpr(cast<AffineDimExpr>(expr));
|
|
case AffineExprKind::SymbolId:
|
|
return self->visitSymbolExpr(cast<AffineSymbolExpr>(expr));
|
|
}
|
|
llvm_unreachable("Unknown AffineExpr");
|
|
}
|
|
|
|
private:
|
|
// Walk the operands - each operand is itself walked in post order.
|
|
RetTy walkOperandsPostOrder(AffineBinaryOpExpr expr) {
|
|
if constexpr (std::is_same<RetTy, WalkResult>::value) {
|
|
if (walkPostOrder(expr.getLHS()).wasInterrupted())
|
|
return WalkResult::interrupt();
|
|
} else {
|
|
walkPostOrder(expr.getLHS());
|
|
}
|
|
if constexpr (std::is_same<RetTy, WalkResult>::value) {
|
|
if (walkPostOrder(expr.getRHS()).wasInterrupted())
|
|
return WalkResult::interrupt();
|
|
return WalkResult::advance();
|
|
} else {
|
|
return walkPostOrder(expr.getRHS());
|
|
}
|
|
}
|
|
};
|
|
|
|
template <typename SubClass>
|
|
class AffineExprVisitor<SubClass, LogicalResult>
|
|
: public AffineExprVisitorBase<SubClass, LogicalResult> {
|
|
//===--------------------------------------------------------------------===//
|
|
// Interface code - This is the public interface of the AffineExprVisitor
|
|
// that you use to visit affine expressions...
|
|
public:
|
|
// Function to walk an AffineExpr (in post order).
|
|
LogicalResult walkPostOrder(AffineExpr expr) {
|
|
static_assert(std::is_base_of<AffineExprVisitor, SubClass>::value,
|
|
"Must instantiate with a derived type of AffineExprVisitor");
|
|
auto self = static_cast<SubClass *>(this);
|
|
switch (expr.getKind()) {
|
|
case AffineExprKind::Add: {
|
|
auto binOpExpr = cast<AffineBinaryOpExpr>(expr);
|
|
if (failed(walkOperandsPostOrder(binOpExpr)))
|
|
return failure();
|
|
return self->visitAddExpr(binOpExpr);
|
|
}
|
|
case AffineExprKind::Mul: {
|
|
auto binOpExpr = cast<AffineBinaryOpExpr>(expr);
|
|
if (failed(walkOperandsPostOrder(binOpExpr)))
|
|
return failure();
|
|
return self->visitMulExpr(binOpExpr);
|
|
}
|
|
case AffineExprKind::Mod: {
|
|
auto binOpExpr = cast<AffineBinaryOpExpr>(expr);
|
|
if (failed(walkOperandsPostOrder(binOpExpr)))
|
|
return failure();
|
|
return self->visitModExpr(binOpExpr);
|
|
}
|
|
case AffineExprKind::FloorDiv: {
|
|
auto binOpExpr = cast<AffineBinaryOpExpr>(expr);
|
|
if (failed(walkOperandsPostOrder(binOpExpr)))
|
|
return failure();
|
|
return self->visitFloorDivExpr(binOpExpr);
|
|
}
|
|
case AffineExprKind::CeilDiv: {
|
|
auto binOpExpr = cast<AffineBinaryOpExpr>(expr);
|
|
if (failed(walkOperandsPostOrder(binOpExpr)))
|
|
return failure();
|
|
return self->visitCeilDivExpr(binOpExpr);
|
|
}
|
|
case AffineExprKind::Constant:
|
|
return self->visitConstantExpr(cast<AffineConstantExpr>(expr));
|
|
case AffineExprKind::DimId:
|
|
return self->visitDimExpr(cast<AffineDimExpr>(expr));
|
|
case AffineExprKind::SymbolId:
|
|
return self->visitSymbolExpr(cast<AffineSymbolExpr>(expr));
|
|
}
|
|
llvm_unreachable("Unknown AffineExpr");
|
|
}
|
|
|
|
private:
|
|
// Walk the operands - each operand is itself walked in post order.
|
|
LogicalResult walkOperandsPostOrder(AffineBinaryOpExpr expr) {
|
|
if (failed(walkPostOrder(expr.getLHS())))
|
|
return failure();
|
|
if (failed(walkPostOrder(expr.getRHS())))
|
|
return failure();
|
|
return success();
|
|
}
|
|
};
|
|
|
|
// This class is used to flatten a pure affine expression (AffineExpr,
|
|
// which is in a tree form) into a sum of products (w.r.t constants) when
|
|
// possible, and in that process simplifying the expression. For a modulo,
|
|
// floordiv, or a ceildiv expression, an additional identifier, called a local
|
|
// identifier, is introduced to rewrite the expression as a sum of product
|
|
// affine expression. Each local identifier is always and by construction a
|
|
// floordiv of a pure add/mul affine function of dimensional, symbolic, and
|
|
// other local identifiers, in a non-mutually recursive way. Hence, every local
|
|
// identifier can ultimately always be recovered as an affine function of
|
|
// dimensional and symbolic identifiers (involving floordiv's); note however
|
|
// that by AffineExpr construction, some floordiv combinations are converted to
|
|
// mod's. The result of the flattening is a flattened expression and a set of
|
|
// constraints involving just the local variables.
|
|
//
|
|
// d2 + (d0 + d1) floordiv 4 is flattened to d2 + q where 'q' is the local
|
|
// variable introduced, with localVarCst containing 4*q <= d0 + d1 <= 4*q + 3.
|
|
//
|
|
// The simplification performed includes the accumulation of contributions for
|
|
// each dimensional and symbolic identifier together, the simplification of
|
|
// floordiv/ceildiv/mod expressions and other simplifications that in turn
|
|
// happen as a result. A simplification that this flattening naturally performs
|
|
// is of simplifying the numerator and denominator of floordiv/ceildiv, and
|
|
// folding a modulo expression to a zero, if possible. Three examples are below:
|
|
//
|
|
// (d0 + 3 * d1) + d0) - 2 * d1) - d0 simplified to d0 + d1
|
|
// (d0 - d0 mod 4 + 4) mod 4 simplified to 0
|
|
// (3*d0 + 2*d1 + d0) floordiv 2 + d1 simplified to 2*d0 + 2*d1
|
|
//
|
|
// The way the flattening works for the second example is as follows: d0 % 4 is
|
|
// replaced by d0 - 4*q with q being introduced: the expression then simplifies
|
|
// to: (d0 - (d0 - 4q) + 4) = 4q + 4, modulo of which w.r.t 4 simplifies to
|
|
// zero. Note that an affine expression may not always be expressible purely as
|
|
// a sum of products involving just the original dimensional and symbolic
|
|
// identifiers due to the presence of modulo/floordiv/ceildiv expressions that
|
|
// may not be eliminated after simplification; in such cases, the final
|
|
// expression can be reconstructed by replacing the local identifiers with their
|
|
// corresponding explicit form stored in 'localExprs' (note that each of the
|
|
// explicit forms itself would have been simplified).
|
|
//
|
|
// The expression walk method here performs a linear time post order walk that
|
|
// performs the above simplifications through visit methods, with partial
|
|
// results being stored in 'operandExprStack'. When a parent expr is visited,
|
|
// the flattened expressions corresponding to its two operands would already be
|
|
// on the stack - the parent expression looks at the two flattened expressions
|
|
// and combines the two. It pops off the operand expressions and pushes the
|
|
// combined result (although this is done in-place on its LHS operand expr).
|
|
// When the walk is completed, the flattened form of the top-level expression
|
|
// would be left on the stack.
|
|
//
|
|
// A flattener can be repeatedly used for multiple affine expressions that bind
|
|
// to the same operands, for example, for all result expressions of an
|
|
// AffineMap or AffineValueMap. In such cases, using it for multiple expressions
|
|
// is more efficient than creating a new flattener for each expression since
|
|
// common identical div and mod expressions appearing across different
|
|
// expressions are mapped to the same local identifier (same column position in
|
|
// 'localVarCst').
|
|
class SimpleAffineExprFlattener
|
|
: public AffineExprVisitor<SimpleAffineExprFlattener, LogicalResult> {
|
|
public:
|
|
// Flattened expression layout: [dims, symbols, locals, constant]
|
|
// Stack that holds the LHS and RHS operands while visiting a binary op expr.
|
|
// In the future, consider adding a prepass to determine how big the
|
|
// SmallVector's will be, and linearize this to std::vector<int64_t> to
|
|
// prevent SmallVector moves on re-allocation.
|
|
std::vector<SmallVector<int64_t, 8>> operandExprStack;
|
|
|
|
unsigned numDims;
|
|
unsigned numSymbols;
|
|
|
|
// Number of newly introduced identifiers to flatten mod/floordiv/ceildiv's.
|
|
unsigned numLocals;
|
|
|
|
// AffineExpr's corresponding to the floordiv/ceildiv/mod expressions for
|
|
// which new identifiers were introduced; if the latter do not get canceled
|
|
// out, these expressions can be readily used to reconstruct the AffineExpr
|
|
// (tree) form. Note that these expressions themselves would have been
|
|
// simplified (recursively) by this pass. Eg. d0 + (d0 + 2*d1 + d0) ceildiv 4
|
|
// will be simplified to d0 + q, where q = (d0 + d1) ceildiv 2. (d0 + d1)
|
|
// ceildiv 2 would be the local expression stored for q.
|
|
SmallVector<AffineExpr, 4> localExprs;
|
|
|
|
SimpleAffineExprFlattener(unsigned numDims, unsigned numSymbols);
|
|
|
|
virtual ~SimpleAffineExprFlattener() = default;
|
|
|
|
// Visitor method overrides.
|
|
LogicalResult visitMulExpr(AffineBinaryOpExpr expr);
|
|
LogicalResult visitAddExpr(AffineBinaryOpExpr expr);
|
|
LogicalResult visitDimExpr(AffineDimExpr expr);
|
|
LogicalResult visitSymbolExpr(AffineSymbolExpr expr);
|
|
LogicalResult visitConstantExpr(AffineConstantExpr expr);
|
|
LogicalResult visitCeilDivExpr(AffineBinaryOpExpr expr);
|
|
LogicalResult visitFloorDivExpr(AffineBinaryOpExpr expr);
|
|
|
|
//
|
|
// 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.
|
|
LogicalResult visitModExpr(AffineBinaryOpExpr expr);
|
|
|
|
protected:
|
|
// 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.
|
|
virtual void addLocalFloorDivId(ArrayRef<int64_t> dividend, int64_t divisor,
|
|
AffineExpr localExpr);
|
|
|
|
/// Add a local identifier (needed to flatten a mod, floordiv, ceildiv, mul
|
|
/// expr) when the rhs is a symbolic expression. The local identifier added
|
|
/// may be a floordiv, ceildiv, mul or mod of a pure affine/semi-affine
|
|
/// function of other identifiers, coefficients of which are specified in the
|
|
/// lhs of the mod, floordiv, ceildiv or mul expression and with respect to a
|
|
/// symbolic rhs expression. `localExpr` is the simplified tree expression
|
|
/// (AffineExpr) corresponding to the quantifier.
|
|
virtual LogicalResult addLocalIdSemiAffine(ArrayRef<int64_t> lhs,
|
|
ArrayRef<int64_t> rhs,
|
|
AffineExpr localExpr);
|
|
|
|
private:
|
|
/// Adds `localExpr`, which may be mod, ceildiv, floordiv or mod expression
|
|
/// representing the affine expression corresponding to the quantifier
|
|
/// introduced as the local variable corresponding to `localExpr`. If the
|
|
/// quantifier is already present, we put the coefficient in the proper index
|
|
/// of `result`, otherwise we add a new local variable and put the coefficient
|
|
/// there.
|
|
LogicalResult addLocalVariableSemiAffine(ArrayRef<int64_t> lhs,
|
|
ArrayRef<int64_t> rhs,
|
|
AffineExpr localExpr,
|
|
SmallVectorImpl<int64_t> &result,
|
|
unsigned long resultSize);
|
|
|
|
// 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
|
|
LogicalResult visitDivExpr(AffineBinaryOpExpr expr, bool isCeil);
|
|
|
|
int findLocalId(AffineExpr localExpr);
|
|
|
|
inline unsigned getNumCols() const {
|
|
return numDims + numSymbols + numLocals + 1;
|
|
}
|
|
inline unsigned getConstantIndex() const { return getNumCols() - 1; }
|
|
inline unsigned getLocalVarStartIndex() const { return numDims + numSymbols; }
|
|
inline unsigned getSymbolStartIndex() const { return numDims; }
|
|
inline unsigned getDimStartIndex() const { return 0; }
|
|
};
|
|
|
|
} // namespace mlir
|
|
|
|
#endif // MLIR_IR_AFFINEEXPRVISITOR_H
|