9b5ef2bea8
This commit implements `LoopLikeOpInterface` on `scf.while`. This enables LICM (and potentially other transforms) on `scf.while`. `LoopLikeOpInterface::getLoopBody()` is renamed to `getLoopRegions` and can now return multiple regions. Also fix a bug in the default implementation of `LoopLikeOpInterface::isDefinedOutsideOfLoop()`, which returned "false" for some values that are defined outside of the loop (in a nested op, in such a way that the value does not dominate the loop). This interface is currently only used for LICM and there is no way to trigger this bug, so no test is added.
253 lines
10 KiB
C++
253 lines
10 KiB
C++
//===----------- MultiBuffering.cpp ---------------------------------------===//
|
|
//
|
|
// 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 implements multi buffering transformation.
|
|
//
|
|
//===----------------------------------------------------------------------===//
|
|
|
|
#include "mlir/Dialect/Affine/IR/AffineOps.h"
|
|
#include "mlir/Dialect/Arith/Utils/Utils.h"
|
|
#include "mlir/Dialect/MemRef/IR/MemRef.h"
|
|
#include "mlir/Dialect/MemRef/Transforms/Passes.h"
|
|
#include "mlir/Dialect/MemRef/Transforms/Transforms.h"
|
|
#include "mlir/IR/AffineExpr.h"
|
|
#include "mlir/IR/BuiltinAttributes.h"
|
|
#include "mlir/IR/Dominance.h"
|
|
#include "mlir/IR/PatternMatch.h"
|
|
#include "mlir/IR/ValueRange.h"
|
|
#include "mlir/Interfaces/LoopLikeInterface.h"
|
|
#include "llvm/ADT/STLExtras.h"
|
|
#include "llvm/Support/Debug.h"
|
|
|
|
using namespace mlir;
|
|
|
|
#define DEBUG_TYPE "memref-transforms"
|
|
#define DBGS() (llvm::dbgs() << "[" DEBUG_TYPE "]: ")
|
|
#define DBGSNL() (llvm::dbgs() << "\n")
|
|
|
|
/// Return true if the op fully overwrite the given `buffer` value.
|
|
static bool overrideBuffer(Operation *op, Value buffer) {
|
|
auto copyOp = dyn_cast<memref::CopyOp>(op);
|
|
if (!copyOp)
|
|
return false;
|
|
return copyOp.getTarget() == buffer;
|
|
}
|
|
|
|
/// Replace the uses of `oldOp` with the given `val` and for subview uses
|
|
/// propagate the type change. Changing the memref type may require propagating
|
|
/// it through subview ops so we cannot just do a replaceAllUse but need to
|
|
/// propagate the type change and erase old subview ops.
|
|
static void replaceUsesAndPropagateType(RewriterBase &rewriter,
|
|
Operation *oldOp, Value val) {
|
|
SmallVector<Operation *> opsToDelete;
|
|
SmallVector<OpOperand *> operandsToReplace;
|
|
|
|
// Save the operand to replace / delete later (avoid iterator invalidation).
|
|
// TODO: can we use an early_inc iterator?
|
|
for (OpOperand &use : oldOp->getUses()) {
|
|
// Non-subview ops will be replaced by `val`.
|
|
auto subviewUse = dyn_cast<memref::SubViewOp>(use.getOwner());
|
|
if (!subviewUse) {
|
|
operandsToReplace.push_back(&use);
|
|
continue;
|
|
}
|
|
|
|
// `subview(old_op)` is replaced by a new `subview(val)`.
|
|
OpBuilder::InsertionGuard g(rewriter);
|
|
rewriter.setInsertionPoint(subviewUse);
|
|
Type newType = memref::SubViewOp::inferRankReducedResultType(
|
|
subviewUse.getType().getShape(), cast<MemRefType>(val.getType()),
|
|
subviewUse.getStaticOffsets(), subviewUse.getStaticSizes(),
|
|
subviewUse.getStaticStrides());
|
|
Value newSubview = rewriter.create<memref::SubViewOp>(
|
|
subviewUse->getLoc(), cast<MemRefType>(newType), val,
|
|
subviewUse.getMixedOffsets(), subviewUse.getMixedSizes(),
|
|
subviewUse.getMixedStrides());
|
|
|
|
// Ouch recursion ... is this really necessary?
|
|
replaceUsesAndPropagateType(rewriter, subviewUse, newSubview);
|
|
|
|
opsToDelete.push_back(use.getOwner());
|
|
}
|
|
|
|
// Perform late replacement.
|
|
// TODO: can we use an early_inc iterator?
|
|
for (OpOperand *operand : operandsToReplace) {
|
|
Operation *op = operand->getOwner();
|
|
rewriter.startRootUpdate(op);
|
|
operand->set(val);
|
|
rewriter.finalizeRootUpdate(op);
|
|
}
|
|
|
|
// Perform late op erasure.
|
|
// TODO: can we use an early_inc iterator?
|
|
for (Operation *op : opsToDelete)
|
|
rewriter.eraseOp(op);
|
|
}
|
|
|
|
// Transformation to do multi-buffering/array expansion to remove dependencies
|
|
// on the temporary allocation between consecutive loop iterations.
|
|
// Returns success if the transformation happened and failure otherwise.
|
|
// This is not a pattern as it requires propagating the new memref type to its
|
|
// uses and requires updating subview ops.
|
|
FailureOr<memref::AllocOp>
|
|
mlir::memref::multiBuffer(RewriterBase &rewriter, memref::AllocOp allocOp,
|
|
unsigned multiBufferingFactor,
|
|
bool skipOverrideAnalysis) {
|
|
LLVM_DEBUG(DBGS() << "Start multibuffering: " << allocOp << "\n");
|
|
DominanceInfo dom(allocOp->getParentOp());
|
|
LoopLikeOpInterface candidateLoop;
|
|
for (Operation *user : allocOp->getUsers()) {
|
|
auto parentLoop = user->getParentOfType<LoopLikeOpInterface>();
|
|
if (!parentLoop) {
|
|
if (isa<memref::DeallocOp>(user)) {
|
|
// Allow dealloc outside of any loop.
|
|
// TODO: The whole precondition function here is very brittle and will
|
|
// need to rethought an isolated into a cleaner analysis.
|
|
continue;
|
|
}
|
|
LLVM_DEBUG(DBGS() << "--no parent loop -> fail\n");
|
|
LLVM_DEBUG(DBGS() << "----due to user: " << *user << "\n");
|
|
return failure();
|
|
}
|
|
if (!skipOverrideAnalysis) {
|
|
/// Make sure there is no loop-carried dependency on the allocation.
|
|
if (!overrideBuffer(user, allocOp.getResult())) {
|
|
LLVM_DEBUG(DBGS() << "--Skip user: found loop-carried dependence\n");
|
|
continue;
|
|
}
|
|
// If this user doesn't dominate all the other users keep looking.
|
|
if (llvm::any_of(allocOp->getUsers(), [&](Operation *otherUser) {
|
|
return !dom.dominates(user, otherUser);
|
|
})) {
|
|
LLVM_DEBUG(
|
|
DBGS() << "--Skip user: does not dominate all other users\n");
|
|
continue;
|
|
}
|
|
} else {
|
|
if (llvm::any_of(allocOp->getUsers(), [&](Operation *otherUser) {
|
|
return !isa<memref::DeallocOp>(otherUser) &&
|
|
!parentLoop->isProperAncestor(otherUser);
|
|
})) {
|
|
LLVM_DEBUG(
|
|
DBGS()
|
|
<< "--Skip user: not all other users are in the parent loop\n");
|
|
continue;
|
|
}
|
|
}
|
|
candidateLoop = parentLoop;
|
|
break;
|
|
}
|
|
|
|
if (!candidateLoop) {
|
|
LLVM_DEBUG(DBGS() << "Skip alloc: no candidate loop\n");
|
|
return failure();
|
|
}
|
|
|
|
std::optional<Value> inductionVar = candidateLoop.getSingleInductionVar();
|
|
std::optional<OpFoldResult> lowerBound = candidateLoop.getSingleLowerBound();
|
|
std::optional<OpFoldResult> singleStep = candidateLoop.getSingleStep();
|
|
if (!inductionVar || !lowerBound || !singleStep ||
|
|
!llvm::hasSingleElement(candidateLoop.getLoopRegions())) {
|
|
LLVM_DEBUG(DBGS() << "Skip alloc: no single iv, lb, step or region\n");
|
|
return failure();
|
|
}
|
|
|
|
if (!dom.dominates(allocOp.getOperation(), candidateLoop)) {
|
|
LLVM_DEBUG(DBGS() << "Skip alloc: does not dominate candidate loop\n");
|
|
return failure();
|
|
}
|
|
|
|
LLVM_DEBUG(DBGS() << "Start multibuffering loop: " << candidateLoop << "\n");
|
|
|
|
// 1. Construct the multi-buffered memref type.
|
|
ArrayRef<int64_t> originalShape = allocOp.getType().getShape();
|
|
SmallVector<int64_t, 4> multiBufferedShape{multiBufferingFactor};
|
|
llvm::append_range(multiBufferedShape, originalShape);
|
|
LLVM_DEBUG(DBGS() << "--original type: " << allocOp.getType() << "\n");
|
|
MemRefType mbMemRefType = MemRefType::Builder(allocOp.getType())
|
|
.setShape(multiBufferedShape)
|
|
.setLayout(MemRefLayoutAttrInterface());
|
|
LLVM_DEBUG(DBGS() << "--multi-buffered type: " << mbMemRefType << "\n");
|
|
|
|
// 2. Create the multi-buffered alloc.
|
|
Location loc = allocOp->getLoc();
|
|
OpBuilder::InsertionGuard g(rewriter);
|
|
rewriter.setInsertionPoint(allocOp);
|
|
auto mbAlloc = rewriter.create<memref::AllocOp>(
|
|
loc, mbMemRefType, ValueRange{}, allocOp->getAttrs());
|
|
LLVM_DEBUG(DBGS() << "--multi-buffered alloc: " << mbAlloc << "\n");
|
|
|
|
// 3. Within the loop, build the modular leading index (i.e. each loop
|
|
// iteration %iv accesses slice ((%iv - %lb) / %step) % %mb_factor).
|
|
rewriter.setInsertionPointToStart(
|
|
&candidateLoop.getLoopRegions().front()->front());
|
|
Value ivVal = *inductionVar;
|
|
Value lbVal = getValueOrCreateConstantIndexOp(rewriter, loc, *lowerBound);
|
|
Value stepVal = getValueOrCreateConstantIndexOp(rewriter, loc, *singleStep);
|
|
AffineExpr iv, lb, step;
|
|
bindDims(rewriter.getContext(), iv, lb, step);
|
|
Value bufferIndex = affine::makeComposedAffineApply(
|
|
rewriter, loc, ((iv - lb).floorDiv(step)) % multiBufferingFactor,
|
|
{ivVal, lbVal, stepVal});
|
|
LLVM_DEBUG(DBGS() << "--multi-buffered indexing: " << bufferIndex << "\n");
|
|
|
|
// 4. Build the subview accessing the particular slice, taking modular
|
|
// rotation into account.
|
|
int64_t mbMemRefTypeRank = mbMemRefType.getRank();
|
|
IntegerAttr zero = rewriter.getIndexAttr(0);
|
|
IntegerAttr one = rewriter.getIndexAttr(1);
|
|
SmallVector<OpFoldResult> offsets(mbMemRefTypeRank, zero);
|
|
SmallVector<OpFoldResult> sizes(mbMemRefTypeRank, one);
|
|
SmallVector<OpFoldResult> strides(mbMemRefTypeRank, one);
|
|
// Offset is [bufferIndex, 0 ... 0 ].
|
|
offsets.front() = bufferIndex;
|
|
// Sizes is [1, original_size_0 ... original_size_n ].
|
|
for (int64_t i = 0, e = originalShape.size(); i != e; ++i)
|
|
sizes[1 + i] = rewriter.getIndexAttr(originalShape[i]);
|
|
// Strides is [1, 1 ... 1 ].
|
|
auto dstMemref =
|
|
cast<MemRefType>(memref::SubViewOp::inferRankReducedResultType(
|
|
originalShape, mbMemRefType, offsets, sizes, strides));
|
|
Value subview = rewriter.create<memref::SubViewOp>(loc, dstMemref, mbAlloc,
|
|
offsets, sizes, strides);
|
|
LLVM_DEBUG(DBGS() << "--multi-buffered slice: " << subview << "\n");
|
|
|
|
// 5. Due to the recursive nature of replaceUsesAndPropagateType , we need to
|
|
// handle dealloc uses separately..
|
|
for (OpOperand &use : llvm::make_early_inc_range(allocOp->getUses())) {
|
|
auto deallocOp = dyn_cast<memref::DeallocOp>(use.getOwner());
|
|
if (!deallocOp)
|
|
continue;
|
|
OpBuilder::InsertionGuard g(rewriter);
|
|
rewriter.setInsertionPoint(deallocOp);
|
|
auto newDeallocOp =
|
|
rewriter.create<memref::DeallocOp>(deallocOp->getLoc(), mbAlloc);
|
|
(void)newDeallocOp;
|
|
LLVM_DEBUG(DBGS() << "----Created dealloc: " << newDeallocOp << "\n");
|
|
rewriter.eraseOp(deallocOp);
|
|
}
|
|
|
|
// 6. RAUW with the particular slice, taking modular rotation into account.
|
|
replaceUsesAndPropagateType(rewriter, allocOp, subview);
|
|
|
|
// 7. Finally, erase the old allocOp.
|
|
rewriter.eraseOp(allocOp);
|
|
|
|
return mbAlloc;
|
|
}
|
|
|
|
FailureOr<memref::AllocOp>
|
|
mlir::memref::multiBuffer(memref::AllocOp allocOp,
|
|
unsigned multiBufferingFactor,
|
|
bool skipOverrideAnalysis) {
|
|
IRRewriter rewriter(allocOp->getContext());
|
|
return multiBuffer(rewriter, allocOp, multiBufferingFactor,
|
|
skipOverrideAnalysis);
|
|
}
|