nnicolosi/Raptor
1
0
Fork 0
Code Issues Pull Requests Actions 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/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);