HPCG on the WSE-3

This directory implements the HPCG benchmark for the Cerebras Wafer-Scale Engine (WSE-3) in Tungsten (hpcg.w). HPCG solves a sparse linear system from a 27-point stencil with conjugate gradient (CG), preconditioned by a multigrid V-cycle whose smoother is symmetric Gauss–Seidel (SYMGS).

The implementation rests on three ideas:

• it maps the 3-D grid onto the 2-D array of processing elements (PEs), one grid column per PE;
• it keeps all data in the PEs' local memory;
• it runs Gauss–Seidel, which is inherently sequential, as a pipelined wavefront across the PEs.

A C reference (hpcg_ref1.c) runs the same algorithm in the same order, and every test checks the fabric's result against it. The code is fp32 only, so a run is not an official HPCG result (TestSymmetry is scaled by DBL_EPSILON).

Sections 1–5 describe the algorithm and its mapping onto the fabric. Section 6 covers the program structure, and Section 7 how to build and run it.

1. The benchmark

HPCG discretizes a 3-D grid of nx × ny × nz points, with one unknown per point. Row i of the matrix A couples point i with the points of its 3 × 3 × 3 neighbourhood (the 27-point stencil). In the reference problem the diagonal is 26 and the off-diagonals are −1.

CG solves A x = b. Each iteration:

• applies the preconditioner, z = M⁻¹ r;
• does one sparse matrix-vector product (SpMV), A p;
• computes three dot products.

The preconditioner is one V-cycle of a 4-level geometric multigrid (Figure 3, Section 4). Each coarser level halves the grid in x, y and z. On every level except the coarsest, one SYMGS runs before the coarse-grid correction and one after it. The coarsest level gets a single SYMGS.

Three properties of the workload shape any implementation:

• Low arithmetic intensity. SpMV and SYMGS use each matrix coefficient for a single multiply-add per pass. The work is in moving A and x to the arithmetic, not in the arithmetic itself.

• SYMGS is sequential. The forward sweep visits the points in order, and each point uses the values its earlier neighbours received in this sweep:

  x_i ← (r_i − Σ_{j<i} a_ij x_j^new − Σ_{j>i} a_ij x_j^old) / a_ii

  The backward sweep does the same in reverse order. HPCG lets every rank sweep its own subdomain, reading the other ranks' values as they were at the start of the sweep. This implementation does the same with its GS domains (Section 3.2).

• Global reductions. The dot products need reductions over the whole machine, and the next step of CG waits for their results.

2. Mapping the grid onto the fabric

Figure 1. (A) Each (x, y) column of the grid goes to one PE. (B) The stencil neighbours of a column are the columns of the 8 adjacent PEs. (C) A PE holds its column's data for all levels, and every operand of a stencil row is local or one hop away.

The x and y axes of the grid map onto the PE array, and z stays inside the PE. PE (ix, iy) holds the whole column of Z points at (ix, iy), for every multigrid level:

• the matrix coefficients (26 off-diagonal per row, plus the diagonal);
• x and r;
• the CG vectors.

Nothing is stored off the PEs. A point's z neighbours are in the same PE; its x and y neighbours are in the 8 adjacent PEs, one fabric hop away.

With this mapping, SpMV and Gauss–Seidel have the same shape:

1. Every PE streams its column to its 8 neighbours, as a window x[z−1..z+1] of 3 values per row.
2. Every PE receives 8 such streams. For each row z it multiplies the 24 incoming values by 24 coefficients from its memory and adds the 3 terms of its own column.

The fabric delivers the windows in row order, so a PE consumes them as they arrive and never stores a neighbour's column. The routing is fixed for the duration of a phase: before each phase, every PE loads the router configuration of that phase (Section 6.4).

3. Gauss–Seidel as a wavefront

3.1 Order and depth

The GS order is z fastest, then x, then y. For a point at (ix, iy, z), the neighbours that come earlier in this order are:

• all three z offsets in the columns to the sw, s, se, w;
• z−1 in its own column.

The later ones are e, nw, n, ne, and z+1.

In the forward sweep, PE (i, j) can therefore update row z only after its w, sw, s and se neighbours have updated rows up to z+1. Its start time is set by the longest chain of such dependencies from the corner PE (0, 0). With depth d = 1 + the largest depth among the four predecessors, this gives d(i, j) = i + 2j. The binding predecessor is the se neighbour (i+1, j−1), at depth i + 2j − 1.

