-
Notifications
You must be signed in to change notification settings - Fork 280
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Add command line argument parsing for reduction dimensions in Triton …
…sum kernel Summary: Add argument parsing for the command line in order to pass in dimension(s) across which the kernel reduces and enable more rigorous testing of different versions of the sum kernel, referencing [torchbenchmark/operators/fb/flash_attention/operator.py](https://www.internalfb.com/code/fbsource/[864a578ce44afdba619d50a352c8ca3b783e05ef]/fbcode/pytorch/benchmark/torchbenchmark/operators/fb/flash_attention/operator.py?lines=84). Inherit the `__init__` function from the parent class `BenchmarkOperator` in order to facilitate command line argument parsing. Change `dim` type to `list` to avoid type issues resulting from `tl.constexpr`. Modify equality checks in kernel and operator to satisfy type requirements for `dim`. Reviewed By: xuzhao9 Differential Revision: D58212366
- Loading branch information
1 parent
2d8999b
commit c9532af
Showing
2 changed files
with
49 additions
and
14 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters