Skip to content

Add a portable _fft_r2c kernel - #22055

Open
msluszniak wants to merge 1 commit into
pytorch:mainfrom
msluszniak:ms/fft-r2c-portable
Open

Add a portable _fft_r2c kernel#22055
msluszniak wants to merge 1 commit into
pytorch:mainfrom
msluszniak:ms/fft-r2c-portable

Conversation

@msluszniak

Copy link
Copy Markdown
Contributor

Summary

Addresses the first half of #21950.

_fft_r2c.out only had an optimized (pocketfft-backed) kernel. A program that leaves it undelegated builds fine and writes a .pte, but a runtime carrying only the portable kernels cannot load it. The failure surfaces as 0x14 OperatorMissing from load_method, long after the export reported success.

_fft_r2c.out and _fft_c2r.out are two of only three ops in optimized.yaml with no counterpart in the portable functions.yaml:

portable ops: 206   optimized ops: 23
in optimized but NOT in portable (3):
   _fft_c2r.out
   _fft_r2c.out
   linear.out

The third, linear.out, ExecuTorch decomposes anyway, so the FFT pair is the real gap.

Approach

This is a direct O(n^2) evaluation of the transform sum, not a fast Fourier transform, in keeping with the portable library's role as the dependency-free reference; kernels/optimized keeps the asymptotically faster pocketfft path for anyone who links it. Audio front-ends transforming a few hundred points per frame, which is where this op tends to show up, are the intended case.

Two details worth calling out:

  • Multi-dimensional transforms run the real transform along the last requested dimension and complex transforms along the rest, matching pocketfft's multi-axis r2c. A complex pass has to read a whole line before overwriting it; lines up to 128 elements (2 KB for double) use a stack buffer and longer ones ask the runtime for temporary memory. The single-dimension case, which is what torch.fft.rfft lowers to, needs no line buffer at all.
  • The quarter-turn twiddle factors are returned exactly rather than through cos/sin, so a real input's Nyquist bin comes out with a zero imaginary part instead of rounding noise around 1e-16. That is what pocketfft produces and what the existing tests expect.

Test plan

Registers the shared op_fft_r2c_test for portable alongside aten and optimized. All five cases pass against both kernel libraries:

$ ./cmake-out/kernels/test/portable_kernels_test --gtest_filter='OpFftR2c*'
[  PASSED  ] 5 tests.
$ ./cmake-out/kernels/test/optimized_kernels_test --gtest_filter='OpFftR2c*'
[  PASSED  ] 5 tests.

End to end on the model from the issue (torch.fft.rfft(x).abs().pow(2) over [8, 512], exported with no partitioner so the op stays on the CPU), run through executor_runner built with portable kernels only:

result
before exit 134 (abort on load)
after exit 0, output [262144., 1.2e-27, 2.4e-27, ...]

which is the correct abs(rfft(ones))**2: 512**2 in bin 0 and zero elsewhere.

Accuracy against numpy, since the tests above use exactly representable values:

case max abs error
rfft, length 5 (odd, no Nyquist bin, no quarter-turn twiddles) 2.5e-15
rfft, 4x512 1.8e-13 (1.4e-15 relative to the largest output)
rfft2, 6x8 1.8e-14

I wanted to add the length-5 case as a regular test, but tensors_are_close() falls through to a bitwise memcmp for complex dtypes, so a complex comparison cannot have a tolerance and only exactly representable expected values can pass. That is why the existing cases all use small integers. Happy to fix that separately if it would be welcome.

Not addressed here

The issue also asks for lowering to fail when an op is left undelegated and no kernel exists, rather than emitting a .pte that dies at load. That one is a change to the AOT/runtime contract, since the export side does not know which kernel library the runtime will link (portable, optimized, or a selective build), so it needs its own discussion rather than riding along here.

_fft_r2c.out only had an optimized (pocketfft-backed) kernel. A program
that leaves it undelegated builds fine and writes a .pte, but a runtime
carrying only the portable kernels cannot load it: the failure surfaces
as 0x14 OperatorMissing from load_method, or an abort in
executor_runner, long after the export reported success.

_fft_r2c.out and _fft_c2r.out are two of only three ops in
optimized.yaml with no counterpart in the portable functions.yaml (the
third is linear.out, which ExecuTorch decomposes anyway).

This is a direct O(n^2) evaluation of the transform sum, not a fast
Fourier transform, in keeping with the portable library's role as the
dependency-free reference: kernels/optimized keeps the asymptotically
faster pocketfft path for anyone who links it. Audio front-ends
transforming a few hundred points per frame, which is where this op
tends to show up, are the intended case.

Multi-dimensional transforms run the real transform along the last
requested dimension and complex transforms along the rest, matching
pocketfft's multi-axis r2c. A complex pass has to read a whole line
before overwriting it; lines up to 128 elements use a stack buffer, and
longer ones ask the runtime for temporary memory. The single-dimension
case, which is what torch.fft.rfft lowers to, needs no line buffer at
all.

The quarter-turn twiddle factors are returned exactly instead of through
cos/sin, so a real input's Nyquist bin comes out with a zero imaginary
part rather than rounding noise around 1e-16, which is what pocketfft
produces and what the existing tests expect.

Checked against numpy: max absolute error 2.5e-15 for a length-5
transform, 1.8e-13 on a 4x512 transform (1.4e-15 relative to the largest
output), and 1.8e-14 for a 6x8 rfft2.

Registers the shared op_fft_r2c_test for portable as well as aten and
optimized; all five cases pass against both kernel libraries.
@pytorch-bot

pytorch-bot Bot commented Aug 22, 2026

Copy link
Copy Markdown

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/22055

Note: Links to docs will display an error until the docs builds have been completed.

⚠️ 20 Awaiting Approval

As of commit f123158 with merge base fbd4bbf (image):

AWAITING APPROVAL - The following workflows need approval before CI can run:

This comment was automatically generated by Dr. CI and updates every 15 minutes.

@meta-cla meta-cla Bot added the CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. label Aug 22, 2026
@github-actions

Copy link
Copy Markdown

This PR needs a release notes: label

If your change should be included in the release notes (i.e. would users of this library care about this change?), please use a label starting with release notes:. This helps us keep track and include your important work in the next release notes.

To add a label, you can comment to pytorchbot, for example
@pytorchbot label "release notes: none"

For more information, see
https://github.com/pytorch/pytorch/wiki/PyTorch-AutoLabel-Bot#why-categorize-for-release-notes-and-how-does-it-work.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed.

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants