From f123158a509b2197cdbc540ed886ccd47362ac52 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Mateusz=20S=C5=82uszniak?= Date: Sat, 22 Aug 2026 19:53:10 +0200 Subject: [PATCH] Add a portable _fft_r2c kernel _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. --- kernels/portable/cpu/op_fft_r2c.cpp | 303 ++++++++++++++++++ kernels/portable/functions.yaml | 5 + kernels/test/CMakeLists.txt | 1 + kernels/test/targets.bzl | 2 +- .../kernels/portable/op_registration_util.bzl | 7 + 5 files changed, 317 insertions(+), 1 deletion(-) create mode 100644 kernels/portable/cpu/op_fft_r2c.cpp diff --git a/kernels/portable/cpu/op_fft_r2c.cpp b/kernels/portable/cpu/op_fft_r2c.cpp new file mode 100644 index 00000000000..da23342b4a0 --- /dev/null +++ b/kernels/portable/cpu/op_fft_r2c.cpp @@ -0,0 +1,303 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#include +#include + +#include +#include +#include +#include + +namespace torch::executor::native { + +namespace { + +constexpr double kTwoPi = 6.283185307179586476925286766559; + +// A complex-to-complex pass has to read a whole line before it can overwrite +// it. Lines up to this length go through a stack buffer, which keeps the buffer +// at 2 KB for double; longer ones ask the runtime for temporary memory. Only +// multi-dimensional transforms reach this path at all: the single-dimension +// case, which is what torch.fft.rfft lowers to, needs no line buffer. +constexpr size_t kStackLineLimit = 128; + +// Mirrors ATen's fft_norm_mode (ATen/native/SpectralOpsUtils.h), which is how +// the normalization argument is encoded. +enum class fft_norm_mode { + none, // No normalization + by_root_n, // Divide by sqrt(signal_size) + by_n, // Divide by signal_size +}; + +template +std::optional compute_fct( + KernelRuntimeContext& ctx, + const Tensor& t, + IntArrayRef dim, + int64_t normalization) { + constexpr auto one = static_cast(1); + const auto mode = static_cast(normalization); + if (mode == fft_norm_mode::none) { + return one; + } + int64_t n = 1; + for (auto idx : dim) { + n *= t.sizes()[idx]; + } + switch (mode) { + case fft_norm_mode::none: + return one; + case fft_norm_mode::by_n: + return one / static_cast(n); + case fft_norm_mode::by_root_n: + return one / std::sqrt(static_cast(n)); + } + ET_KERNEL_CHECK_MSG( + ctx, + false, + InvalidArgument, + std::nullopt, + "Unsupported normalization type: %" PRId64, + normalization); +} + +// cos and sin of -2*pi*idx/n. +// +// The quarter turns are returned exactly rather than through cos/sin, so that a +// real input's Nyquist bin comes out with a zero imaginary part instead of +// rounding noise on the order of 1e-16. Reducing idx modulo n first also keeps +// the angle inside one period, which matters for accuracy once k * j is large. +void twiddle(size_t idx, size_t n, double& cos_out, double& sin_out) { + idx %= n; + if (idx == 0) { + cos_out = 1.0; + sin_out = 0.0; + } else if (2 * idx == n) { + cos_out = -1.0; + sin_out = 0.0; + } else if (4 * idx == n) { + cos_out = 0.0; + sin_out = -1.0; + } else if (4 * idx == 3 * n) { + cos_out = 0.0; + sin_out = 1.0; + } else { + const double angle = + -kTwoPi * static_cast(idx) / static_cast(n); + cos_out = std::cos(angle); + sin_out = std::sin(angle); + } +} + +// Offset of the start of the line_index'th line along `axis`, for a tensor with +// the given sizes and strides. Lines are enumerated over every dimension except +// `axis`, so the same index names the same logical line in two tensors that +// agree on all dimensions but that one. +size_t line_offset( + size_t line_index, + ArrayRef sizes, + ArrayRef strides, + size_t axis) { + size_t offset = 0; + for (size_t d = sizes.size(); d-- > 0;) { + if (d == axis) { + continue; + } + const size_t size = static_cast(sizes[d]); + offset += (line_index % size) * static_cast(strides[d]); + line_index /= size; + } + return offset; +} + +// Forward real-to-complex DFT along `axis`, writing the onesided output. +// The normalization factor is folded in here: every later pass is linear, so +// applying it once at the front scales the whole transform. +template +void dft_r2c_axis(const Tensor& in, Tensor& out, size_t axis, T fct) { + using C = executorch::runtime::etensor::complex; + const T* const in_data = in.const_data_ptr(); + C* const out_data = out.mutable_data_ptr(); + + const size_t n = static_cast(in.size(axis)); + const size_t n_out = static_cast(out.size(axis)); + const size_t in_stride = static_cast(in.strides()[axis]); + const size_t out_stride = static_cast(out.strides()[axis]); + const size_t num_lines = n == 0 ? 0 : static_cast(in.numel()) / n; + + for (size_t line = 0; line < num_lines; ++line) { + const size_t in_off = line_offset(line, in.sizes(), in.strides(), axis); + const size_t out_off = line_offset(line, out.sizes(), out.strides(), axis); + for (size_t k = 0; k < n_out; ++k) { + double real = 0; + double imag = 0; + for (size_t j = 0; j < n; ++j) { + double c = 0; + double s = 0; + twiddle(k * j, n, c, s); + const double x = static_cast(in_data[in_off + j * in_stride]); + real += x * c; + imag += x * s; + } + out_data[out_off + k * out_stride] = + C{static_cast(real * static_cast(fct)), + static_cast(imag * static_cast(fct))}; + } + } +} + +// In-place forward complex-to-complex DFT along `axis`. `scratch` must hold at +// least out.size(axis) elements. +template +void dft_c2c_axis_(Tensor& out, size_t axis, void* scratch) { + using C = executorch::runtime::etensor::complex; + C* const out_data = out.mutable_data_ptr(); + C* const line_buf = static_cast(scratch); + + const size_t n = static_cast(out.size(axis)); + const size_t stride = static_cast(out.strides()[axis]); + const size_t num_lines = n == 0 ? 0 : static_cast(out.numel()) / n; + + for (size_t line = 0; line < num_lines; ++line) { + const size_t off = line_offset(line, out.sizes(), out.strides(), axis); + for (size_t j = 0; j < n; ++j) { + line_buf[j] = out_data[off + j * stride]; + } + for (size_t k = 0; k < n; ++k) { + double real = 0; + double imag = 0; + for (size_t j = 0; j < n; ++j) { + double c = 0; + double s = 0; + twiddle(k * j, n, c, s); + const double xr = static_cast(line_buf[j].real_); + const double xi = static_cast(line_buf[j].imag_); + real += xr * c - xi * s; + imag += xr * s + xi * c; + } + out_data[off + k * stride] = + C{static_cast(real), static_cast(imag)}; + } + } +} + +} // namespace + +// Reference discrete Fourier transform. +// +// This is a direct O(n^2) evaluation of the transform sum, not a fast Fourier +// transform. kernels/optimized provides a pocketfft-backed _fft_r2c.out that is +// asymptotically faster; this exists so that a graph containing _fft_r2c can be +// run by a build that only has the portable kernels, rather than failing to +// load with OperatorMissing. Audio front-ends that transform a few hundred +// points per frame are the intended case. +Tensor& _fft_r2c_out( + KernelRuntimeContext& ctx, + const Tensor& in, + IntArrayRef dim, + int64_t normalization, + bool onesided, + Tensor& out) { + auto in_sizes = in.sizes(); + ET_KERNEL_CHECK(ctx, in.dim() <= kTensorDimensionLimit, InvalidArgument, out); + ET_KERNEL_CHECK(ctx, !dim.empty(), InvalidArgument, out); + ET_KERNEL_CHECK( + ctx, tensors_have_same_dim_order(in, out), InvalidArgument, out); + + ET_KERNEL_CHECK_MSG( + ctx, + onesided, + InvalidArgument, + out, + "onesided=False is not supported yet in _fft_r2c"); + + ET_KERNEL_CHECK_MSG( + ctx, + out.scalar_type() == executorch::runtime::toComplexType(in.scalar_type()), + InvalidArgument, + out, + "the output type for _fft_r2c must be the Complex type corresponding to the input type"); + + for (auto d : dim) { + ET_KERNEL_CHECK_MSG( + ctx, + d >= 0 && d < in.dim(), + InvalidArgument, + out, + "dims must be in bounds (got %" PRId64 ")", + d); + } + + std::array out_sizes_storage; + executorch::runtime::Span out_sizes( + out_sizes_storage.data(), in_sizes.size()); + std::copy(in_sizes.begin(), in_sizes.end(), out_sizes.begin()); + out_sizes[dim.back()] = out_sizes[dim.back()] / 2 + 1; + + ET_KERNEL_CHECK_MSG( + ctx, + resize_tensor( + out, + executorch::runtime::ArrayRef( + out_sizes.data(), out_sizes.size())) == Error::Ok, + InvalidArgument, + out, + "Failed to resize output tensor (last dim %d).", + out_sizes[dim.back()]); + + // NOTE: as of this writing, upstream PyTorch only supports float/double, so + // we follow suit. + ET_SWITCH_FLOAT_TYPES(in.scalar_type(), ctx, "_fft_r2c.out", CTYPE_IN, [&] { + auto fct = compute_fct(ctx, in, dim, normalization); + if (!fct) { + // Check failed, just bail out of the lambda. + return; + } + + // The real transform runs along the last requested dimension, which is the + // one that is halved; the remaining dimensions are complex transforms of + // the result, matching pocketfft's multi-axis r2c. + const size_t real_axis = static_cast(dim.back()); + dft_r2c_axis(in, out, real_axis, *fct); + + if (dim.size() == 1) { + return; + } + + using Complex = executorch::runtime::etensor::complex; + size_t max_line = 0; + for (size_t i = 0; i + 1 < dim.size(); ++i) { + max_line = std::max(max_line, static_cast(out.size(dim[i]))); + } + + std::array stack_buf; + void* line_buf = stack_buf.data(); + if (max_line > kStackLineLimit) { + Result scratch = ctx.allocate_temp(max_line * sizeof(Complex)); + ET_KERNEL_CHECK_MSG( + ctx, + scratch.ok(), + MemoryAllocationFailed, + , + "_fft_r2c needs %zu bytes of temporary memory to transform a " + "dimension of length %zu, but no temp allocator is available", + max_line * sizeof(Complex), + max_line); + line_buf = scratch.get(); + } + + for (size_t i = 0; i + 1 < dim.size(); ++i) { + dft_c2c_axis_(out, static_cast(dim[i]), line_buf); + } + }); + + return out; +} + +} // namespace torch::executor::native diff --git a/kernels/portable/functions.yaml b/kernels/portable/functions.yaml index ecf62ee3606..61f32677b99 100644 --- a/kernels/portable/functions.yaml +++ b/kernels/portable/functions.yaml @@ -32,6 +32,11 @@ - arg_meta: null kernel_name: torch::executor::_conj_physical_out +- op: _fft_r2c.out + kernels: + - arg_meta: null + kernel_name: torch::executor::_fft_r2c_out + - op: _log_softmax.out kernels: - arg_meta: null diff --git a/kernels/test/CMakeLists.txt b/kernels/test/CMakeLists.txt index deb9fbf92e4..14f76b3964b 100644 --- a/kernels/test/CMakeLists.txt +++ b/kernels/test/CMakeLists.txt @@ -207,6 +207,7 @@ set(all_test_sources "op_exp_test.cpp" "op_expand_copy_test.cpp" "op_expm1_test.cpp" + "op_fft_r2c_test.cpp" "op_fill_test.cpp" "op_flip_test.cpp" "op_floor_divide_test.cpp" diff --git a/kernels/test/targets.bzl b/kernels/test/targets.bzl index 837c7327c4f..a02958850fa 100644 --- a/kernels/test/targets.bzl +++ b/kernels/test/targets.bzl @@ -252,7 +252,7 @@ def define_common_targets(): _common_op_test("op_expand_copy_test", ["aten", "portable"]) _common_op_test("op_expm1_test", ["aten", "portable"]) _common_op_test("op_fft_c2r_test", ["aten", "optimized"]) - _common_op_test("op_fft_r2c_test", ["aten", "optimized"]) + _common_op_test("op_fft_r2c_test", ["aten", "portable", "optimized"]) _common_op_test("op_fill_test", ["aten", "portable"]) _common_op_test("op_flip_test", ["aten", "portable"]) _common_op_test("op_floor_divide_test", ["aten", "portable"]) diff --git a/shim_et/xplat/executorch/kernels/portable/op_registration_util.bzl b/shim_et/xplat/executorch/kernels/portable/op_registration_util.bzl index f1a46616295..a6149f4106a 100644 --- a/shim_et/xplat/executorch/kernels/portable/op_registration_util.bzl +++ b/shim_et/xplat/executorch/kernels/portable/op_registration_util.bzl @@ -574,6 +574,13 @@ ATEN_OPS = ( "//executorch/kernels/portable/cpu/pattern:pattern", ], ), + op_target( + name = "op_fft_r2c", + deps = [ + "//executorch/runtime/core/exec_aten/util:scalar_type_util", + "//executorch/runtime/core/exec_aten/util:tensor_util", + ], + ), op_target( name = "op_fill", deps = [