PoC: Make the codelets operate entirely in registers

Cargo.toml

@@ -61,4 +61,4 @@ debug = true
 [package.metadata.docs.rs]
 all-features = true
 
-[lints.rust]
+[lints.rust]

src/algorithms/dit.rs

@@ -17,7 +17,7 @@
 use fearless_simd::{dispatch, Simd};
 
 use crate::algorithms::bravo::{bit_rev_bravo_f32, bit_rev_bravo_f64};
-use crate::kernels::codelets::{fft_dit_codelet_32_f32, fft_dit_codelet_32_f64};
+use crate::kernels::codelets::{fft_dit_codelet_16_f64, fft_dit_codelet_32_f32};
 use crate::kernels::dit::*;
 use crate::options::Options;
 use crate::parallel::run_maybe_in_parallel;
@@ -42,10 +42,11 @@ fn recursive_dit_fft_f64<S: Simd>(
     let log_size = size.ilog2() as usize;
 
     if size <= L1_BLOCK_SIZE {
-        // Use FFT-32 codelet to fuse stages 0-4 into a single pass per 32-element chunk
-        let start_stage = if planner.use_codelet_32 {
-            fft_dit_codelet_32_f64(simd, &mut reals[..size], &mut imags[..size]);
-            5
+        // Use FFT-16 codelet to fuse stages 0-3 into a single pass per 16-element chunk
+        let codelet_stages = 4;
+        let start_stage = if stage_twiddle_idx == 0 && size >= power_of_two(codelet_stages) {
+            fft_dit_codelet_16_f64(simd, &mut reals[..size], &mut imags[..size]);
+            codelet_stages
         } else {
             0
         };
@@ -109,9 +110,10 @@ fn recursive_dit_fft_f32<S: Simd>(
 
     if size <= L1_BLOCK_SIZE {
         // Use FFT-32 codelet to fuse stages 0-4 into a single pass per 32-element chunk
-        let start_stage = if planner.use_codelet_32 {
+        let codelet_stages = 5;
+        let start_stage = if stage_twiddle_idx == 0 && size >= power_of_two(codelet_stages) {
             fft_dit_codelet_32_f32(simd, &mut reals[..size], &mut imags[..size]);
-            5
+            codelet_stages
         } else {
             0
         };
@@ -392,3 +394,10 @@ fn fft_32_dit_with_planner_and_opts_impl<S: Simd>(
         }
     }
 }
+
+#[inline]
+fn power_of_two(power: usize) -> usize {
+    // 2.pow() requires a lot of ugly type annotations so here's a helper function
+    debug_assert!(power < usize::BITS as usize);
+    1 << power
+}