Refactor code for sum Triton kernels

jananisriram commented 3 months ago

Summary: Refactor code to improve readability and logical flow for cases which select the sum Triton kernel implementation to run. Create helper functions for the following cases:

Reduce N-dimensional input to scalar output
Reduce 2-dimensional input to 1-dimensional output
Reduce 3-dimensional input along dimension 1 to 2-dimensional output

Add command line argument parsing for the input_dim parameter, which specifies the number of dimensions desired in kernel inputs.

Modify absolute tolerance to account for floating-point operation error.

Reviewed By: jbschlosser

Differential Revision: D58488137

facebook-github-bot commented 3 months ago

This pull request was exported from Phabricator. Differential Revision: D58488137

facebook-github-bot commented 3 months ago

This pull request has been merged in pytorch/benchmark@339ccfdec0d47a5372d6fd64f8ddbe7667b030bf.

pytorch / benchmark

Refactor code for sum Triton kernels #2303