FFT

The fast Fourier transform (FFT) computes the discrete Fourier transform (DFT) of $N$ complex values $x_0, \dots, x_{N-1}$:

$$X_k = \sum_{j=0}^{N-1} x_j \, \omega_N^{jk}, \qquad \omega_N = e^{-2\pi i/N}, \qquad k = 0, \dots, N-1,$$

in $O(N \log N)$ operations.

This example computes one 1D FFT of length $N = P_X P_Y M$ on a $P_X \times P_Y$ grid of PEs, with $M$ complex fp16 values on each PE. It is a variation of the 4-step FFT: local FFTs are computed on each PE, separated by all-to-all exchanges between PEs.

The algorithm

The 4-step FFT

Split the length as $N = N_0 N_1$ and write the indices as $j = j_0 N_1 + j_1$ and $k = k_0 + N_0 k_1$. The DFT becomes

$$X_{k_0 + N_0 k_1} = \sum_{j_1=0}^{N_1-1} \omega_{N_1}^{j_1 k_1} \; \underbrace{\omega_N^{j_1 k_0}}_{\text{twiddle}} \; \underbrace{\sum_{j_0=0}^{N_0-1} x_{j_0 N_1 + j_1} \, \omega_{N_0}^{j_0 k_0}}_{\text{FFT of length } N_0}.$$

Equivalently, this becomes $N_1$ FFTs of length $N_0$, a pointwise multiply by twiddle factors, then $N_0$ FFTs of length $N_1$. If each small FFT runs on a single PE, the only communication is a reshuffle between the two steps: an all-to-all.

Rounds

However, if the number of elements per PE ($M$) is small and the total number of elements ($N$) is large, two steps are not enough. We can apply this same split repeatedly, giving $N = \prod_{t=0}^{R-1} N_t = N_0 N_1 \cdots N_{R-1}$, with every $N_t \leq M$ fitting into local PE memory. The FFT then runs

$$R = \left\lceil \frac{\log_2 N}{\log_2 M} \right\rceil$$

total rounds, each with three phases:

1. Exchange (not in round 0): an all-to-all among a small group of PEs,
2. Local FFT: $M/N_t$ interleaved FFTs of length $N_t$ on each PE,
3. Twiddle (not in the last round): local pointwise multiplication by a table of twiddle factors.

Conceptually, the bits of global indices $j$ are broken up to determine each data point's PE row, PE column, and local slot within that PE. Between rounds, a few column (or row) bits are interchanged with slot bits; each exchange therefore stays within a single row (or column) of PEs.

For example, a $4 \times 4$ PE grid with $M=4$ computes an FFT of length $N = 64$ in three rounds of length-4 FFTs. Every exchange here is between pairs of PEs: each PE swaps half of its values with one partner.

The diagram below follows the bits of the input index $j$ and the output index $k$ through the rounds.

                      PE row (py)   PE column (px)   slot (s)
                      1   0         1   0            1   0
round 0   input       j1  j3        j0  j2           j5  j4
          FFT         j1  j3        j0  j2           k1  k0
round 1   exchange    j1  k1        j0  k0           j3  j2
          FFT         j1  k1        j0  k0           k3  k2
round 2   exchange    k3  k1        k2  k0           j1  j0
          FFT         k3  k1        k2  k0           k5  k4

The twiddle multiplies after rounds 0 and 1 change values but move no bits. Round 1 exchanges bit 0 of $p_x$ and $p_y$, so its partners are neighbouring PEs. Round 2 exchanges bit 1, so its partners are 2 apart. A PE keeps the slots whose exchanged slot bit equals its own PE bit, and swaps the others, the shaded slots in the figure.

Larger rounds exchange more bits at once, among larger groups. A $16 \times 16$ grid with $M=64$ runs rounds of $2^6$, $2^4$ and $2^4$ points. Round 1 exchanges $p_x$ and $p_y$ bits 1–0, among groups of 4 neighbouring PEs. Round 2 exchanges bits 3–2, among groups of 4 PEs spaced 4 apart. python plan.py PX PY M prints the plan for any grid.

The exchanges and rounds are precomputed on the host by plan.py: each PE gets a table plan_table[t] of each round's FFT size and exchanged bits, and its twiddle tables.

The local FFT

Each round's FFT runs on one PE, over its $M$ slots. When $N_t < M$, the PE holds $B = M/N_t$ independent FFTs, interleaved: batch $b$ is the values in slots $b, b+B, b+2B, \dots$. So the round transforms the top $\log_2 N_t$ slot bits and leaves the low $\log_2 B$ bits alone. In the $16 \times 16$, $M=64$ plan above, rounds 1 and 2 each hold $B = 4$ FFTs of length 16.

batched_fft() is an in-place, self-sorting radix-2 FFT: input and output are both in natural order, so no bit-reversal pass is needed. The buffer x is split into halves x_lo and x_hi. Each stage multiplies x_hi by twiddles, does the butterflies $a \pm wb$ across the two halves, and re-shuffles so the next stage again finds its pairs one half apart. A full $M$-point FFT takes $\log_2 M$ stages. Stopping after $\log_2 N_t$ stages leaves exactly the $B$ interleaved FFTs of length $N_t$.

Communication

Two networks, each on its own color:

