nx.gpu

GPU-resident mirror of nx.fft* — Stockham radix-2 FFT compute shaders.

Loaded by nx automatically as nx.gpu. Operates on GPU-resident substrate buffer handles — the same typed buffer primitive the CPU nx.* ops use, with kind = "gpu" so a compute dispatch binds them.

Buffer layout

Complex samples are interleaved vec2<f32>, held as f32 elements: a length-N complex spectrum is 2 * N f32s. A compute shader makes one and hands back the handle:

local fft = asset.resolve("@builtin::shaders.compute_fft_stage", "computeShader")
local specIn  = fft:createBuffer("spec_in",  { type = "f32", len = 2 * n, usage = { "readback" } })
local specOut = fft:createBuffer("spec_out", { type = "f32", len = 2 * n, usage = { "readback" } })

2D layout is row-major: element (y, x) lives at index (y * width + x) * 2.

API

nx.gpu.fft1d({ input = bufIn, output = bufOut, size = n })
nx.gpu.ifft1d({ input = bufIn, output = bufOut, size = n })
nx.gpu.fft2d({ input = bufIn, output = bufOut, width = w, height = h })
nx.gpu.ifft2d({ input = bufIn, output = bufOut, width = w, height = h })

input and output are buffer handles; pass the same handle for both to transform in place. Sizes (and width / height) must be powers of two ≥ 2. Non-power-of-two GPU FFT (Bluestein chirp-z) is a follow-up.

Example

Verify a forward+inverse roundtrip via readback:

local n = 16
local fft = asset.resolve("@builtin::shaders.compute_fft_stage", "computeShader")
local exIn  = fft:createBuffer("ex_in",  { type = "f32", len = 2 * n, usage = { "readback" } })
local exOut = fft:createBuffer("ex_out", { type = "f32", len = 2 * n, usage = { "readback" } })

-- 16-sample sine at bin 3 (interleaved re, im)
local samples = {}
for i = 0, n - 1 do
    samples[2 * i + 1] = math.sin(2 * math.pi * 3 * i / n)
    samples[2 * i + 2] = 0
end
exIn:write(samples)

nx.gpu.fft1d({ input = exIn, output = exOut, size = n })
-- After dispatch, read `exOut` back — bins 3 and 13 should peak at mag N/2.

Implementation

One generic Stockham auto-sort single-stage WGSL shader handles 1D, 2D row pass, and 2D col pass via line/element stride parameters. Each log2(N)-stage pass ping-pongs an internal scratch buffer with the user's output buffer; an extra copy is inserted if the parity puts data in the wrong buffer. Inverse passes are scaled by 1/N (1D) or 1/(width*height) (2D) via a final scale dispatch.

See src/lua/lib/shaders/compute_fft_{stage,copy,scale}.computeShader/ for the WGSL sources.