nnicolosi/Raptor
1
0
Fork 0
Code Issues Pull Requests Actions 5 Packages Projects Releases Wiki Activity

Raptor sync wait

Browse Source
This commit is contained in:
ilgeco
2026-08-06 14:32:46 +02:00
parent a963009855
commit a39fdba366
48 changed files with 3357 additions and 96 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/PIM/Conversion/SpatialToPim/Common.cpp
+46
View File
@@ -1,7 +1,10 @@
#include "mlir/IR/ValueRange.h"
#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/Dialect/Func/IR/FuncOps.h"
#include "mlir/Dialect/MemRef/IR/MemRef.h"
#include "mlir/Dialect/SCF/IR/SCF.h"
#include "mlir/IR/BuiltinOps.h"
#include "llvm/ADT/STLExtras.h"
@@ -28,6 +31,49 @@ FailureOr<IntegerAttr> getTensorSizeInBytesAttr(Builder& builder, Operation* anc
return pim::getCheckedI32Attr(builder, anchor, *byteSize, "tensor byte size");
}
LogicalResult materializePipelineHostBuffer(
func::FuncOp funcOp, RewriterBase &rewriter) {
auto bytes = funcOp->getAttrOfType<IntegerAttr>(
kPipelineHostBufferBytesAttrName);
if (!bytes)
return success();
if (bytes.getInt() <= 0)
return funcOp.emitOpError(
"pipeline host transfer buffer must be positive");
ModuleOp moduleOp = funcOp->getParentOfType<ModuleOp>();
if (moduleOp.lookupSymbol<memref::GlobalOp>(kPipelineHostBufferName))
return funcOp.emitOpError(
"pipeline host transfer buffer symbol already exists");
auto type = MemRefType::get(
{bytes.getInt()}, rewriter.getI8Type());
OpBuilder::InsertionGuard guard(rewriter);
rewriter.setInsertionPointToStart(moduleOp.getBody());
memref::GlobalOp::create(
rewriter, funcOp.getLoc(),
rewriter.getStringAttr(kPipelineHostBufferName),
rewriter.getStringAttr("private"), TypeAttr::get(type), Attribute(),
UnitAttr(), IntegerAttr());
return success();
}
FailureOr<mlir::Value> getPipelineHostBuffer(
OpBuilder &builder, Operation *anchor) {
auto funcOp = anchor->getParentOfType<func::FuncOp>();
auto moduleOp = anchor->getParentOfType<ModuleOp>();
auto bytes = funcOp
? funcOp->getAttrOfType<IntegerAttr>(kPipelineHostBufferBytesAttrName)
: IntegerAttr();
auto global = moduleOp
? moduleOp.lookupSymbol<memref::GlobalOp>(kPipelineHostBufferName)
: memref::GlobalOp();
if (!bytes || !global)
return anchor->emitOpError(
"requires the pipeline host transfer buffer"), failure();
auto type = MemRefType::get({bytes.getInt()}, builder.getI8Type());
return memref::GetGlobalOp::create(
builder, anchor->getLoc(), type, kPipelineHostBufferName).getResult();
}
Operation* getEarliestUserWithinBlock(mlir::Value value) {
auto users = value.getUsers();
src/PIM/Conversion/SpatialToPim/Common.hpp
+7
View File
@@ -10,6 +10,7 @@
#include "mlir/IR/Builders.h"
#include "mlir/IR/Value.h"
#include "mlir/Dialect/Tensor/IR/Tensor.h"
#include "mlir/Dialect/Func/IR/FuncOps.h"
#include "mlir/Support/LogicalResult.h"
#include "src/Accelerators/PIM/Common/PimCommon.hpp"
@@ -23,6 +24,12 @@ namespace onnx_mlir {
mlir::FailureOr<mlir::IntegerAttr>
getTensorSizeInBytesAttr(mlir::Builder& builder, mlir::Operation* anchor, mlir::Value value);
mlir::LogicalResult materializePipelineHostBuffer(
mlir::func::FuncOp funcOp, mlir::RewriterBase &rewriter);
mlir::FailureOr<mlir::Value> getPipelineHostBuffer(
mlir::OpBuilder &builder, mlir::Operation *anchor);
template <class T>
size_t rangeLength(const mlir::iterator_range<T> range) {
return std::distance(range.begin(), range.end());
src/PIM/Conversion/SpatialToPim/CoreLoweringPatterns.cpp
+28 -8
View File
@@ -345,20 +345,39 @@ LogicalResult raptor::SpatialToPimPass::lowerComputeOp(spatial::SpatScheduledCom
auto blockArg = computeOp.getInputArgument(inputIndex);
if (!blockArg)
return computeOp.emitOpError("expected compute input block arguments during lowering");
auto receiveOp = dyn_cast_or_null<spatial::SpatChannelReceiveOp>(input.getDefiningOp());
auto channelReceive = dyn_cast_or_null<spatial::SpatChannelReceiveOp>(
input.getDefiningOp());
auto hostWaitLoad = dyn_cast_or_null<spatial::SpatHostWaitLoadOp>(
input.getDefiningOp());
Operation *receiveOp = channelReceive
? channelReceive.getOperation() : hostWaitLoad.getOperation();
if (receiveOp && !blockArg->use_empty()) {
rewriter.setInsertionPoint(getEarliestUserWithinBlock(*blockArg));
auto outputType = cast<ShapedType>(blockArg->getType());
auto outputBuffer = createEmptyTensorFromShaped(rewriter, receiveOp.getLoc(), outputType);
auto outputBuffer = createEmptyTensorFromShaped(
rewriter, receiveOp->getLoc(), outputType);
auto sizeAttr = getTensorSizeInBytesAttr(rewriter, computeOp.getOperation(), *blockArg);
if (failed(sizeAttr))
return failure();
Value received =
PimReceiveOp::create(
rewriter, receiveOp.getLoc(), outputBuffer.getType(), outputBuffer,
arith::ConstantIndexOp::create(rewriter, receiveOp.getLoc(), 0),
*sizeAttr, receiveOp.getSourceCoreId())
Value zero = arith::ConstantIndexOp::create(
rewriter, receiveOp->getLoc(), 0);
Value received;
if (hostWaitLoad) {
auto hostBuffer = getPipelineHostBuffer(rewriter, hostWaitLoad);
if (failed(hostBuffer))
return failure();
PimWaitOp::create(
rewriter, receiveOp->getLoc(), hostWaitLoad.getEventRegister(),
rewriter.getI32IntegerAttr(1));
received = PimMemCopyHostToDevOp::create(
rewriter, receiveOp->getLoc(), outputBuffer.getType(), zero,
hostWaitLoad.getHostOffset(), outputBuffer, *hostBuffer, *sizeAttr)
.getOutput();
} else {
received = PimReceiveOp::create(
rewriter, receiveOp->getLoc(), outputBuffer.getType(), outputBuffer,
zero, *sizeAttr, channelReceive.getSourceCoreId()).getOutput();
}
blockArg->replaceAllUsesWith(received);
markOpToRemove(receiveOp);
continue;
@@ -383,7 +402,8 @@ LogicalResult raptor::SpatialToPimPass::lowerComputeOp(spatial::SpatScheduledCom
if (rangeLength(resultUses) == 1) {
OpOperand& resultUse = *resultUses.begin();
Operation* resultUser = resultUse.getOwner();
if (isa<spatial::SpatChannelSendOp>(resultUser))
if (isa<spatial::SpatChannelSendOp,
spatial::SpatHostStoreSyncOp>(resultUser))
continue;
}
src/PIM/Conversion/SpatialToPim/Patterns/ChannelLowering.cpp
+61 -9
View File
@@ -57,10 +57,29 @@ struct ChannelSendLowering : OpRewritePattern<spatial::SpatChannelSendOp> {
}
};
struct ChannelReceiveLowering : OpRewritePattern<spatial::SpatChannelReceiveOp> {
struct HostStoreSyncLowering : OpRewritePattern<spatial::SpatHostStoreSyncOp> {
using OpRewritePattern::OpRewritePattern;
LogicalResult matchAndRewrite(spatial::SpatChannelReceiveOp op, PatternRewriter& rewriter) const override {
LogicalResult matchAndRewrite(spatial::SpatHostStoreSyncOp op, PatternRewriter& rewriter) const override {
auto sizeAttr = getTensorSizeInBytesAttr(rewriter, op.getOperation(), op.getInput());
auto hostBuffer = getPipelineHostBuffer(rewriter, op);
if (failed(sizeAttr) || failed(hostBuffer))
return failure();
Value zero = arith::ConstantIndexOp::create(rewriter, op.getLoc(), 0);
pim::PimMemCopyDevToHostOp::create(
rewriter, op.getLoc(), hostBuffer->getType(), op.getHostOffset(), zero,
*hostBuffer, op.getInput(), *sizeAttr);
auto sync = pim::PimSyncOp::create(
rewriter, op.getLoc(), op.getTargetCoreId(), op.getEventRegister());
copyRaptorDebugAttrs(op.getOperation(), sync.getOperation());
rewriter.eraseOp(op);
return success();
}
};
template <typename ReceiveOp, typename CreateReceive>
static LogicalResult lowerReceive(
ReceiveOp op, PatternRewriter& rewriter, CreateReceive createReceive) {
if (op->use_empty()) {
rewriter.eraseOp(op);
return success();
@@ -86,12 +105,11 @@ struct ChannelReceiveLowering : OpRewritePattern<spatial::SpatChannelReceiveOp>
if (failed(sizeAttr))
return failure();
Value zero = arith::ConstantIndexOp::create(rewriter, op.getLoc(), 0);
auto receive = pim::PimReceiveOp::create(
rewriter, op.getLoc(), op.getResult().getType(), outputBuffer, zero, *sizeAttr, op.getSourceCoreId());
copyRaptorDebugAttrs(op.getOperation(), receive.getOperation());
Value received = receive.getOutput();
auto received = createReceive(outputBuffer, zero, *sizeAttr);
if (failed(received))
return failure();
if (!destinationInsert) {
rewriter.replaceOp(op, received);
rewriter.replaceOp(op, *received);
return success();
}
@@ -99,10 +117,42 @@ struct ChannelReceiveLowering : OpRewritePattern<spatial::SpatChannelReceiveOp>
Value targetOffset = createDestinationByteOffset(rewriter, destinationInsert);
auto copy = pim::PimMemCopyOp::create(
rewriter, op.getLoc(), destinationInsert.getDestType(), targetOffset, zero,
destinationInsert.getDest(), received, *sizeAttr);
destinationInsert.getDest(), *received, *sizeAttr);
rewriter.replaceOp(destinationInsert, copy.getOutput());
rewriter.eraseOp(op);
return success();
}
struct ChannelReceiveLowering : OpRewritePattern<spatial::SpatChannelReceiveOp> {
using OpRewritePattern::OpRewritePattern;
LogicalResult matchAndRewrite(spatial::SpatChannelReceiveOp op, PatternRewriter& rewriter) const override {
return lowerReceive(op, rewriter, [&](Value outputBuffer, Value zero, IntegerAttr sizeAttr) -> FailureOr<Value> {
auto receive = pim::PimReceiveOp::create(
rewriter, op.getLoc(), op.getResult().getType(), outputBuffer, zero,
sizeAttr, op.getSourceCoreId());
copyRaptorDebugAttrs(op.getOperation(), receive.getOperation());
return receive.getOutput();
});
}
};
struct HostWaitLoadLowering : OpRewritePattern<spatial::SpatHostWaitLoadOp> {
using OpRewritePattern::OpRewritePattern;
LogicalResult matchAndRewrite(spatial::SpatHostWaitLoadOp op, PatternRewriter& rewriter) const override {
return lowerReceive(op, rewriter, [&](Value outputBuffer, Value zero, IntegerAttr sizeAttr) -> FailureOr<Value> {
auto hostBuffer = getPipelineHostBuffer(rewriter, op);
if (failed(hostBuffer))
return failure();
auto wait = pim::PimWaitOp::create(
rewriter, op.getLoc(), op.getEventRegister(),
rewriter.getI32IntegerAttr(1));
copyRaptorDebugAttrs(op.getOperation(), wait.getOperation());
return pim::PimMemCopyHostToDevOp::create(
rewriter, op.getLoc(), outputBuffer.getType(), zero,
op.getHostOffset(), outputBuffer, *hostBuffer, sizeAttr).getOutput();
});
}
};
@@ -148,7 +198,9 @@ struct ConcatLowering : OpRewritePattern<spatial::SpatConcatOp> {
} // namespace
void populateChannelLoweringPatterns(RewritePatternSet& patterns) {
patterns.add<ChannelSendLowering, ChannelReceiveLowering, ExtractRowsLowering, ConcatLowering>(patterns.getContext());
patterns.add<ChannelSendLowering, ChannelReceiveLowering,
HostStoreSyncLowering, HostWaitLoadLowering,
ExtractRowsLowering, ConcatLowering>(patterns.getContext());
}
} // namespace onnx_mlir
src/PIM/Conversion/SpatialToPim/ReturnPathNormalization.cpp
+4
View File
@@ -859,6 +859,10 @@ void raptor::SpatialToPimPass::replaceReturnWithOutputBuffers(func::ReturnOp ret
markOpToRemove(receiveOp);
return;
}
if (auto receiveOp = dyn_cast<spatial::SpatHostWaitLoadOp>(op)) {
markOpToRemove(receiveOp);
return;
}
};
SmallVector<Value> originalOperands(returnOp.getOperands().begin(), returnOp.getOperands().end());
src/PIM/Conversion/SpatialToPim/SpatialToPimPass.cpp
+23
View File
@@ -126,6 +126,8 @@ void onnx_mlir::raptor::SpatialToPimPass::runOnOperation() {
spatial::SpatConcatOp,
spatial::SpatChannelReceiveOp,
spatial::SpatChannelSendOp,
spatial::SpatHostStoreSyncOp,
spatial::SpatHostWaitLoadOp,
spatial::SpatExtractRowsOp>();
RewritePatternSet initialPatterns(ctx);
@@ -140,6 +142,12 @@ void onnx_mlir::raptor::SpatialToPimPass::runOnOperation() {
populateGlobalTensorMaterializationPatterns(globalTensorPatterns);
walkAndApplyPatterns(moduleOp, std::move(globalTensorPatterns));
if (funcOp->hasAttr(kPipelineHostBufferBytesAttrName)
&& failed(materializePipelineHostBuffer(funcOp, rewriter))) {
signalPassFailure();
return;
}
auto returnOp = cast<func::ReturnOp>(funcOp.front().getTerminator());
addReturnOutputBuffers(returnOp, rewriter);
if (failed(allocateAndInitializeCoreLocalVariables(funcOp, rewriter))) {
@@ -182,6 +190,17 @@ void onnx_mlir::raptor::SpatialToPimPass::runOnOperation() {
continue;
}
}
SmallVector<spatial::SpatHostWaitLoadOp> hostWaitLoadOps;
for (auto op : funcOp.getOps<spatial::SpatHostWaitLoadOp>())
hostWaitLoadOps.push_back(op);
for (auto op : hostWaitLoadOps) {
bool onlyPendingRemovalUsers = llvm::all_of(
op->getUsers(), [&](Operation* user) {
return llvm::is_contained(operationsToRemove, user);
});
if (onlyPendingRemovalUsers)
markOpToRemove(op);
}
RewritePatternSet coreBodyPatterns(ctx);
populateCoreBodyPatterns(coreBodyPatterns);
@@ -202,6 +221,8 @@ void onnx_mlir::raptor::SpatialToPimPass::runOnOperation() {
spatial::SpatConcatOp,
spatial::SpatChannelReceiveOp,
spatial::SpatChannelSendOp,
spatial::SpatHostStoreSyncOp,
spatial::SpatHostWaitLoadOp,
spatial::SpatExtractRowsOp>();
SmallVector<pim::PimCoreOp> coreOps;
@@ -251,6 +272,8 @@ void onnx_mlir::raptor::SpatialToPimPass::runOnOperation() {
communicationTarget.addIllegalOp<spatial::SpatConcatOp,
spatial::SpatChannelReceiveOp,
spatial::SpatChannelSendOp,
spatial::SpatHostStoreSyncOp,
spatial::SpatHostWaitLoadOp,
spatial::SpatExtractRowsOp>();
RewritePatternSet communicationPatterns(ctx);