So PEs on a line i + 2j = const work at the same time, each on its own row. A PE trails its predecessors by a few rows and streams its new rows to its successors as it computes them. The sweep is a pipelined wavefront: a plane z = t − d(i, j) that moves through the 3-D grid (Figure 2).

Figure 2. One SYMGS on a domain of 6 × 6 PEs with 24 rows per PE. The animation is schematic: one row per step. The forward sweep starts at PE (0, 0), and PE (i, j) starts i + 2j steps later. Once the deepest PE has started, every PE is busy until the first PEs finish. The backward sweep runs the same wave from the far corner back to PE (0, 0). The plot at the bottom shows the fraction of busy PEs: ramp-up, full utilization, ramp-down.

An N × N domain has depth 3(N − 1). The fill and drain of every sweep therefore grow with the domain edge, while the fully busy part grows with Z.

3.2 GS domains and z-ranks

The PE plane is split into GS domains of NGS × NGS PEs (16 × 16 by default), which sweep concurrently. A wavefront never crosses a domain edge. As across HPCG ranks, values from other domains are frozen at the start of the SYMGS (BC=frozen). The domain edge cannot be smaller than 16, because HPCG requires local subdomains of at least 16 points in every dimension.

SYMGS is affine in (r, x), so

SYMGS_frozen(r, x) = x + SYMGS_severed(r − A x, 0)

A GS that does not start from x = 0 therefore runs as a SpMV on the full stencil, followed by a GS from zero on the routing with the domain edges cut.

The z column is cut the same way, into z-ranks of ZR points. zf_patch zeroes the coefficients that couple across a z-rank boundary before a GS, and restores them before a SpMV. The declared HPCG rank is therefore NGS × NGS PEs × ZR points.

3.3 One SYMGS in two routed phases

Every SYMGS of the V-cycle starts from x = 0: pre-smoothing by construction, post-smoothing through the identity above. From x = 0 a SYMGS needs two routed phases:

phase	streams read	order	result
gs_fwd_down_fused	sw, s, se, w, own z−1 (new x)	wavefront from PE (0, 0), z ↑	new x[z]; ax[z] = lower sum
gs_bwd_up	ne, n, nw, e, own z+1 (new x)	wavefront from the far corner, z ↓	new x[z] = (r − ax − upper sum) / diag

Two facts remove the other phases:

• In the forward sweep, the upper neighbours still hold 0, so their sum vanishes.
• The lower sum of the forward sweep is also the lower sum of the backward sweep. Those values do not change in between: the lower neighbours keep their forward values until the backward wave reaches them. So the forward phase stores the lower sum in ax, and the backward phase reads it from there.

4. Multigrid on a fixed PE array

Figure 3. One HPCG iteration, as in the HPCG reference (CG_ref and ComputeMG_ref). Each level halves the grid in x, y and z. Restriction and prolongation are injections.

Level L has (nx × ny × nz) / 8^L points, but the PE array does not shrink. At level L the points sit on every 2^L-th PE in x and y (the owners), each with a column of Z / 2^L points. The PEs in between hold no level-L points and relay the streams between owners.

A coarse point is also a point of the finer level and has the same owner. Restriction and prolongation are therefore local to the owner's column: every second z point.

Left like this, level 1 would keep one PE in four busy, and level 2 one in sixteen. Co-sets put all PEs to work (L1_REDIST, COSET23):

• At level 1 the four PEs of each 2 × 2 block split the owner's column into four z-chunks of Z/8 points. Chunk k moves to PE k of the block (owner, E, N, NE).
• Chunk k on all blocks forms an interleaved copy of the level-1 grid, a co-set. All four co-sets sweep at once (Figure 4).
• The chunk boundaries fall on z-rank boundaries (hence Z ≥ 4·ZR), so a GS treats them like any other z-rank boundary.
• Levels 2 and 3 stay in the co-sets: co-set k keeps chunk k down to the coarsest level, so the restriction from level 1 to 2 and the matching prolongation are local. Before each SpMV at levels 1 and 2, an exchange phase (xchg) copies the x values at the chunk boundaries into the neighbouring chunks' pads.
• The data move into the co-sets after restriction to level 1 (redist) and back to the owners before prolongation to level 0 (reduce).

Figure 4. Co-sets at level 1. The animation is schematic: plain row sweeps, without the wavefront fill. Without co-sets, one PE in four updates all Z/2 rows of its column. With co-sets, every PE updates Z/8 rows.

5. Conjugate gradient and the dot product

One CG iteration, after the V-cycle z = MG(r):

