a pixel shader FFT

5 Nov 2026

It’s Thursday, and the deadline for ShaderFes is Saturday. I decided last minute that I want to submit my FFT water. To do that, I need an FFT that runs in a pixel shader. Let’s do it.

Deriving a 1-dimensional fft

Our pixel shader will conceptually look like this:

uniform uint _Stage;
uniform uint _Big_N;
uniform uint _Stride;
uniform Texture2D _Signal;

float4 px(v2f i, uint2 px_pos : VPOS) {
    uint2 res = fft_stage(_Stage, _Signal, px_pos.x, px_pos.y);
    return float4(res, 0, 1);
}

So all we have to do is define our fft_stage based on the output positions.

Our old iterative Rust code is perhaps a useful starting point:

// Get odd and even elements.
let p = input[start_idx + 2 * k * stride];
let q = input[start_idx + (2 * k + 1) * stride];
// Combine.
let factor = twiddles[k * stride];
output[start_idx + k * stride] = p + q * factor;
output[start_idx + (k + big_n / 2) * stride] = p - q * factor;

So from a given x, y, big_n, and stride, we have to derive start_idx and k.

Recall how start_idx and k are set in the iterative code:

for start_idx in 0..stride {
for k in 0..big_n / N {

Start by considering the 1D case, so no y. Since we want to output at x, we know that one of these equations is true:

start_idx + k * stride               = x
start_idx + (k + big_n / 2) * stride = x

How do we know which one is true? In other words, which output is the current input going to produce?

Consider the 0th stage, where stride = input.len() / 2, and big_n = 2.

start_idx + k * stride       = x
start_idx + (k + 1) * stride = x

We know in the 0th stage that k can only be zero, so:

start_idx          = x
start_idx + stride = x

In either case, we can clearly see that start_idx = x % stride.

Generalizing to any k:

start_idx + k * stride               = x
start_idx + (k + big_n / 2) * stride = x

The same identity clearly still applies for start_idx.

If we subtract start_idx from both sides and divide by stride:

k               = (x - start_idx) / stride
(k + big_n / 2) = (x - start_idx) / stride

Applying the same modulus trick:

k = ((x - start_idx) / stride) % (big_n / 2)

Plugging in our expression for start_idx:

k = ((x - (x % stride)) / stride) % (big_n / 2)

Because we’re dividing by an integer, subtracting the remainder does nothing:

k = (x / stride) % (big_n / 2)

Now we have closed-form expressions for start_idx and k:

start_idx = x % stride
k         = (x / stride) % (big_n / 2)

Circling back to our output expressions:

output[start_idx + k * stride] = p + q * factor;
output[start_idx + (k + big_n / 2) * stride] = p - q * factor;

We know from the expression for k that either:

x / stride = k
x / stride = k + big_n / 2

So we can just use this to determine which output we’re producing:

if (x / stride == k) {
    output[x] = p + q * factor;
} else {
    output[x] = p + q * factor;
}

Putting it all together:

uniform uint _Stage;
uniform uint _Big_N;
uniform uint _Stride;
uniform Texture2D _Signal;
// Fill this in for your input size, or plumb it as a uniform.
static const float2 kTwiddles[...] = ...;

float2 cmul(float2 x, float2 y) {
    return float2(x[0] * y[0] - x[1] * y[1], x[0] * y[1] + x[1] * y[0]);
}

float2 fft_stage(uint x) {
    uint start_idx = x % _Stride;
    uint k = (x / _Stride) % (_Big_N / 2);
    uint2 p_coords = uint2(start_idx + 2 * k * _Stride, 0);
    uint2 q_coords = uint2(start_idx + (2 * k + 1) * _Stride, 0);
    float2 p = _Signal.Load(uint3(p_coords, 0));
    float2 q = _Signal.Load(uint3(q_coords, 0));
    // Combine.
    float2 factor = kTwiddles[k * _Stride];
    if (x / _Stride == k) {
        return p + cmul(q, factor);
    } else {
        return p - cmul(q, factor);
    }
}

float4 px(v2f i, uint2 px_pos : VPOS) {
    float2 res = fft_stage(px_pos.x);
    return float4(res, 0, 1);
}

Going 2D

A 2-dimensional FFT is just a horizontal FFT followed by a vertical one. So we just need to add a y parameter to our fft_stage along with a direction.

#define DIR_HORIZONTAL 0
#define DIR_VERTICAL   1

uniform uint _Stage;
uniform uint _Big_N;
uniform uint _Stride;
uniform uint _Direction;
uniform Texture2D _Signal;
// Fill this in for your input size, or plumb it as a uniform.
static const float2 kTwiddles[...] = ...;

uint compute_start_idx(uint x) {
    return x % _Stride;
}

uint compute_k(uint x) {
    return (x / _Stride) % (_Big_N / 2);
}

float2 cmul(float2 x, float2 y) {
    return float2(x[0] * y[0] - x[1] * y[1], x[0] * y[1] + x[1] * y[0]);
}

float2 fft_stage(uint x, uint y) {
    uint main_coord = _Direction == DIR_HORIZONTAL ? x : y;
    uint alt_coord  = _Direction == DIR_HORIZONTAL ? y : x;
    uint start_idx = compute_start_idx(main_coord);
    uint k = compute_k(main_coord);
    uint2 p_coords = uint2(start_idx + 2 * k * _Stride, alt_coord);
    uint2 q_coords = uint2(start_idx + (2 * k + 1) * _Stride, alt_coord);
    if (_Direction == DIR_VERTICAL) {
        p_coords = p_coords.yx;
        q_coords = q_coords.yx;
    }
    float2 p = _Signal.Load(uint3(p_coords, 0));
    float2 q = _Signal.Load(uint3(q_coords, 0));
    // Combine.
    float2 factor = kTwiddles[k * _Stride];
    if (main_coord / _Stride == k) {
        return p + cmul(q, factor);
    } else {
        return p - cmul(q, factor);
    }
}

float4 px(v2f i, uint2 px_pos : VPOS) {
    float2 res = fft_stage(px_pos.x, px_pos.y);
    return float4(res, 0, 1);
}

back to main page