pub const FFT_GPU_WGSL: &str = "// RLX \u{2014} versatile ML compiler + runtime.\n// Copyright (C) 2026 Eugene Hauptmann, Nataliya Kosmyna.\n//\n// This program is free software: you can redistribute it and/or modify\n// it under the terms of the GNU General Public License as published by\n// the Free Software Foundation, version 3.\n//\n// This program is distributed in the hope that it will be useful,\n// but WITHOUT ANY WARRANTY; without even the implied warranty of\n// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the\n// GNU General Public License for more details.\n//\n// You should have received a copy of the GNU General Public License\n// along with this program. If not, see <https://www.gnu.org/licenses/>.\n\n// Multi-kernel f32 FFT (gpu-fft strategy), RLX 2N real-block layout.\n\nstruct Params {\n off: u32,\n dst_off: u32,\n n: u32,\n log2n: u32,\n inverse: u32,\n norm_scale: f32,\n outer: u32,\n tile: u32,\n inner_stages: u32,\n q_or_hs: u32,\n};\n\n@group(0) @binding(0) var<storage, read_write> arena: array<f32>;\n@group(0) @binding(1) var<uniform> params: Params;\n\nfn re_at(base: u32, k: u32, n: u32) -> f32 {\n return arena[base + k];\n}\nfn im_at(base: u32, k: u32, n: u32) -> f32 {\n return arena[base + n + k];\n}\nfn set_re(base: u32, k: u32, n: u32, v: f32) {\n arena[base + k] = v;\n}\nfn set_im(base: u32, k: u32, n: u32, v: f32) {\n arena[base + n + k] = v;\n}\n\n// Single-kernel path (N <= 1024): bit-reverse load + all stages in TG mem.\nvar<workgroup> sre_full: array<f32, 1024>;\nvar<workgroup> sim_full: array<f32, 1024>;\n\n@compute @workgroup_size(256)\nfn fft_radix2_full(\n @builtin(workgroup_id) wgid: vec3<u32>,\n @builtin(local_invocation_id) lid: vec3<u32>,\n) {\n let n = params.n;\n let log2n = params.log2n;\n let row = wgid.y;\n if (row >= params.outer) { return; }\n let src_base = params.off + row * 2u * n;\n let dst_base = params.dst_off + row * 2u * n;\n let tid = lid.x;\n let tg_size = 256u;\n\n var k: u32 = tid;\n loop {\n if (k >= n) { break; }\n let rev = reverseBits(k) >> (32u - log2n);\n sre_full[rev] = re_at(src_base, k, n);\n sim_full[rev] = im_at(src_base, k, n);\n k = k + tg_size;\n }\n workgroupBarrier();\n\n let sign = select(-1.0, 1.0, params.inverse != 0u);\n let two_pi = 6.28318530717958647692;\n var len: u32 = 2u;\n loop {\n if (len > n) { break; }\n let h2 = len >> 1u;\n let theta_base = sign * two_pi / f32(len);\n var b: u32 = tid;\n loop {\n if (b >= n / 2u) { break; }\n let group = b / h2;\n let k_in = b % h2;\n let i_lo = group * len + k_in;\n let i_hi = i_lo + h2;\n let theta = theta_base * f32(k_in);\n let wre = cos(theta);\n let wim = sin(theta);\n let t_re = wre * sre_full[i_hi] - wim * sim_full[i_hi];\n let t_im = wre * sim_full[i_hi] + wim * sre_full[i_hi];\n let u_re = sre_full[i_lo];\n let u_im = sim_full[i_lo];\n sre_full[i_lo] = u_re + t_re;\n sim_full[i_lo] = u_im + t_im;\n sre_full[i_hi] = u_re - t_re;\n sim_full[i_hi] = u_im - t_im;\n b = b + tg_size;\n }\n workgroupBarrier();\n len = len << 1u;\n }\n\n k = tid;\n loop {\n if (k >= n) { break; }\n set_re(dst_base, k, n, sre_full[k] * params.norm_scale);\n set_im(dst_base, k, n, sim_full[k] * params.norm_scale);\n k = k + tg_size;\n }\n}\n\n// Bit-reverse one row before multi-kernel outer stages.\n@compute @workgroup_size(256)\nfn fft_bit_reverse(\n @builtin(global_invocation_id) gid: vec3<u32>,\n @builtin(workgroup_id) wgid: vec3<u32>,\n) {\n let row = wgid.y;\n if (row >= params.outer) { return; }\n let n = params.n;\n let k = gid.x;\n if (k >= n) { return; }\n let base = params.off + row * 2u * n;\n let rev = reverseBits(k) >> (32u - params.log2n);\n if (k >= rev) { return; }\n let tr = re_at(base, k, n);\n let ti = im_at(base, k, n);\n set_re(base, k, n, re_at(base, rev, n));\n set_im(base, k, n, im_at(base, rev, n));\n set_re(base, rev, n, tr);\n set_im(base, rev, n, ti);\n}\n\n// Inner shared-memory tile (tile <= 1024).\nvar<workgroup> sre_in: array<f32, 1024>;\nvar<workgroup> sim_in: array<f32, 1024>;\n\n@compute @workgroup_size(512)\nfn fft_inner(\n @builtin(workgroup_id) wgid: vec3<u32>,\n @builtin(local_invocation_id) lid: vec3<u32>,\n) {\n let row = wgid.y;\n if (row >= params.outer) { return; }\n let n = params.n;\n let tile = params.tile;\n let half_tile = tile / 2u;\n let tile_id = wgid.x;\n let num_tiles = (n + tile - 1u) / tile;\n if (tile_id >= num_tiles) { return; }\n let local = lid.x;\n if (local >= half_tile) { return; }\n\n let row_base = params.off + row * 2u * n;\n let tile_base = tile_id * tile;\n\n if (tile_base + local < n) {\n sre_in[local] = re_at(row_base, tile_base + local, n);\n sim_in[local] = im_at(row_base, tile_base + local, n);\n }\n if (tile_base + local + half_tile < n) {\n sre_in[local + half_tile] = re_at(row_base, tile_base + local + half_tile, n);\n sim_in[local + half_tile] = im_at(row_base, tile_base + local + half_tile, n);\n }\n workgroupBarrier();\n\n let sign = select(-1.0, 1.0, params.inverse != 0u);\n let pi = 3.14159265358979323846;\n for (var s: u32 = 0u; s < params.inner_stages; s = s + 1u) {\n let hs = 1u << s;\n let k = local % hs;\n let i = (local / hs) * (hs * 2u) + k;\n let j = i + hs;\n let angle = sign * pi * f32(k) / f32(hs);\n let cos_a = cos(angle);\n let sin_a = sin(angle);\n let ur = sre_in[i];\n let ui = sim_in[i];\n let vr = cos_a * sre_in[j] - sin_a * sim_in[j];\n let vi = sin_a * sre_in[j] + cos_a * sim_in[j];\n sre_in[i] = ur + vr;\n sim_in[i] = ui + vi;\n sre_in[j] = ur - vr;\n sim_in[j] = ui - vi;\n workgroupBarrier();\n }\n\n let scale = params.norm_scale;\n if (tile_base + local < n) {\n set_re(row_base, tile_base + local, n, sre_in[local] * scale);\n set_im(row_base, tile_base + local, n, sim_in[local] * scale);\n }\n if (tile_base + local + half_tile < n) {\n set_re(row_base, tile_base + local + half_tile, n, sre_in[local + half_tile] * scale);\n set_im(row_base, tile_base + local + half_tile, n, sim_in[local + half_tile] * scale);\n }\n}\n\n@compute @workgroup_size(256)\nfn fft_outer_r4(\n @builtin(global_invocation_id) gid: vec3<u32>,\n @builtin(workgroup_id) wgid: vec3<u32>,\n) {\n let row = wgid.y;\n if (row >= params.outer) { return; }\n let n = params.n;\n let q = params.q_or_hs;\n let tid = gid.x;\n if (tid >= n / 4u) { return; }\n\n let base = params.off + row * 2u * n;\n let k = tid % q;\n let group = tid / q;\n let p = group * (q * 4u) + k;\n\n let ar = re_at(base, p, n);\n let ai = im_at(base, p, n);\n let br = re_at(base, p + q, n);\n let bi = im_at(base, p + q, n);\n let cr = re_at(base, p + q * 2u, n);\n let ci = im_at(base, p + q * 2u, n);\n let dr = re_at(base, p + q * 3u, n);\n let di = im_at(base, p + q * 3u, n);\n\n let sign = select(-1.0, 1.0, params.inverse != 0u);\n let neg_sign = select(1.0, -1.0, params.inverse != 0u);\n let angle1 = sign * 3.14159265358979323846 * f32(k) / f32(q);\n let cos1 = cos(angle1);\n let sin1 = sin(angle1);\n let w1b_r = cos1 * br - sin1 * bi;\n let w1b_i = sin1 * br + cos1 * bi;\n let w1d_r = cos1 * dr - sin1 * di;\n let w1d_i = sin1 * dr + cos1 * di;\n\n let u0r = ar + w1b_r;\n let u0i = ai + w1b_i;\n let u1r = ar - w1b_r;\n let u1i = ai - w1b_i;\n let u2r = cr + w1d_r;\n let u2i = ci + w1d_i;\n let u3r = cr - w1d_r;\n let u3i = ci - w1d_i;\n\n let angle2a = sign * 3.14159265358979323846 * f32(k) / f32(q * 2u);\n let cos2a = cos(angle2a);\n let sin2a = sin(angle2a);\n let cos2b = neg_sign * sin2a;\n let sin2b = sign * cos2a;\n\n let w2a_u2r = cos2a * u2r - sin2a * u2i;\n let w2a_u2i = sin2a * u2r + cos2a * u2i;\n let w2b_u3r = cos2b * u3r - sin2b * u3i;\n let w2b_u3i = sin2b * u3r + cos2b * u3i;\n\n let scale = params.norm_scale;\n set_re(base, p, n, (u0r + w2a_u2r) * scale);\n set_im(base, p, n, (u0i + w2a_u2i) * scale);\n set_re(base, p + q * 2u, n, (u0r - w2a_u2r) * scale);\n set_im(base, p + q * 2u, n, (u0i - w2a_u2i) * scale);\n set_re(base, p + q, n, (u1r + w2b_u3r) * scale);\n set_im(base, p + q, n, (u1i + w2b_u3i) * scale);\n set_re(base, p + q * 3u, n, (u1r - w2b_u3r) * scale);\n set_im(base, p + q * 3u, n, (u1i - w2b_u3i) * scale);\n}\n\n@compute @workgroup_size(256)\nfn fft_outer_r2(\n @builtin(global_invocation_id) gid: vec3<u32>,\n @builtin(workgroup_id) wgid: vec3<u32>,\n) {\n let row = wgid.y;\n if (row >= params.outer) { return; }\n let n = params.n;\n let half_stride = params.q_or_hs;\n let tid = gid.x;\n if (tid >= n / 2u) { return; }\n\n let base = params.off + row * 2u * n;\n let k = tid % half_stride;\n let i = (tid / half_stride) * (half_stride * 2u) + k;\n let j = i + half_stride;\n\n let sign = select(-1.0, 1.0, params.inverse != 0u);\n let angle = sign * 3.14159265358979323846 * f32(k) / f32(half_stride);\n let cos_a = cos(angle);\n let sin_a = sin(angle);\n\n let ur = re_at(base, i, n);\n let ui = im_at(base, i, n);\n let vr = cos_a * re_at(base, j, n) - sin_a * im_at(base, j, n);\n let vi = sin_a * re_at(base, j, n) + cos_a * im_at(base, j, n);\n let scale = params.norm_scale;\n set_re(base, i, n, (ur + vr) * scale);\n set_im(base, i, n, (ui + vi) * scale);\n set_re(base, j, n, (ur - vr) * scale);\n set_im(base, j, n, (ui - vi) * scale);\n}\n";