Repository navigation
[ROCm] Work around a rocFFT twiddle cache bug in multi-dimensional real FFTs - #40389
Merged
copybara-service[bot] merged 1 commit intoSep 2, 2026
Conversation
magaonka-amd
marked this pull request as ready for review
September 2, 2026 05:38
magaonka-amd
force-pushed
the
fix/rocm-multidim-real-fft-twiddle-cache
branch
from
September 2, 2026 06:31
2197e3e to
37bfd65
Compare
âĶal FFTs This is a temporary workaround in the JAX lowering. The defect is in rocFFT and the proper fix belongs there; once that fix ships, gate this on the ROCm version rather than applying it unconditionally. - rocFFT keys its 2D twiddle table cache without the attach_halfN flags that select the table a fused real pre/post-processing kernel needs. A multi-dimensional RFFT or IRFFT therefore reuses a plain complex 2D kernel's table when one of matching lengths and radices was planned earlier in the process and is still alive, and returns silently wrong results. Whichever kernel is planned first wins, so which transform is affected depends on the order transforms run in, and nothing is reported by rocFFT or HIP. - Work around it by lowering multi-dimensional RFFT and IRFFT on ROCm to a 1-D real transform plus a complex transform over the outer axes, which is also NumPy's convention and so needs no Hermitian symmetrization of the input. rocFFT keys its 1D twiddle cache correctly, so a 1-D real transform cannot pick up a complex kernel's table. CUDA is unaffected and keeps the single multi-dimensional C2R lowering. - Fixes tests/fft_test.py::FftTest::testFftn5, which fails on the gfx950 8-GPU CI job whenever an earlier test in the same pytest-xdist worker process plans a matching complex 2D kernel.
magaonka-amd
force-pushed
the
fix/rocm-multidim-real-fft-twiddle-cache
branch
from
September 2, 2026 06:32
37bfd65 to
9034a8a
Compare
tsrw2048
approved these changes
Sep 2, 2026
google-ml-butler
Bot
added
kokoro:force-run
pull ready
mminutoli
approved these changes
Sep 2, 2026
hawkinsp
approved these changes
Sep 2, 2026
This file contains hidden or 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
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
This is a workaround in the JAX lowering. The defect is in rocFFT and I'm working landing a fix there, version gate this once that fix has shipped in a ROCm release.
rocFFT keys its 2D twiddle table cache without the attach_halfN flags that select the table a fused real pre/post-processing kernel needs. A multi-dimensional RFFT or IRFFT therefore reuses a plain complex 2D kernel's table when one of matching lengths and radices was planned earlier in the process and is still alive, and returns silently wrong results. Whichever kernel is planned first wins, so which transform is affected depends on the order transforms run in, and nothing is reported by rocFFT or HIP.
Work around it by lowering multi-dimensional RFFT and IRFFT on ROCm to a 1-D real transform plus a complex transform over the outer axes. rocFFT keys its 1D twiddle cache correctly, so a 1-D real transform cannot pick up a complex kernel's table.
Fixes tests/fft_test.py::FftTest::testFftn5, which fails on the gfx950 8-GPU CI job whenever an earlier test in the same pytest-xdist worker process plans a matching complex 2D kernel.
Added TODO to myself to version gate the workaround once fix lands in ROCm.
My attempt to fix this in our rocfft is WIP : fix(rocfft): include attach_halfN in the 2D twiddle table cache key ROCm/rocm-libraries#11595