17 #include "llvm/Support/Debug.h"
18 #include "llvm/Support/DebugLog.h"
19 #include "llvm/Support/InterleavedRange.h"
24 #define DEBUG_TYPE "decompose-affine-ops"
32 while (
auto loopOp = currentOp->
getParentOfType<LoopLikeOpInterface>()) {
33 if (!loopOp.isDefinedOutsideOfLoop(operand.
get()))
44 llvm::map_range(op->getOpOperands(), [&](
OpOperand &operand) {
45 return numEnclosingInvariantLoops(operand);
48 int64_t numOperands = op.getNumOperands();
50 llvm::to_vector(llvm::seq<int64_t>(0, numOperands));
51 llvm::stable_sort(operandPositions, [&numInvariant](
size_t i1,
size_t i2) {
52 return numInvariant[i1] > numInvariant[i2];
57 for (int64_t i = 0; i < numOperands; ++i) {
58 operands[i] = op.getOperand(operandPositions[i]);
75 op->setOperands(operands);
91 return AffineApplyOp::create(rewriter, originalOp.getLoc(), rhsMap,
104 auto binExpr = dyn_cast<AffineBinaryOpExpr>(remainingExp);
108 if (!isa<AffineBinaryOpExpr>(binExpr.getLHS()) &&
109 !isa<AffineBinaryOpExpr>(binExpr.getRHS()))
116 op,
"only add or mul binary expr can be reassociated");
118 LDBG() <<
"Start decomposeIntoFinerGrainedOps: " << op;
125 auto currentBinExpr = dyn_cast<AffineBinaryOpExpr>(remainingExp);
126 if (!currentBinExpr || currentBinExpr.getKind() != binExpr.getKind()) {
127 subExpressions.push_back(remainingExp);
128 LDBG() <<
"--terminal: " << subExpressions.back();
131 subExpressions.push_back(currentBinExpr.getRHS());
132 LDBG() <<
"--subExpr: " << subExpressions.back();
133 remainingExp = currentBinExpr.getLHS();
140 auto getMaxSymbol = [&](
AffineExpr e) -> int64_t {
142 if (e.isFunctionOfSymbol(i))
147 return getMaxSymbol(e1) < getMaxSymbol(e2);
149 LDBG() <<
"--sorted subexprs: " << llvm::interleaved(subExpressions);
159 for (int64_t i = 1, e = subExpressions.size(); i < e; ++i) {
161 current = AffineApplyOp::create(rewriter, op.getLoc(), binMap,
163 LDBG() <<
"--reassociate into: " << current;
167 rewriter.
replaceOp(op, current.getResult());
static AffineApplyOp createSubApply(RewriterBase &rewriter, AffineApplyOp originalOp, AffineExpr expr)
Build an affine.apply that is a subexpression expr of originalOps affine map and with the same operan...
Base type for affine expression.
MLIRContext * getContext() const
A multi-dimensional affine map Affine map's are immutable like Type's, and they are uniqued.
static AffineMap get(MLIRContext *context)
Returns a zero result affine map with no dimensions or symbols: () -> ().
unsigned getNumSymbols() const
unsigned getNumDims() const
AffineMap replaceDimsAndSymbols(ArrayRef< AffineExpr > dimReplacements, ArrayRef< AffineExpr > symReplacements, unsigned numResultDims, unsigned numResultSyms) const
This method substitutes any uses of dimensions and symbols (e.g.
AffineExpr getResult(unsigned idx) const
IRValueT get() const
Return the current value being used by this operand.
MLIRContext is the top-level object for a collection of MLIR operations.
This class represents an operand of an operation.
Operation is the basic unit of execution within MLIR.
OpTy getParentOfType()
Return the closest surrounding parent operation that is of type 'OpTy'.
This class coordinates the application of a rewrite on a set of IR, providing a way for clients to tr...
std::enable_if_t<!std::is_convertible< CallbackT, Twine >::value, LogicalResult > notifyMatchFailure(Location loc, CallbackT &&reasonCallback)
Used to notify the listener that the IR failed to be rewritten because of a match failure,...
virtual void replaceOp(Operation *op, ValueRange newValues)
Replace the results of the given (original) operation with the specified list of values (replacements...
virtual void finalizeOpModification(Operation *op)
This method is used to signal the end of an in-place modification of the given operation.
virtual void startOpModification(Operation *op)
This method is used to notify the rewriter that an in-place operation modification is about to happen...
This class provides an abstraction over the different types of ranges over Values.
This class represents an instance of an SSA value in the MLIR system, representing a computable value...
Operation * getOwner() const
Return the owner of this operand.
void canonicalizeMapAndOperands(AffineMap *map, SmallVectorImpl< Value > *operands)
Modifies both map and operands in-place so as to:
FailureOr< AffineApplyOp > decompose(RewriterBase &rewriter, AffineApplyOp op)
Split an "affine.apply" operation into smaller ops.
int64_t numEnclosingInvariantLoops(OpOperand &operand)
Count the number of loops surrounding operand such that operand could be hoisted above.
void reorderOperandsByHoistability(RewriterBase &rewriter, AffineApplyOp op)
Helper function to rewrite op's affine map and reorder its operands such that they are in increasing ...
Include the generated interface declarations.
@ Mul
RHS of mul is always a constant or a symbolic expression.
AffineExpr getAffineBinaryOpExpr(AffineExprKind kind, AffineExpr lhs, AffineExpr rhs)
AffineExpr simplifyAffineExpr(AffineExpr expr, unsigned numDims, unsigned numSymbols)
Simplify an affine expression by flattening and some amount of simple analysis.
AffineExpr getAffineSymbolExpr(unsigned position, MLIRContext *context)