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

2
include/triton/Dialect/TritonGPU/Transforms/Passes.h
View File

@@ -4,7 +4,7 @@
#include "mlir/Pass/Pass.h"
namespace mlir {
std::unique_ptr<Pass> createTritonGPUPipelinePass();
std::unique_ptr<Pass> createTritonGPUPipelinePass(int numStages);
namespace triton {
namespace gpu {

6
include/triton/Dialect/TritonGPU/Transforms/Passes.td
View File

@@ -24,6 +24,12 @@ def TritonGPUPipeline : Pass<"tritongpu-pipeline", "mlir::ModuleOp"> {
let dependentDialects = ["mlir::triton::gpu::TritonGPUDialect",
"mlir::scf::SCFDialect",
"mlir::arith::ArithmeticDialect"];
let options = [
Option<"numStages", "num-stages",
"int32_t", /*default*/"2",
"number of pipeline stages">
];
}
def TritonGPUCombineOps : Pass<"tritongpu-combine", "mlir::ModuleOp"> {