MLIR  22.0.0git
DecomposeAffineOps.cpp
Go to the documentation of this file.
1 //===- DecomposeAffineOps.cpp - Decompose affine ops into finer-grained ---===//
2 //
3 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4 // See https://llvm.org/LICENSE.txt for license information.
5 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6 //
7 //===----------------------------------------------------------------------===//
8 //
9 // This file implements functionality to progressively decompose coarse-grained
10 // affine ops into finer-grained ops.
11 //
12 //===----------------------------------------------------------------------===//
13 
16 #include "mlir/IR/PatternMatch.h"
17 #include "llvm/Support/Debug.h"
18 #include "llvm/Support/DebugLog.h"
19 #include "llvm/Support/InterleavedRange.h"
20 
21 using namespace mlir;
22 using namespace mlir::affine;
23 
24 #define DEBUG_TYPE "decompose-affine-ops"
25 
26 /// Count the number of loops surrounding `operand` such that operand could be
27 /// hoisted above.
28 /// Stop counting at the first loop over which the operand cannot be hoisted.
29 static int64_t numEnclosingInvariantLoops(OpOperand &operand) {
30  int64_t count = 0;
31  Operation *currentOp = operand.getOwner();
32  while (auto loopOp = currentOp->getParentOfType<LoopLikeOpInterface>()) {
33  if (!loopOp.isDefinedOutsideOfLoop(operand.get()))
34  break;
35  currentOp = loopOp;
36  count++;
37  }
38  return count;
39 }
40 
42  AffineApplyOp op) {
43  SmallVector<int64_t> numInvariant = llvm::to_vector(
44  llvm::map_range(op->getOpOperands(), [&](OpOperand &operand) {
45  return numEnclosingInvariantLoops(operand);
46  }));
47 
48  int64_t numOperands = op.getNumOperands();
49  SmallVector<int64_t> operandPositions =
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];
53  });
54 
55  SmallVector<AffineExpr> replacements(numOperands);
56  SmallVector<Value> operands(numOperands);
57  for (int64_t i = 0; i < numOperands; ++i) {
58  operands[i] = op.getOperand(operandPositions[i]);
59  replacements[operandPositions[i]] = getAffineSymbolExpr(i, op.getContext());
60  }
61 
62  AffineMap map = op.getAffineMap();
63  ArrayRef<AffineExpr> repls{replacements};
64  map = map.replaceDimsAndSymbols(repls.take_front(map.getNumDims()),
65  repls.drop_front(map.getNumDims()),
66  /*numResultDims=*/0,
67  /*numResultSyms=*/numOperands);
68  map = AffineMap::get(0, numOperands,
69  simplifyAffineExpr(map.getResult(0), 0, numOperands),
70  op->getContext());
71  canonicalizeMapAndOperands(&map, &operands);
72 
73  rewriter.startOpModification(op);
74  op.setMap(map);
75  op->setOperands(operands);
76  rewriter.finalizeOpModification(op);
77 }
78 
79 /// Build an affine.apply that is a subexpression `expr` of `originalOp`s affine
80 /// map and with the same operands.
81 /// Canonicalize the map and operands to deduplicate and drop dead operands
82 /// before returning but do not perform maximal composition of AffineApplyOp
83 /// which would defeat the purpose.
84 static AffineApplyOp createSubApply(RewriterBase &rewriter,
85  AffineApplyOp originalOp, AffineExpr expr) {
86  MLIRContext *ctx = originalOp->getContext();
87  AffineMap m = originalOp.getAffineMap();
88  auto rhsMap = AffineMap::get(m.getNumDims(), m.getNumSymbols(), expr, ctx);
89  SmallVector<Value> rhsOperands = originalOp->getOperands();
90  canonicalizeMapAndOperands(&rhsMap, &rhsOperands);
91  return AffineApplyOp::create(rewriter, originalOp.getLoc(), rhsMap,
92  rhsOperands);
93 }
94 
95 FailureOr<AffineApplyOp> mlir::affine::decompose(RewriterBase &rewriter,
96  AffineApplyOp op) {
97  // 1. Preconditions: only handle dimensionless AffineApplyOp maps with a
98  // top-level binary expression that we can reassociate (i.e. add or mul).
99  AffineMap m = op.getAffineMap();
100  if (m.getNumDims() > 0)
101  return rewriter.notifyMatchFailure(op, "expected no dims");
102 
103  AffineExpr remainingExp = m.getResult(0);
104  auto binExpr = dyn_cast<AffineBinaryOpExpr>(remainingExp);
105  if (!binExpr)
106  return rewriter.notifyMatchFailure(op, "terminal affine.apply");
107 
108  if (!isa<AffineBinaryOpExpr>(binExpr.getLHS()) &&
109  !isa<AffineBinaryOpExpr>(binExpr.getRHS()))
110  return rewriter.notifyMatchFailure(op, "terminal affine.apply");
111 
112  bool supportedKind = ((binExpr.getKind() == AffineExprKind::Add) ||
113  (binExpr.getKind() == AffineExprKind::Mul));
114  if (!supportedKind)
115  return rewriter.notifyMatchFailure(
116  op, "only add or mul binary expr can be reassociated");
117 
118  LDBG() << "Start decomposeIntoFinerGrainedOps: " << op;
119 
120  // 2. Iteratively extract the RHS subexpressions while the top-level binary
121  // expr kind remains the same.
122  MLIRContext *ctx = op->getContext();
123  SmallVector<AffineExpr> subExpressions;
124  while (true) {
125  auto currentBinExpr = dyn_cast<AffineBinaryOpExpr>(remainingExp);
126  if (!currentBinExpr || currentBinExpr.getKind() != binExpr.getKind()) {
127  subExpressions.push_back(remainingExp);
128  LDBG() << "--terminal: " << subExpressions.back();
129  break;
130  }
131  subExpressions.push_back(currentBinExpr.getRHS());
132  LDBG() << "--subExpr: " << subExpressions.back();
133  remainingExp = currentBinExpr.getLHS();
134  }
135 
136  // 3. Reorder subExpressions by the min symbol they are a function of.
137  // This also takes care of properly reordering local variables.
138  // This however won't be able to split expression that cannot be reassociated
139  // such as ones that involve divs and multiple symbols.
140  auto getMaxSymbol = [&](AffineExpr e) -> int64_t {
141  for (int64_t i = m.getNumSymbols(); i >= 0; --i)
142  if (e.isFunctionOfSymbol(i))
143  return i;
144  return -1;
145  };
146  llvm::stable_sort(subExpressions, [&](AffineExpr e1, AffineExpr e2) {
147  return getMaxSymbol(e1) < getMaxSymbol(e2);
148  });
149  LDBG() << "--sorted subexprs: " << llvm::interleaved(subExpressions);
150 
151  // 4. Merge sorted subExpressions iteratively, thus achieving reassociation.
152  auto s0 = getAffineSymbolExpr(0, ctx);
153  auto s1 = getAffineSymbolExpr(1, ctx);
154  AffineMap binMap = AffineMap::get(
155  /*dimCount=*/0, /*symbolCount=*/2,
156  getAffineBinaryOpExpr(binExpr.getKind(), s0, s1), ctx);
157 
158  auto current = createSubApply(rewriter, op, subExpressions[0]);
159  for (int64_t i = 1, e = subExpressions.size(); i < e; ++i) {
160  Value tmp = createSubApply(rewriter, op, subExpressions[i]);
161  current = AffineApplyOp::create(rewriter, op.getLoc(), binMap,
162  ValueRange{current, tmp});
163  LDBG() << "--reassociate into: " << current;
164  }
165 
166  // 5. Replace original op.
167  rewriter.replaceOp(op, current.getResult());
168  return current;
169 }
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.
Definition: AffineExpr.h:68
MLIRContext * getContext() const
Definition: AffineExpr.cpp:31
A multi-dimensional affine map Affine map's are immutable like Type's, and they are uniqued.
Definition: AffineMap.h:46
static AffineMap get(MLIRContext *context)
Returns a zero result affine map with no dimensions or symbols: () -> ().
unsigned getNumSymbols() const
Definition: AffineMap.cpp:394
unsigned getNumDims() const
Definition: AffineMap.cpp:390
AffineMap replaceDimsAndSymbols(ArrayRef< AffineExpr > dimReplacements, ArrayRef< AffineExpr > symReplacements, unsigned numResultDims, unsigned numResultSyms) const
This method substitutes any uses of dimensions and symbols (e.g.
Definition: AffineMap.cpp:496
AffineExpr getResult(unsigned idx) const
Definition: AffineMap.cpp:407
IRValueT get() const
Return the current value being used by this operand.
Definition: UseDefLists.h:160
MLIRContext is the top-level object for a collection of MLIR operations.
Definition: MLIRContext.h:63
This class represents an operand of an operation.
Definition: Value.h:257
Operation is the basic unit of execution within MLIR.
Definition: Operation.h:88
OpTy getParentOfType()
Return the closest surrounding parent operation that is of type 'OpTy'.
Definition: Operation.h:238
This class coordinates the application of a rewrite on a set of IR, providing a way for clients to tr...
Definition: PatternMatch.h:358
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,...
Definition: PatternMatch.h:716
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...
Definition: PatternMatch.h:612
This class provides an abstraction over the different types of ranges over Values.
Definition: ValueRange.h:387
This class represents an instance of an SSA value in the MLIR system, representing a computable value...
Definition: Value.h:96
Operation * getOwner() const
Return the owner of this operand.
Definition: UseDefLists.h:38
void canonicalizeMapAndOperands(AffineMap *map, SmallVectorImpl< Value > *operands)
Modifies both map and operands in-place so as to:
Definition: AffineOps.cpp:1619
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.
Definition: LoopUtils.cpp:2817
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)
Definition: AffineExpr.cpp:68
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)
Definition: AffineExpr.cpp:629