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

make numStages an option in PipelinePass

Browse Source
This commit is contained in:
Yan Da
2022-05-23 12:47:55 +08:00
parent 39b1235082
commit 36c45ec687
6 changed files with 23 additions and 9 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -311,9 +311,13 @@ scf::ForOp LoopPipeliner::createNewForOp() {
// ref: mlir/lib/Dialect/SCF/Transforms/LoopPipelining.cpp
struct PipelinePass : public TritonGPUPipelineBase<PipelinePass> {
PipelinePass() = default;
PipelinePass(int numStages) {
this->numStages = numStages;
}
void runOnOperation() override {
// TODO: collect numStages from ModuleOp
int numStages = 2;
int numStages = this->numStages;
if (numStages <= 1)
return;
@@ -337,6 +341,6 @@ struct PipelinePass : public TritonGPUPipelineBase<PipelinePass> {
};
} // anonymous namespace
std::unique_ptr<Pass> mlir::createTritonGPUPipelinePass() {
return std::make_unique<PipelinePass>();
std::unique_ptr<Pass> mlir::createTritonGPUPipelinePass(int numStages) {
return std::make_unique<PipelinePass>(numStages);
}