MLIR  22.0.0git
MultiBuffer.cpp
Go to the documentation of this file.
1 //===----------- MultiBuffering.cpp ---------------------------------------===//
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 multi buffering transformation.
10 //
11 //===----------------------------------------------------------------------===//
12 
17 #include "mlir/IR/AffineExpr.h"
19 #include "mlir/IR/Dominance.h"
20 #include "mlir/IR/PatternMatch.h"
21 #include "mlir/IR/ValueRange.h"
23 #include "llvm/ADT/STLExtras.h"
24 #include "llvm/Support/Debug.h"
25 
26 using namespace mlir;
27 
28 #define DEBUG_TYPE "memref-transforms"
29 #define DBGS() (llvm::dbgs() << "[" DEBUG_TYPE "]: ")
30 #define DBGSNL() (llvm::dbgs() << "\n")
31 
32 /// Return true if the op fully overwrite the given `buffer` value.
33 static bool overrideBuffer(Operation *op, Value buffer) {
34  auto copyOp = dyn_cast<memref::CopyOp>(op);
35  if (!copyOp)
36  return false;
37  return copyOp.getTarget() == buffer;
38 }
39 
40 /// Replace the uses of `oldOp` with the given `val` and for subview uses
41 /// propagate the type change. Changing the memref type may require propagating
42 /// it through subview ops so we cannot just do a replaceAllUse but need to
43 /// propagate the type change and erase old subview ops.
45  Operation *oldOp, Value val) {
46  // Iterate with early_inc to erase current user inside the loop.
47  for (OpOperand &use : llvm::make_early_inc_range(oldOp->getUses())) {
48  Operation *user = use.getOwner();
49  if (auto subviewUse = dyn_cast<memref::SubViewOp>(user)) {
50  // `subview(old_op)` is replaced by a new `subview(val)`.
51  OpBuilder::InsertionGuard g(rewriter);
52  rewriter.setInsertionPoint(subviewUse);
53  MemRefType newType = memref::SubViewOp::inferRankReducedResultType(
54  subviewUse.getType().getShape(), cast<MemRefType>(val.getType()),
55  subviewUse.getStaticOffsets(), subviewUse.getStaticSizes(),
56  subviewUse.getStaticStrides());
57  Value newSubview = memref::SubViewOp::create(
58  rewriter, subviewUse->getLoc(), newType, val,
59  subviewUse.getMixedOffsets(), subviewUse.getMixedSizes(),
60  subviewUse.getMixedStrides());
61 
62  // Ouch recursion ... is this really necessary?
63  replaceUsesAndPropagateType(rewriter, subviewUse, newSubview);
64 
65  // Safe to erase.
66  rewriter.eraseOp(subviewUse);
67  continue;
68  }
69  // Non-subview: replace with new value.
70  rewriter.startOpModification(user);
71  use.set(val);
72  rewriter.finalizeOpModification(user);
73  }
74 }
75 
76 // Transformation to do multi-buffering/array expansion to remove dependencies
77 // on the temporary allocation between consecutive loop iterations.
78 // Returns success if the transformation happened and failure otherwise.
79 // This is not a pattern as it requires propagating the new memref type to its
80 // uses and requires updating subview ops.
81 FailureOr<memref::AllocOp>
82 mlir::memref::multiBuffer(RewriterBase &rewriter, memref::AllocOp allocOp,
83  unsigned multiBufferingFactor,
84  bool skipOverrideAnalysis) {
85  LLVM_DEBUG(DBGS() << "Start multibuffering: " << allocOp << "\n");
86  DominanceInfo dom(allocOp->getParentOp());
87  LoopLikeOpInterface candidateLoop;
88  for (Operation *user : allocOp->getUsers()) {
89  auto parentLoop = user->getParentOfType<LoopLikeOpInterface>();
90  if (!parentLoop) {
91  if (isa<memref::DeallocOp>(user)) {
92  // Allow dealloc outside of any loop.
93  // TODO: The whole precondition function here is very brittle and will
94  // need to rethought an isolated into a cleaner analysis.
95  continue;
96  }
97  LLVM_DEBUG(DBGS() << "--no parent loop -> fail\n");
98  LLVM_DEBUG(DBGS() << "----due to user: " << *user << "\n");
99  return failure();
100  }
101  if (!skipOverrideAnalysis) {
102  /// Make sure there is no loop-carried dependency on the allocation.
103  if (!overrideBuffer(user, allocOp.getResult())) {
104  LLVM_DEBUG(DBGS() << "--Skip user: found loop-carried dependence\n");
105  continue;
106  }
107  // If this user doesn't dominate all the other users keep looking.
108  if (llvm::any_of(allocOp->getUsers(), [&](Operation *otherUser) {
109  return !dom.dominates(user, otherUser);
110  })) {
111  LLVM_DEBUG(
112  DBGS() << "--Skip user: does not dominate all other users\n");
113  continue;
114  }
115  } else {
116  if (llvm::any_of(allocOp->getUsers(), [&](Operation *otherUser) {
117  return !isa<memref::DeallocOp>(otherUser) &&
118  !parentLoop->isProperAncestor(otherUser);
119  })) {
120  LLVM_DEBUG(
121  DBGS()
122  << "--Skip user: not all other users are in the parent loop\n");
123  continue;
124  }
125  }
126  candidateLoop = parentLoop;
127  break;
128  }
129 
130  if (!candidateLoop) {
131  LLVM_DEBUG(DBGS() << "Skip alloc: no candidate loop\n");
132  return failure();
133  }
134 
135  std::optional<Value> inductionVar = candidateLoop.getSingleInductionVar();
136  std::optional<OpFoldResult> lowerBound = candidateLoop.getSingleLowerBound();
137  std::optional<OpFoldResult> singleStep = candidateLoop.getSingleStep();
138  if (!inductionVar || !lowerBound || !singleStep ||
139  !llvm::hasSingleElement(candidateLoop.getLoopRegions())) {
140  LLVM_DEBUG(DBGS() << "Skip alloc: no single iv, lb, step or region\n");
141  return failure();
142  }
143 
144  if (!dom.dominates(allocOp.getOperation(), candidateLoop)) {
145  LLVM_DEBUG(DBGS() << "Skip alloc: does not dominate candidate loop\n");
146  return failure();
147  }
148 
149  LLVM_DEBUG(DBGS() << "Start multibuffering loop: " << candidateLoop << "\n");
150 
151  // 1. Construct the multi-buffered memref type.
152  ArrayRef<int64_t> originalShape = allocOp.getType().getShape();
153  SmallVector<int64_t, 4> multiBufferedShape{multiBufferingFactor};
154  llvm::append_range(multiBufferedShape, originalShape);
155  LLVM_DEBUG(DBGS() << "--original type: " << allocOp.getType() << "\n");
156  MemRefType mbMemRefType = MemRefType::Builder(allocOp.getType())
157  .setShape(multiBufferedShape)
158  .setLayout(MemRefLayoutAttrInterface());
159  LLVM_DEBUG(DBGS() << "--multi-buffered type: " << mbMemRefType << "\n");
160 
161  // 2. Create the multi-buffered alloc.
162  Location loc = allocOp->getLoc();
163  OpBuilder::InsertionGuard g(rewriter);
164  rewriter.setInsertionPoint(allocOp);
165  auto mbAlloc = memref::AllocOp::create(rewriter, loc, mbMemRefType,
166  ValueRange{}, allocOp->getAttrs());
167  LLVM_DEBUG(DBGS() << "--multi-buffered alloc: " << mbAlloc << "\n");
168 
169  // 3. Within the loop, build the modular leading index (i.e. each loop
170  // iteration %iv accesses slice ((%iv - %lb) / %step) % %mb_factor).
171  rewriter.setInsertionPointToStart(
172  &candidateLoop.getLoopRegions().front()->front());
173  Value ivVal = *inductionVar;
174  Value lbVal = getValueOrCreateConstantIndexOp(rewriter, loc, *lowerBound);
175  Value stepVal = getValueOrCreateConstantIndexOp(rewriter, loc, *singleStep);
176  AffineExpr iv, lb, step;
177  bindDims(rewriter.getContext(), iv, lb, step);
179  rewriter, loc, ((iv - lb).floorDiv(step)) % multiBufferingFactor,
180  {ivVal, lbVal, stepVal});
181  LLVM_DEBUG(DBGS() << "--multi-buffered indexing: " << bufferIndex << "\n");
182 
183  // 4. Build the subview accessing the particular slice, taking modular
184  // rotation into account.
185  int64_t mbMemRefTypeRank = mbMemRefType.getRank();
186  IntegerAttr zero = rewriter.getIndexAttr(0);
187  IntegerAttr one = rewriter.getIndexAttr(1);
188  SmallVector<OpFoldResult> offsets(mbMemRefTypeRank, zero);
189  SmallVector<OpFoldResult> sizes(mbMemRefTypeRank, one);
190  SmallVector<OpFoldResult> strides(mbMemRefTypeRank, one);
191  // Offset is [bufferIndex, 0 ... 0 ].
192  offsets.front() = bufferIndex;
193  // Sizes is [1, original_size_0 ... original_size_n ].
194  for (int64_t i = 0, e = originalShape.size(); i != e; ++i)
195  sizes[1 + i] = rewriter.getIndexAttr(originalShape[i]);
196  // Strides is [1, 1 ... 1 ].
197  MemRefType dstMemref = memref::SubViewOp::inferRankReducedResultType(
198  originalShape, mbMemRefType, offsets, sizes, strides);
199  Value subview = memref::SubViewOp::create(rewriter, loc, dstMemref, mbAlloc,
200  offsets, sizes, strides);
201  LLVM_DEBUG(DBGS() << "--multi-buffered slice: " << subview << "\n");
202 
203  // 5. Due to the recursive nature of replaceUsesAndPropagateType , we need
204  // to handle dealloc uses separately..
205  for (OpOperand &use : llvm::make_early_inc_range(allocOp->getUses())) {
206  auto deallocOp = dyn_cast<memref::DeallocOp>(use.getOwner());
207  if (!deallocOp)
208  continue;
209  OpBuilder::InsertionGuard g(rewriter);
210  rewriter.setInsertionPoint(deallocOp);
211  auto newDeallocOp =
212  memref::DeallocOp::create(rewriter, deallocOp->getLoc(), mbAlloc);
213  (void)newDeallocOp;
214  LLVM_DEBUG(DBGS() << "----Created dealloc: " << newDeallocOp << "\n");
215  rewriter.eraseOp(deallocOp);
216  }
217 
218  // 6. RAUW with the particular slice, taking modular rotation into account.
219  replaceUsesAndPropagateType(rewriter, allocOp, subview);
220 
221  // 7. Finally, erase the old allocOp.
222  rewriter.eraseOp(allocOp);
223 
224  return mbAlloc;
225 }
226 
227 FailureOr<memref::AllocOp>
228 mlir::memref::multiBuffer(memref::AllocOp allocOp,
229  unsigned multiBufferingFactor,
230  bool skipOverrideAnalysis) {
231  IRRewriter rewriter(allocOp->getContext());
232  return multiBuffer(rewriter, allocOp, multiBufferingFactor,
233  skipOverrideAnalysis);
234 }
static bool overrideBuffer(Operation *op, Value buffer)
Return true if the op fully overwrite the given buffer value.
Definition: MultiBuffer.cpp:33
static void replaceUsesAndPropagateType(RewriterBase &rewriter, Operation *oldOp, Value val)
Replace the uses of oldOp with the given val and for subview uses propagate the type change.
Definition: MultiBuffer.cpp:44
#define DBGS()
Definition: MultiBuffer.cpp:29
Base type for affine expression.
Definition: AffineExpr.h:68
IntegerAttr getIndexAttr(int64_t value)
Definition: Builders.cpp:103
MLIRContext * getContext() const
Definition: Builders.h:55
A class for computing basic dominance information.
Definition: Dominance.h:140
bool dominates(Operation *a, Operation *b) const
Return true if operation A dominates operation B, i.e.
Definition: Dominance.h:158
This class coordinates rewriting a piece of IR outside of a pattern rewrite, providing a way to keep ...
Definition: PatternMatch.h:764
This class defines the main interface for locations in MLIR and acts as a non-nullable wrapper around...
Definition: Location.h:76
This is a builder type that keeps local references to arguments.
Definition: BuiltinTypes.h:182
Builder & setShape(ArrayRef< int64_t > newShape)
Definition: BuiltinTypes.h:193
RAII guard to reset the insertion point of the builder when destroyed.
Definition: Builders.h:346
void setInsertionPointToStart(Block *block)
Sets the insertion point to the start of the specified block.
Definition: Builders.h:429
void setInsertionPoint(Block *block, Block::iterator insertPoint)
Set the insertion point to the specified location.
Definition: Builders.h:396
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
user_range getUsers()
Returns a range of all users.
Definition: Operation.h:873
use_range getUses()
Returns a range of all uses, which is useful for iterating over all uses.
Definition: Operation.h:846
This class coordinates the application of a rewrite on a set of IR, providing a way for clients to tr...
Definition: PatternMatch.h:358
virtual void finalizeOpModification(Operation *op)
This method is used to signal the end of an in-place modification of the given operation.
virtual void eraseOp(Operation *op)
This method erases an operation that is known to have no uses.
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
Type getType() const
Return the type of this value.
Definition: Value.h:105
AffineApplyOp makeComposedAffineApply(OpBuilder &b, Location loc, AffineMap map, ArrayRef< OpFoldResult > operands, bool composeAffineMin=false)
Returns a composed AffineApplyOp by composing map and operands with other AffineApplyOps supplying th...
Definition: AffineOps.cpp:1276
FailureOr< memref::AllocOp > multiBuffer(RewriterBase &rewriter, memref::AllocOp allocOp, unsigned multiplier, bool skipOverrideAnalysis=false)
Transformation to do multi-buffering/array expansion to remove dependencies on the temporary allocati...
Definition: MultiBuffer.cpp:82
Include the generated interface declarations.
void bindDims(MLIRContext *ctx, AffineExprTy &...exprs)
Bind a list of AffineExpr references to DimExpr at positions: [0 .
Definition: AffineExpr.h:311
Value getValueOrCreateConstantIndexOp(OpBuilder &b, Location loc, OpFoldResult ofr)
Converts an OpFoldResult to a Value.
Definition: Utils.cpp:111