gaspersic/triton
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Fix op mapping in pipeline.cpp

Browse Source
This commit is contained in:
Yan Da
2022-05-26 13:57:01 +08:00
parent c529b462f5
commit 41d338d848
1 changed files with 7 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files

15
lib/Dialect/TritonGPU/Transforms/Pipeline.cpp
View File

@@ -163,6 +163,10 @@ LogicalResult LoopPipeliner::initialize() {
} }
void LoopPipeliner::emitPrologue() { void LoopPipeliner::emitPrologue() {
// llvm::errs() << "loads to pipeline...:\n";
// for (Value load : loads)
// llvm::errs() << load << "\n";
OpBuilder builder(forOp); OpBuilder builder(forOp);
for (BlockArgument &arg : forOp.getRegionIterArgs()) { for (BlockArgument &arg : forOp.getRegionIterArgs()) {
OpOperand &operand = forOp.getOpOperandForRegionIterArg(arg); OpOperand &operand = forOp.getOpOperandForRegionIterArg(arg);
@@ -302,7 +306,7 @@ scf::ForOp LoopPipeliner::createNewForOp() {
for (size_t idx = 0; idx < loads.size(); ++idx) { for (size_t idx = 0; idx < loads.size(); ++idx) {
Value load = loads[idx]; Value load = loads[idx];
mapping.lookup(load).replaceAllUsesWith( mapping.lookup(load).replaceAllUsesWith(
newForOp.getRegionIterArgs()[loadIdx+idx]); newForOp.getRegionIterArgs()[loadIdx + idx*(numStages-1)]);
} }
@@ -377,17 +381,12 @@ scf::ForOp LoopPipeliner::createNewForOp() {
SmallVector<Value> yieldValues; SmallVector<Value> yieldValues;
for (Value v : forOp.getBody()->getTerminator()->getOperands()) for (Value v : forOp.getBody()->getTerminator()->getOperands())
yieldValues.push_back(mapping.lookup(v)); yieldValues.push_back(mapping.lookup(v));
// for (int i = 1; i < numStages - 1; ++i) // shift pipelined args by 1
// yieldValues.push_back(newForOp.getRegionIterArgs()[aArgIdx + i]);
// yieldValues.push_back(nextMapping.lookup(info.dotOp.a()));
// for (int i = 1; i < numStages - 1; ++i)
// yieldValues.push_back(newForOp.getRegionIterArgs()[bArgIdx + i]);
// yieldValues.push_back(nextMapping.lookup(info.dotOp.b()));
for (size_t idx = 0; idx < loads.size(); ++idx) { for (size_t idx = 0; idx < loads.size(); ++idx) {
Value load = loads[idx]; Value load = loads[idx];
for (int stage = 1; stage < numStages - 1; ++stage) { for (int stage = 1; stage < numStages - 1; ++stage) {
yieldValues.push_back(newForOp.getRegionIterArgs()[ yieldValues.push_back(newForOp.getRegionIterArgs()[
loadIdx + idx*(numStages-1) + stage-1 loadIdx + idx*(numStages-1) + stage
]); ]);
} }
yieldValues.push_back(nextMapping.lookup(load)); yieldValues.push_back(nextMapping.lookup(load));