nx_fft_f32.nx source

1// nx_fft_f32.nx -- f32 radix-2 FFT + f32 FNet 2D-DFT token mixer. The bridge that lets the sub-quadratic FNet 2// mixer (FNET-001/002, which ran in i64 Q14 fixed-point) plug into the f32 autograd tape so the block becomes 3// TRAINABLE. Twiddles are synthesized by the sovereign CORDIC (fx_cos/fx_sin, Q16.16) then converted to f32; 4// because f32 twiddles are unit-magnitude actual values, the butterfly needs NO per-stage descaling -- cleaner 5// (and more accurate) than the fixed-point version. DRY: reuses nx_fft_log2 + nx_fft_bit_reverse. 6// license_tier: ORIGINAL 7import "nx_fft.nx" // nx_fft_log2 + nx_fft_bit_reverse (integer, length-generic) 8import "fx.nx" // CORDIC fx_cos/fx_sin + FX_TWO_PI 9import "nx_f32.nx" // f32 add/sub/mul 10import "nx_f32_div.nx" // nx_f32_div 11import "nx_f32_cvt.nx" // nx_i32_to_f32 12import "nx_syscalls.nx" 13const FFTF_MAGIC_65536: i64 = 65536 14 15const FFTF_F32_ZERO: i64 = 0 16 17// Q16.16 (CORDIC) integer -> f32 value (q / 65536) 18func fftf_q16_to_f32(q: i64) -> i64 { return nx_f32_div(nx_i32_to_f32(q), nx_i32_to_f32(FFTF_MAGIC_65536)) } 19 20// length-n twiddle table (n/2 entries) in f32: tw_re[k]=cos(2pi k/n), tw_im[k]=-sin(2pi k/n) 21func fftf_twiddles(n: i64, tw_re: *i64, tw_im: *i64) -> i64 { 22 let half: i64 = n / 2 23 var k: i64 = 0 24 while k < half { 25 let ang: i64 = FX_TWO_PI * k / n 26 tw_re[k] = fftf_q16_to_f32(fx_cos(ang)) 27 tw_im[k] = fftf_q16_to_f32(0 - fx_sin(ang)) 28 k = k + 1 29 } 30 return 0 31} 32 33// in-place f32 radix-2 forward FFT (re/im are parallel f32 arrays of length n; tw built by fftf_twiddles(n,..)) 34func fftf_fwd(re: *i64, im: *i64, n: i64, tw_re: *i64, tw_im: *i64) -> i64 { 35 let log2n: i64 = nx_fft_log2(n) 36 var i: i64 = 0 37 while i < n { 38 let j: i64 = nx_fft_bit_reverse(i, log2n) 39 if i < j { 40 let tr: i64 = re[i]; let ti: i64 = im[i] 41 re[i] = re[j]; im[i] = im[j]; re[j] = tr; im[j] = ti 42 } 43 i = i + 1 44 } 45 var m: i64 = 1 46 while m < n { 47 let m2: i64 = m * 2 48 let twstride: i64 = n / m2 49 var k: i64 = 0 50 while k < n { 51 var j2: i64 = 0 52 while j2 < m { 53 let idx: i64 = j2 * twstride 54 let w_r: i64 = tw_re[idx]; let w_i: i64 = tw_im[idx] 55 let x_r: i64 = re[k + j2 + m]; let x_i: i64 = im[k + j2 + m] 56 let t_r: i64 = nx_f32_sub(nx_f32_mul(w_r, x_r), nx_f32_mul(w_i, x_i)) 57 let t_i: i64 = nx_f32_add(nx_f32_mul(w_r, x_i), nx_f32_mul(w_i, x_r)) 58 let u_r: i64 = re[k + j2]; let u_i: i64 = im[k + j2] 59 re[k + j2] = nx_f32_add(u_r, t_r) 60 im[k + j2] = nx_f32_add(u_i, t_i) 61 re[k + j2 + m] = nx_f32_sub(u_r, t_r) 62 im[k + j2 + m] = nx_f32_sub(u_i, t_i) 63 j2 = j2 + 1 64 } 65 k = k + m2 66 } 67 m = m2 68 } 69 return 0 70} 71 72// FNet 2D-DFT real-part mixer over an [n,d] row-major f32 matrix. Row pass (length d) then column pass (length 73// n); real part stays in x_re, x_im is scratch. Self-adjoint (proven FNET-002), so backprop reuses this fn. 74func fnet_mix_f32(x_re: *i64, x_im: *i64, n: i64, d: i64) -> i64 { 75 let twr_d: *i64 = (sys_mmap(d * 8)) as *i64 76 let twi_d: *i64 = (sys_mmap(d * 8)) as *i64 77 fftf_twiddles(d, twr_d, twi_d) 78 let rr: *i64 = (sys_mmap(d * 8)) as *i64 79 let ri: *i64 = (sys_mmap(d * 8)) as *i64 80 var y: i64 = 0 81 while y < n { 82 var c: i64 = 0 83 while c < d { rr[c] = x_re[y * d + c]; ri[c] = x_im[y * d + c]; c = c + 1 } 84 fftf_fwd(rr, ri, d, twr_d, twi_d) 85 c = 0 86 while c < d { x_re[y * d + c] = rr[c]; x_im[y * d + c] = ri[c]; c = c + 1 } 87 y = y + 1 88 } 89 let twr_n: *i64 = (sys_mmap(n * 8)) as *i64 90 let twi_n: *i64 = (sys_mmap(n * 8)) as *i64 91 fftf_twiddles(n, twr_n, twi_n) 92 let cr: *i64 = (sys_mmap(n * 8)) as *i64 93 let ci: *i64 = (sys_mmap(n * 8)) as *i64 94 var xc: i64 = 0 95 while xc < d { 96 var r2: i64 = 0 97 while r2 < n { cr[r2] = x_re[r2 * d + xc]; ci[r2] = x_im[r2 * d + xc]; r2 = r2 + 1 } 98 fftf_fwd(cr, ci, n, twr_n, twi_n) 99 r2 = 0 100 while r2 < n { x_re[r2 * d + xc] = cr[r2]; x_im[r2 * d + xc] = ci[r2]; r2 = r2 + 1 } 101 xc = xc + 1 102 } 103 return 0 104} 105 106// smoke: FFT of [1,1,1,1] -> bin0 = 4, others ~0 107func main() -> i64 { 108 let re: *i64 = (sys_mmap(4 * 8)) as *i64 109 let im: *i64 = (sys_mmap(4 * 8)) as *i64 110 var i: i64 = 0 111 while i < 4 { re[i] = nx_i32_to_f32(1); im[i] = FFTF_F32_ZERO; i = i + 1 } 112 let twr: *i64 = (sys_mmap(2 * 8)) as *i64 113 let twi: *i64 = (sys_mmap(2 * 8)) as *i64 114 fftf_twiddles(4, twr, twi) 115 fftf_fwd(re, im, 4, twr, twi) 116 return 0 117}