[BE] Accumulator init optimization #4680
Conversation
| @@ -169,4 +169,14 @@ def TritonGPUCombineTensorSelectAndIf: Pass<"tritongpu-combine-tensor-select-and | |||
| "mlir::triton::TritonDialect"]; | |||
| } | |||
|
|
|||
| def TritonGPUOptimizeAccumulatorInit: Pass<"tritongpu-optimize-accumulator-init", "mlir::ModuleOp"> { | |||
There was a problem hiding this comment.
should this go in triton/Dialect/TritonNvidiaGPU/Transforms/Passes.td?
I was trying to keep the code generic enough with the helpers so any HW that supports acc scale or acc init flag could be plugged in easily, that's why I kept it in TritonGPU. I would argue that this is more generic than our Matmul pipeliner, that lives in TritonGPU :) But I agree that we don't have anything else right now that can benefit from it. |
fair enough |
Adding a transformation pass that skips filling the accumulator with zero value if the HW supports accumulator scale or init flag. In such case flag value is created and maintained, and passed to the MMA op indicating if accumulator should be taken into an account when calculating the dot product.
The pass is intended to be generic enough to be reusable between different HW platforms, therefore it is not placed in the Nvidia specific folder, even though it is supporting only Hopper MMA for the moment.