f0bf97281f
The memref::multiBuffer transformation replaces an allocation with a multi-buffered allocation and creates a strided memref.subview at each loop iteration. When the original allocation is used through view-like ops, the existing code only handles SubViewOp, leaving other view-like ops with incorrect types. This patch extends replaceUsesAndPropagateType to handle ExpandShapeOp, CollapseShapeOp, and CastOp using TypeSwitch. For each view-like op, we compute the correct result type (or assert on failure) and create a new operation, then recursively propagate the updated type through chains. New FileCheck tests cover expand_shape, collapse_shape, cast, and a chained expand_shape->cast case. A single ViewLikeOpInterface hook is not practical here: view-like ops have distinct type inference and validity rules (e.g., subview uses offset/size/stride inference, expand/collapse use reassociation, cast requires compatibility checks). Ops like memref.view or memref.reinterpret_cast need additional layout/size validation beyond what multi-buffering currently tracks, so this patch handles the common safe cases directly.
298 lines
12 KiB
C++
298 lines
12 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/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/ADT/TypeSwitch.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 view-like uses
|
|
/// propagate the type change. Changing the memref type may require propagating
|
|
/// it through view-like ops (subview, expand_shape, collapse_shape, cast) so
|
|
/// we need to propagate the type change and erase old view ops.
|
|
///
|
|
/// Only view-like ops whose result type can be recomputed from the new source
|
|
/// type and existing op attributes are handled here. Other ops fall back to
|
|
/// operand replacement without type propagation.
|
|
static LogicalResult replaceUsesAndPropagateType(RewriterBase &rewriter,
|
|
Operation *oldOp, Value val) {
|
|
SmallVector<Operation *> opsToErase;
|
|
// Iterate with early_inc to erase current user inside the loop.
|
|
for (OpOperand &use : llvm::make_early_inc_range(oldOp->getUses())) {
|
|
Operation *user = use.getOwner();
|
|
OpBuilder::InsertionGuard g(rewriter);
|
|
rewriter.setInsertionPoint(user);
|
|
MemRefType srcType = cast<MemRefType>(val.getType());
|
|
|
|
// Try to create a new view-like op with updated result type.
|
|
// Each view-like op has its own method to compute the result type.
|
|
bool typeInferenceFailed = false;
|
|
Value replacement =
|
|
llvm::TypeSwitch<Operation *, Value>(user)
|
|
.Case([&](memref::SubViewOp subview) -> Value {
|
|
MemRefType newType =
|
|
memref::SubViewOp::inferRankReducedResultType(
|
|
subview.getType().getShape(), srcType,
|
|
subview.getStaticOffsets(), subview.getStaticSizes(),
|
|
subview.getStaticStrides());
|
|
return memref::SubViewOp::create(
|
|
rewriter, subview->getLoc(), newType, val,
|
|
subview.getMixedOffsets(), subview.getMixedSizes(),
|
|
subview.getMixedStrides());
|
|
})
|
|
.Case([&](memref::ExpandShapeOp expand) -> Value {
|
|
FailureOr<MemRefType> newType =
|
|
memref::ExpandShapeOp::computeExpandedType(
|
|
srcType, expand.getResultType().getShape(),
|
|
expand.getReassociationIndices());
|
|
if (failed(newType)) {
|
|
typeInferenceFailed = true;
|
|
return Value();
|
|
}
|
|
return memref::ExpandShapeOp::create(
|
|
rewriter, expand->getLoc(), *newType, val,
|
|
expand.getReassociationIndices(),
|
|
expand.getMixedOutputShape());
|
|
})
|
|
.Case([&](memref::CollapseShapeOp collapse) -> Value {
|
|
FailureOr<MemRefType> newType =
|
|
memref::CollapseShapeOp::computeCollapsedType(
|
|
srcType, collapse.getReassociationIndices());
|
|
if (failed(newType)) {
|
|
typeInferenceFailed = true;
|
|
return Value();
|
|
}
|
|
return memref::CollapseShapeOp::create(
|
|
rewriter, collapse->getLoc(), *newType, val,
|
|
collapse.getReassociationIndices());
|
|
})
|
|
.Case([&](memref::CastOp cast) -> Value {
|
|
if (!memref::CastOp::areCastCompatible(srcType, cast.getType())) {
|
|
typeInferenceFailed = true;
|
|
return Value();
|
|
}
|
|
return memref::CastOp::create(rewriter, cast->getLoc(),
|
|
cast.getType(), val);
|
|
})
|
|
.Default([&](Operation *) -> Value { return Value(); });
|
|
|
|
if (typeInferenceFailed) {
|
|
user->emitOpError(
|
|
"failed to compute view-like result type after multi-buffering");
|
|
return failure();
|
|
}
|
|
|
|
if (replacement) {
|
|
// Recursively propagate through view-like ops and mark old op for
|
|
// erasure.
|
|
if (failed(replaceUsesAndPropagateType(rewriter, user, replacement)))
|
|
return failure();
|
|
opsToErase.push_back(user);
|
|
} else {
|
|
// Not a view-like op: just replace operand.
|
|
rewriter.startOpModification(user);
|
|
use.set(val);
|
|
rewriter.finalizeOpModification(user);
|
|
}
|
|
}
|
|
|
|
for (Operation *op : opsToErase) {
|
|
rewriter.eraseOp(op);
|
|
}
|
|
|
|
return success();
|
|
}
|
|
|
|
// 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 = memref::AllocOp::create(rewriter, 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 ].
|
|
MemRefType dstMemref = memref::SubViewOp::inferRankReducedResultType(
|
|
originalShape, mbMemRefType, offsets, sizes, strides);
|
|
Value subview = memref::SubViewOp::create(rewriter, 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 =
|
|
memref::DeallocOp::create(rewriter, 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.
|
|
if (failed(replaceUsesAndPropagateType(rewriter, allocOp, subview)))
|
|
return failure();
|
|
|
|
// 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);
|
|
}
|