From f565d0db5b987091e5ac4e8f9470f1ed6b32d95f Mon Sep 17 00:00:00 2001 From: Hiroshi Shinaoka Date: Sun, 19 Jul 2026 00:56:54 +0900 Subject: [PATCH] fix: stabilize interleaved FFT CPU execution --- crates/cubek-fft/src/fft/cfft_interleaved.rs | 1 + crates/cubek-fft/src/fft/irfft_interleaved.rs | 1 + crates/cubek-fft/src/fft/rfft_interleaved.rs | 1 + crates/cubek-fft/tests/fft/interleaved_cfft.rs | 14 ++++++++++---- .../tests/fft/interleaved_validation.rs | 16 ++++++++++++---- 5 files changed, 25 insertions(+), 8 deletions(-) diff --git a/crates/cubek-fft/src/fft/cfft_interleaved.rs b/crates/cubek-fft/src/fft/cfft_interleaved.rs index f279601f5..ad98f9bae 100644 --- a/crates/cubek-fft/src/fft/cfft_interleaved.rs +++ b/crates/cubek-fft/src/fft/cfft_interleaved.rs @@ -376,6 +376,7 @@ fn cfft_interleaved_shared_kernel( output_im.write_checked(k, shared_im[k] * scale); k += threads_per_cube; } + sync_cube(); } /// First four-step pass over the strided N1 dimension of each C32 window. diff --git a/crates/cubek-fft/src/fft/irfft_interleaved.rs b/crates/cubek-fft/src/fft/irfft_interleaved.rs index 6245b2893..3d0c11522 100644 --- a/crates/cubek-fft/src/fft/irfft_interleaved.rs +++ b/crates/cubek-fft/src/fft/irfft_interleaved.rs @@ -299,4 +299,5 @@ fn irfft_interleaved_kernel( signal_view.write_checked(i, shared_re[i] * scale); i += threads_per_cube; } + sync_cube(); } diff --git a/crates/cubek-fft/src/fft/rfft_interleaved.rs b/crates/cubek-fft/src/fft/rfft_interleaved.rs index 7f62ff26c..7c465ab06 100644 --- a/crates/cubek-fft/src/fft/rfft_interleaved.rs +++ b/crates/cubek-fft/src/fft/rfft_interleaved.rs @@ -319,4 +319,5 @@ fn rfft_interleaved_kernel( spectrum_im.write_checked(k, shared_im[k] * scale); k += threads_per_cube; } + sync_cube(); } diff --git a/crates/cubek-fft/tests/fft/interleaved_cfft.rs b/crates/cubek-fft/tests/fft/interleaved_cfft.rs index 5afcd473a..392720a58 100644 --- a/crates/cubek-fft/tests/fft/interleaved_cfft.rs +++ b/crates/cubek-fft/tests/fft/interleaved_cfft.rs @@ -103,7 +103,7 @@ fn round_trip(shape: Vec, dim: usize, normalization: FftNormalization) { run_round_trip(&client, shape, dim, normalization, 1e-4); } -fn test_max_shared_fft_n(client: &ComputeClient) -> usize { +fn device_max_shared_fft_n(client: &ComputeClient) -> usize { let max_elems = client.properties().hardware.max_shared_memory_size / (2 * core::mem::size_of::()); if max_elems.is_power_of_two() { @@ -113,9 +113,15 @@ fn test_max_shared_fft_n(client: &ComputeClient) -> usize { } } +fn standard_test_shared_fft_n(client: &ComputeClient) -> usize { + // The CPU backend reports system RAM as its shared-memory limit. Using that + // value directly would turn this routine-path test into a multi-gigabyte FFT. + device_max_shared_fft_n(client).min(256) +} + #[cfg(feature = "heavy")] fn first_four_step_n(client: &ComputeClient) -> usize { - 2 * test_max_shared_fft_n(client) + 2 * device_max_shared_fft_n(client) } #[test] @@ -190,9 +196,9 @@ fn cfft_interleaved_ortho_round_trip() { } #[test] -fn cfft_interleaved_shared_memory_boundary_round_trip() { +fn cfft_interleaved_shared_memory_path_round_trip() { let client = ::client(&Default::default()); - let n_fft = test_max_shared_fft_n(&client); + let n_fft = standard_test_shared_fft_n(&client); run_round_trip(&client, vec![1, n_fft, 1], 1, FftNormalization::ByN, 0.03); } diff --git a/crates/cubek-fft/tests/fft/interleaved_validation.rs b/crates/cubek-fft/tests/fft/interleaved_validation.rs index 43ace6c20..da3c4c3e0 100644 --- a/crates/cubek-fft/tests/fft/interleaved_validation.rs +++ b/crates/cubek-fft/tests/fft/interleaved_validation.rs @@ -105,7 +105,7 @@ fn non_contiguous_c32_extent_includes_the_last_imaginary_scalar() { assert_eq!(complex.physical_scalar_len(), 16); } -fn first_unsupported_real_fft_n() -> usize { +fn first_unsupported_real_fft_n() -> Option { let client = ::client(&Default::default()); let max_elems = client.properties().hardware.max_shared_memory_size / (2 * core::mem::size_of::()); @@ -114,14 +114,18 @@ fn first_unsupported_real_fft_n() -> usize { } else { max_elems.next_power_of_two() >> 1 }; - max_shared.saturating_mul(max_shared).saturating_mul(4) + max_shared.checked_mul(max_shared)?.checked_mul(4) } #[test] fn oversized_rfft_is_rejected_for_allocating_and_caller_owned_apis_without_allocating_data() { let client = ::client(&Default::default()); let dtype = f32::as_type_native_unchecked().storage_type(); - let n_fft = first_unsupported_real_fft_n(); + let Some(n_fft) = first_unsupported_real_fft_n() else { + // The CPU backend uses system RAM as its shared-memory limit, so its + // reported maximum can exceed every FFT length representable by usize. + return; + }; let signal = TensorHandle::new_contiguous(vec![0, n_fft], client.empty(0), dtype); assert!(matches!( @@ -148,7 +152,11 @@ fn oversized_rfft_is_rejected_for_allocating_and_caller_owned_apis_without_alloc fn oversized_irfft_is_rejected_for_allocating_and_caller_owned_apis_without_allocating_data() { let client = ::client(&Default::default()); let dtype = f32::as_type_native_unchecked().storage_type(); - let n_fft = first_unsupported_real_fft_n(); + let Some(n_fft) = first_unsupported_real_fft_n() else { + // See the RFFT case above: there is no representable unsupported + // length to exercise for a backend with a saturating device limit. + return; + }; let spectrum = ComplexTensorHandle::new_contiguous(vec![0, n_fft / 2 + 1], client.empty(0), dtype) .unwrap();