Use multi-dimensional cuFFT transforms to improve FFT performance (#61203)
Summary:
Benchmark and numerical accuracy tests on A100 and V100 are available at https://github.com/xwang233/code-snippet/tree/master/fft-61203.
I've checked the FFT results for different shapes/dims and different `dim` arg for `rfftn` and `irfftn` before and after this PR, and they all numerically matched.
With this PR, about 10%~15% performance gain is expected on commonly used shapes and dims.
Pull Request resolved: https://github.com/pytorch/pytorch/pull/61203
Reviewed By: heitorschueroff
Differential Revision: D29996244
Pulled By: zou3519
fbshipit-source-id: 02c9862eaa1ad8f2ae9c7f7448aeb9e23bcda276