rtz = <r, z>                    allreduce
p   = z + (rtz / rtz_old) p
Ap  = A p                       SpMV at level 0
pAp = <p, Ap>,  rr = <r, r>     one allreduce of two values
x  += (rtz / pAp) p
r  -= (rtz / pAp) Ap

dot() is an allreduce built from short chains, so no PE adds more than a few values:

1. Neighbouring PEs pair up (color 18).
2. Row chains carry the sums toward the column x = H/2 (colors 16, 17).
3. A column chain carries them toward the row y = V/2 (colors 20, 21).
4. The PE at the center adds everything and broadcasts the result (color 19).

The cost grows with H + V. Work that does not need the result (the x update, zeroing the coarse x) runs while a PE waits for the broadcast. Every PE receives the reduced values and computes the scalar ratios itself. A Newton iteration (inv) computes the reciprocals.

6. Program structure

6.1 Data layout

The global grid is H × V × Z (H × V PEs). Each PE stores all MG levels packed, finest first:

array	per row	content
A	26	4 upper triples, 4 lower triples, z+1, z−1 coefficients
diag, invd	1	diagonal, 1/diagonal
r, ax	1	right-hand side / residual, partial A·x
x	1	iterate, with a zero pad at each end (Dirichlet)

All kernels address the arrays through iterators. set_level(L) patches their bases from a per-level table (lvl_tab, written by the host), so each kernel is written once and runs at every level, including the co-set view of level 1.

6.2 Kernels

The row body of gs_fwd_down_fused:

ra   <- walkr_fwd[] - walkax_rd_fwd[];            // r - upper sum
rinv <- ra * walkinvd_fwd[];
acc  <-  ∑ ∀i3 (sp)(sw[]) *walkA3_down[];         // 3 z offsets per neighbour
acc  <+- ∑ ∀i3 (sp)(s[])  *walkA3_down[];
acc  <+- ∑ ∀i1 walkZ1b_fwd[] *walkA1b_down[];     // own column, z-1
acc  <+- ∑ ∀i3 (sp)(se[]) *walkA3_down[];
acc  <+- ∑ ∀i3 (sp)(w[])  *walkA3_down[];
walky_fwd[]     = rinv - walkinvd2_fwd[] * acc;   // x[z] = (r - ax - acc) / diag
walkax_wr_fwd[] = acc;                            // lower sum for gs_bwd_up
∀i3 wv:tx[] ← tx_window_fd[];                     // send the x window

walk*[] reads or writes the next element of an iterator, sw[], s[], … read the incoming neighbour streams, and ∑ ∀i3 a[] * b[] is a 3-term dot product accumulated in a dot-product register.

Cost of a row (one fp32 multiply-add per cycle):