• Message passing carries the exchanges. Every message has a header with its destination PE, so one bank of message-passing routes serves every exchange, whatever its stride. A message holds a tag (the sender, and which part of its data this is) and up to 16 values; the receiver uses the tag to place them.
• A switched barrier starts every exchange. Without it, a fast PE could send data for its next exchange to a PE still receiving the current one. The barrier uses router switch positions and control wavelets on a static color (see the barrier-broadcast example). Every PE but the origin starts closed, and opens its own switch when it arrives. A probe from the origin can only get back round the grid once every PE is open. The origin then sends a release to everyone and resets the switches.

Files

File	Contents
fft.w	The program: parameters, buffers, the local FFT and twiddle multiply, and main, which runs the rounds
exchange.w	The all-to-all exchange by message passing, and the barrier
bench.paint	The grid, its message-passing bank, and the barrier network
plan.py	The plan and host tables, with a numpy model of the kernel (python plan.py self-tests it)
ref.py	Writes the inputs, and checks the output against numpy
Makefile, test.rc	Build, simulate, and check

Reading the code

Start at main at the bottom of fft.w. It loops over the rounds, reads round $t$'s row of plan_table, and calls one function per phase:

Phase	Function	File
Exchange along the row, then the column	all_to_all(along_row, ...)	exchange.w
Barrier at the start of each exchange	barrier()	exchange.w
Local FFT	batched_fft(log_N_t)	fft.w
Twiddle	twiddle_multiply(t)	fft.w

all_to_all() runs in three steps:

1. Build every outgoing message (header, tag, values) into one envelope buffer.
2. Call barrier(), then send the envelope and receive the other PEs' messages into receive_buffer, at the same time, in a parallel block.
3. Copy the PE's own values (the ones that stay put) into the buffer, then the buffer back to x.

On the host side, plan.py builds the plan, and its docstring gives the index algebra in full. The barrier network's switch positions are in makebarriergrid() in bench.paint.

Lessons from the exchange

Two details of all_to_all() matter. Without them, large configurations deadlock:

• Send everything as one stream. A loop that sends message by message, each from its own send statement, keeps the CE from receiving in between. With 16 or more PEs in a group, the fabric's buffers fill and every PE ends up waiting on a send. So the messages are packed into envelope first, and one statement, ∀i ∈ [0, envelope_words) mp[] ← envelope[i], streams the whole envelope.
• Receive with an iterator loop. The receive arm of the parallel block is a ∀ loop. Written as a while loop, it does not overlap with the send arm, and the exchange deadlocks.

Building and running

The experimental SDK installs this example beside the tutorial, as projects/fft. Run every row of its test.rc with tungsten-tutorials -f projects/fft, or work by hand in a copy of the directory, set up as for the tutorials, where <sdk> is the SDK's directory:

cp -r <sdk>/share/tutorial/projects/fft . && cd fft
python plan.py 16 16 64     # show the plan for one configuration
make TEST=7                 # build + simulate + check row 7 of test.rc
make TEST=grid              # the three grid rows
make -k TEST='*'            # every row
make clean                  # remove all test-*/ directories

$(SIMFABRIC) runs the simulation on the Cerebras SDK simulator. MODE (default sim) only names the run's outputs (out-$(MODE).bin); sim is the only mode the Makefile provides.

Each row of test.rc is one configuration, tagged by size (small, grid, deep, large):

grid PX=16 PY=16 M=64 ARCH=sdr

Variable	Meaning
PX, PY	Grid width and height, powers of 2, at least 2
M	Complex values per PE, a power of 2 from 2 to 1024
MPBANK	Message-passing bank, 0 or 1 (default 0)
CYCLES	Simulation cycle cap (default 50000)
ARCH	sdr (WSE-3) in every row

The build compiles fft.w for the grid and writes the inputs with ref.py. It loads them into the image with miset, simulates it, reads x back with miget and checks it. Each tag appears on more than one row, so test.mk names each run directory by row number: test-1/ to test-12/. Each holds make-log (the real output; make's own stdout is a one-line PASS/FAIL).

The check compares the result with np.fft.fft of the same input, rounded to fp16. It passes when the largest error is within 2% of the largest output magnitude, $\max_k |X_k - X^{\text{ref}}_k| \le 0.02 \, \max_k |X^{\text{ref}}_k|$. The kernel computes in fp16, so the tolerance is loose. A passing run ends with

PASS: max error <e> (relative to max |X|)

and a failing one exits non-zero with FAIL: max error <e> (relative to max |X|) > 0.02. If the run doesn't finish within CYCLES, x is read back mid-computation and the check fails.

Tungsten and Paint concepts used

Tungsten (language reference):

• Parameters (xp param M), set per configuration from Paint
• Imports, to split the program into files, and a main block
• Iterator loops, the vector operations of the local FFT and twiddle multiply, written with the Unicode syntax (∀i ∈ [0, M), ← for sends and receives, ∧)
• Arrays at fixed addresses, both to alias buffers (x as fp16 words and as 32-bit wavelets) and to read the PE's grid position
• Unions, to pack 16-bit fields into wavelets
• Sockets and parallel blocks, to send and receive at the same time
• Message passing and control wavelets

Paint (language reference):

• Code attachment with parameters, and naming for miset/miget
• Message passing: bank(), color(bank) and mp()
• Colors, switch positions, functions and layout combinators (hstack, vstackrep, ...) for the barrier network

Related examples: the message-passing tutorials (tutorial/message-passing/) and bench/app/barrier-broadcast.

Run it

The experimental SDK installs this project beside the tutorial, as projects/fft. Run it like a tutorial: tungsten-tutorials -f projects/fft.