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

fix memory allocation in pim codegen

Browse Source 
fix crossbar allocation to only consider weights from vmm and mvm
This commit is contained in:
NiccoloN
2026-04-21 13:31:10 +02:00
parent 85e2750d6c
commit 25ade1bd63
3 changed files with 65 additions and 22 deletions
Download Patch File Download Diff File Expand all files Collapse all files

17
src/PIM/Dialect/Pim/Transforms/Bufferization/PimBufferizationPass.cpp
View File

@@ -93,15 +93,22 @@ void PimBufferizationPass::runOnOperation() {
}
void PimBufferizationPass::annotateWeightsMemrefs(ModuleOp moduleOp, func::FuncOp funcOp) const {
funcOp.walk([&](memref::GetGlobalOp getGlobalOp) {
bool isAlwaysWeight = !getGlobalOp->getUsers().empty()
&& all_of(getGlobalOp->getUsers(), [](auto user) -> bool { return isa<PimCoreOp>(user); });
if (isAlwaysWeight) {
funcOp.walk([&](PimCoreOp coreOp) {
auto annotateWeight = [&](unsigned weightIndex) {
if (weightIndex >= coreOp.getWeights().size())
return;
Value weight = coreOp.getWeights()[weightIndex];
auto getGlobalOp = weight.getDefiningOp<memref::GetGlobalOp>();
if (!getGlobalOp)
return;
auto globalMemrefOp = lookupGlobalForGetGlobal(moduleOp, getGlobalOp);
assert("Weights must be constants" && globalMemrefOp.getConstant());
markWeightAlways(getGlobalOp);
markWeightAlways(globalMemrefOp);
}
};
coreOp.walk([&](PimMVMOp mvmOp) { annotateWeight(mvmOp.getWeightIndex()); });
coreOp.walk([&](PimVMMOp vmmOp) { annotateWeight(vmmOp.getWeightIndex()); });
});
}