pytorch
ed2eb13d - [inductor] Create triton_helpers module for helper functions (#99880)

Commit
2 years ago
[inductor] Create triton_helpers module for helper functions (#99880) This changes codegen of `torch.prod` from: ```python tl.reduce(tmp2, 1, _prod_accumulate)[:, None] ``` where `_prod_accumulate` is defined elsewhere, to ```python triton_helpers.prod(tmp2, 1)[:, None] ``` A quirk I uncovered though is that `TritonCodeCache` breaks if you define any new symbol beginning with `triton_`, since it assumes that must be the kernel name. Instead, I've made the kernel name an explicit argument to `async_compile.triton` so it doesn't have to guess. Pull Request resolved: https://github.com/pytorch/pytorch/pull/99880 Approved by: https://github.com/ngimel
Author
Committer
Parents
Loading