• the five dot steps are 4·3 + 1 = 13 multiply-adds, 13 cycles;
• the r − ax and ·invd steps, the final multiply-subtract, two stores and the 3-word send take about 9 more, so a row issues in about 22 cycles (12 instructions);
• every row uses the same dot-product register, so a row's chain starts only after the previous row's has drained. Unrolling the loop by two, with the rows alternating between two registers, overlaps the chains; it makes a CG iteration 4-7% faster (16x16 to 32x32 PEs) and is left out here for clarity;
• the terms that arrive last (se, w: neighbours updated in this sweep) come last in the chain, and the own z−1 term (the previous row's result) comes after sw and s.

SpMV uses the same row shape but is not a wavefront: every PE reads its neighbours' x from before the phase, so all PEs stream at once.

• At level 0, SpMV is one phase that reads all 8 neighbours (spmv_full).
• At the coarse levels it is two phases: gs_fwd_up computes the upper sum, and spmv_down adds the lower sum.

gs_fwd_up (13 multiply-adds, a store and a send) issues in about 18 cycles per row.

6.3 Op program

Tungsten inlines every call, so each kernel has a single call site in run_iter. run_iter interprets vprog, a table of 16-bit ops for one CG iteration, generated by vprog.py and loaded into the image. The ops are:

• a routed phase, given its kernel and routing context;
• set_level;
• restrict and prolong;
• the co-set redist, reduce and xchg;
• the z-rank patch;
• the CG scalar steps.

A routed phase loads its routing context (set_router) and the queue colors (set_sockets), then runs the kernel. It ends with a teardown control wavelet. The V-cycle:

GS(0) SpMV(0) restrict   GS(1) SpMV(1) restrict   GS(2) SpMV(2) restrict
GS(3)
prolong GS(2)   prolong GS(1)   prolong GS(0)

Each post-smoothing GS(L) on the way up runs as SpMV(L) and a GS from zero (Section 3.2). The CG steps of Section 5 follow the V-cycle.

6.4 Routing

hpcg_router1.c writes the router configuration of every PE for each of the C routing contexts (17 by default): the GS and SpMV routings of each level, and the co-set redistribute, reduce and exchange. Before each phase set_router loads the context's slots and set_sockets the queue colors. On the level-L grid a PE sends on color (ix + 3·iy) mod 6; a receiver's slot depends only on its offset from the sender, so one table per direction covers the plane. Coarse levels route through the inactive PEs, and the four co-sets use colors 6k .. 6k+5.

A GS phase reads only 4 of the 8 neighbour streams. The streams do not arrive in step: in the forward sweep the sw neighbour runs two wave steps ahead of the binding se neighbour (Section 3.1), so its rows wait in the input queue (ne in the backward sweep). In the GS phases of levels 0 and 1, the 4 idle input queues lend their storage to the active ones (DBLQ).

7. Running

Build the toolchain, cd into this directory (so the environment loads), then:

make check GOAL=test TEST=smoke MODE=sim          # one test
make -j4 -k check GOAL=test TEST=sim MODE=sim     # all tests

On the cluster, prefix with cbrun -t rocky -- srun -c 16 --mem=60G.

Each test builds and simulates in test-<name>/ (log: make-log), then runs three checks:

1. fcmp_check.py compares x after CG_NITER iterations with the reference (abs/rel tolerance 1e-4).
2. rr_history.py compares ‖r‖² of every iteration (rel 1e-3).
3. perf_report.py prints cycles per iteration and the HPCG flop rate.

make check does not rebuild after source edits: remove test-*/bench.obj first.

test	PEs	Z	GS domains	iterations	sim time
smoke	8×8	128	1 (NGS=8)	2	1 min
small	16×16	128	1	10	12 min
small_z192	16×16	192	1	10	15 min
domains2x2	32×32	192	2×2	10	1.5 h

Files

file	role
hpcg.w	PE program: kernels, V-cycle, CG, dot product, timers
bench.paint	fabric layout: one H × V region running hpcg
vprog.py	op program of one CG iteration
hpcg_router1.c	routing contexts and the dot product configuration
pack_routing.py	packs the router tables for hpcg.w
host_levels.py, host_stencil.py	level and queue tables, written into the image
hpcg_ref1.c	C reference: problem, MG hierarchy, CG, input files, expected results, flop count
fcmp_check.py, rr_history.py, perf_report.py	checks and cycle report
Makefile, test.rc	build flow and test list
hpcg_mapping.svg, hpcg_gs_wavefront.gif, hpcg_vcycle.svg, hpcg_cosets.gif	Figures 1–4

Build: paint and compile hpcg.w (bench.obj); run the router, vprog.py, the reference and the host table scripts, and load their output into the image (bench.img); simulate; read xcg, rrh and the timestamps back. bench.obj does not depend on the router or the reference.

Make variables

variable	default	meaning
H, V	–	PEs in x / y
Z	–	level-0 points per PE in z
CG_NITER	10	CG iterations
NGS	16	GS domain edge in PEs
ZR	16	z-rank height; 0 = one z-rank
LEVELS	4	MG levels
MG	1	1 = V-cycle preconditioner, 0 = plain CG
BC	frozen	coupling across GS domains: frozen (HPCG) or sever (block GS)
TSTAMP	0	1 = per-phase timestamps
XTOL	1e-4	tolerance of the x check
REFCACHE	–	cache directory for the reference outputs

The other variables in the Makefile select the configuration described above and are best left at their defaults.

Constraints: H, V, NGS and Z divisible by 8; H, V divisible by NGS; Z divisible by ZR, Z ≥ 4·ZR. HPCG needs a local subdomain of at least 16 in every dimension (NGS ≥ 16, ZR = 16). Code and data must fit in PE memory, which limits Z to 224.

Reference

hpcg_ref1.c is a single-file C port of the HPCG reference with simulated ranks: problem generation, MG hierarchy, SpMV, SYMGS in the fabric's order (domains, z-ranks, co-set chunks), V-cycle, CG, and the HPCG credited-flop count. It writes the fabric's input data and expected results. With --init random (the default here) the matrix is a perturbed, non-symmetric version of the 26/−1 stencil, so any coupling error shows in the checks.

Run it

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