diff --git a/docs/assets/figures/mage-007/neighbor-access-mobile.svg b/docs/assets/figures/mage-007/neighbor-access-mobile.svg new file mode 100644 index 0000000..a6f4b4e --- /dev/null +++ b/docs/assets/figures/mage-007/neighbor-access-mobile.svg @@ -0,0 +1,2168 @@ + + + + + + + + Schematic of the neighbor kernel's access pattern before and after the change; not a measurement. See docs/experiments/mage-007.md. + image/svg+xml + + + Matplotlib v3.11.1, https://matplotlib.org/ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + diff --git a/docs/assets/figures/mage-007/neighbor-access.png b/docs/assets/figures/mage-007/neighbor-access.png new file mode 100644 index 0000000..2a2dcaf Binary files /dev/null and b/docs/assets/figures/mage-007/neighbor-access.png differ diff --git a/docs/assets/figures/mage-007/neighbor-access.svg b/docs/assets/figures/mage-007/neighbor-access.svg new file mode 100644 index 0000000..1cc28a0 --- /dev/null +++ b/docs/assets/figures/mage-007/neighbor-access.svg @@ -0,0 +1,2266 @@ + + + + + + + + Schematic of the neighbor kernel's access pattern before and after the change; not a measurement. See docs/experiments/mage-007.md. + image/svg+xml + + + Matplotlib v3.11.1, https://matplotlib.org/ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + diff --git a/docs/experiments/mage-004.md b/docs/experiments/mage-004.md index 48ba177..9d5c4a6 100644 --- a/docs/experiments/mage-004.md +++ b/docs/experiments/mage-004.md @@ -2,35 +2,29 @@ title: What the compiler chose permalink: /experiments/mage-004/ eyebrow: "Field note 004 / Mathematics on a GPU" -description: Handing the kernel decisions to a tile compiler: where it beat the hand-written kernels, where it lost, and why most of the apparent difference turned out to be the cost of starting work rather than the cost of doing it. +description: "Handing the kernel decisions to a tile compiler — where it beat the hand-written kernels, where it lost, and why most of the apparent difference turned out to be the cost of starting work rather than the cost of doing it." math: true --- -The [earlier notes]({{ '/experiments/mage-003/' | relative_url }}) wrote the kernels by hand. Each -thread owned a small block of results, and the layout of shared memory, the width of each load and -the placement of barriers were all deliberate choices that produced measured wins. - -This note hands those choices to a compiler. [cuTile Rust](https://github.com/NVlabs/cutile-rs) -lets you write a kernel as a single-threaded program over *tiles* — blocks of data — and works out -the threads, the memory layout and the tensor-core instructions itself. The question is what that is -worth on the same five operations, measured the same way. - -The answer has three parts, and the middle one is the surprise. - -**The compiler is competitive.** It beats the hand-written kernel on bias + GELU (8.19 µs against -11.0) and draws level on layer normalization (10.8 against 10.1). It loses on the two -matrix-shaped operations, 1.6× and 1.4×, and by 3.4× on neighbor aggregation, which is irregular -enough that the tile model has no safe way to express it. - -**Most of the gap in the timing column was not the kernel.** Timing that waits for every call to -finish prices the *cost of starting work*, and this runtime starts work expensively: 14–23 µs per -call against 2–3 µs for the hand-written kernels. Queue the calls up, or replay them from a -recorded graph, and the numbers fall onto the kernel times. The kernels were never as far apart as -the column said, and any comparison that blocks on every call will mislead in the same direction. - -**The library's own search beat mine.** Twelve hand-picked tile shapes moved the matrix multiply -from 738 µs to 201 µs of measured time; the bundled autotuner then found a shape that measured -137 µs against my 189 µs. Picking twelve configurations to trace a ratio is not a search. +The [earlier notes]({{ '/experiments/mage-003/' | relative_url }}) wrote these kernels by hand: each +thread owned a fixed block of results, and the width of every load, the layout of shared memory and +the placement of barriers were deliberate choices with measured effects. + +This note gives those decisions to a compiler. A [cuTile Rust](https://github.com/NVlabs/cutile-rs) +kernel is a single-threaded program over *tiles* — blocks of data — and the compiler decides how many +warps receive each tile, which values stay in registers, when loads widen to 128 bits, and when a +tensor-core instruction replaces a multiply. On the same five operations it is competitive: bias + +GELU 8.19 µs against the hand-written 11.0, layer normalization 10.76 against 10.05. It trails where +reuse is highest — matrix multiply 131.56 against 80.00, triangle contraction 116.08 against 80.1 — +and by 3.4× on neighbor aggregation, whose irregular gathers have no safe expression in the tile +model and run through raw device pointers instead. + +Two measurement facts matter more than the ranking. Timing that waits for each call prices the cost +of *starting* work: 14–23 µs per call against 2–3 µs for the hand-written kernel, which vanishes +when ten calls share one measurement — bias + GELU goes from 25.76 µs to 8.50 µs against an 8.19 µs +kernel. And tile shape is a memory decision: the same arithmetic measured 738 µs at the tutorial's +16×16×8 tile and 131.56 µs at 32×128×32, and the compiler's own autotuner found a shape 27% faster +than the best of twelve hand-picked ones. ## The five operations, two views @@ -115,7 +109,7 @@ three launch paths matched the wins and losses are the ones in the first table a same session. The kernel track has since published [mage-006]({{ '/experiments/mage-006/' | relative_url }}) with lower numbers of its own, so the two Rust columns should not be subtracted from each other.* -## The tile shape, and the search that beat me +## Tile shape is a memory decision The matrix multiply started from the upstream tutorial's 16 × 16 × 8 tile and was slow. Twelve hand-picked configurations took it from 738 µs to 201 µs, and the shape was not monotone in any @@ -130,9 +124,10 @@ limit rather than arithmetic. | 64×64×32 | 255.4 | 128×128×8 | 469.8 | Then the library's autotuner searched the same space properly — 36 candidates, each validated -before timing, kernel time alone, 17 seconds — and chose **32 × 128 × 32**, which measured -137.28 µs against my 189.44 µs when both were checked in one session. Every number in this note -was re-measured with it afterwards, which is where the matrix multiply's 131.56 µs comes from. +before timing, kernel time alone, in 17 seconds — and chose **32 × 128 × 32**. Checked in one +session, that shape measured 137.28 µs against the hand-picked 128 × 64 × 8's 189.44 µs. Every +number in this note was re-measured with it afterwards, which is where the matrix multiply's +131.56 µs comes from. ## What the compiler would not let us write diff --git a/docs/experiments/mage-007.md b/docs/experiments/mage-007.md index f8f9103..8dd3cac 100644 --- a/docs/experiments/mage-007.md +++ b/docs/experiments/mage-007.md @@ -2,63 +2,68 @@ title: Wider loads helped two kernels and hurt a third permalink: /experiments/mage-007/ eyebrow: "Field note 007 / Mathematics on a GPU" -description: The same treatment — more data per thread, loaded in one instruction — closed one gap and halved another, and then made a third kernel 2.6x slower. All three were measured the same way, and the replacement was reverted. +description: "One gap closed, one halved, and one attempt that measured 2.6x slower and was reverted, each before-and-after pair taken in one session." math: true --- -The [previous note]({{ '/experiments/mage-006/' | relative_url }}) ended with the Rust kernels -at 68.6 µs of GPU kernel time for the matrix multiply and 9.0 µs for layer normalization, and -with three places where they were still behind. The worst was layer normalization **at wide -rows**: at 4096×4096 it was 1.41× slower than Triton's kernel, because the shape fell back to -an older kernel that handles one row at a time. Two more of the five operations, bias + GELU -and neighbor aggregation, had never been touched since the first comparison. +Neighbor aggregation sums, for each output row, the input rows named by that row's edges. Each +thread read one float per edge — a 32-bit load, one edge at a time — so a row with 16 edges issued +16 narrow loads in a dependent chain: -All three want the same kind of help. A thread that handles four values at once loads them -with one instruction instead of four, and measures have shown that shape winning before: it is -what took layer normalization from 18.6 µs to 11.0 µs, and what made the matrix multiply's -shared reads cheap. So the natural next move was to widen what each thread handles in all -three kernels and measure what happened. - -Two of the three got faster, and the third got much slower. **Neighbor aggregation** went from -9.45 µs to 7.06 µs once each thread took a 128-bit quad of features and two edge gathers -overlapped. **Layer normalization at wide rows** turned out not to need a rewrite at all: the -two-warp kernel that serves narrower rows was already sound at 4096 and was simply excluded by -a stale size limit, so removing the limit took it from 255.41 µs to 186.80 µs. And **bias + -GELU**, given exactly the same treatment, measured 27.00 µs against the 10.29 µs it already -had — 2.6× slower — and the change was reverted. - -That third result is the reason this note exists. Widening a thread's load is not a rule that -transfers between kernels; it trades instruction count for occupancy, and it pays only when the -kernel is short of instructions. The elementwise kernel already has enough threads to saturate -the device, so giving each of them four times the work only removed parallelism. - -The gaps that remain are smaller: neighbor aggregation is 1.2× behind Triton, and bias + GELU -is unchanged at 11.0 µs against Triton's 7.8. - -## Method - -| | | -| --- | --- | -| Instrument | GPU kernel time, one Nsight Systems capture per measurement, mean of 100 launches after 25 warm-up launches | -| Host | RTX 4090 (sm_89), driver 591.74, WSL2 Ubuntu 22.04, device idle | -| Session discipline | Each before/after pair was measured in one session, with the same inputs and the same warm-up, on the same binary path; the device was checked idle before each pass | -| Correctness | Full-output check against PyTorch FP32 with TF32 disabled, `rtol = atol = 1e-4`, before any timing was believed | -| References | Triton and cuTile bars are quoted from their own sessions (mage-004, mage-006) and are drawn hatched in the figure because they are not pairings | - -The figure is [`docs/assets/figures/mage-007/kernel-lane-fixes.svg`](../assets/figures/mage-007/kernel-lane-fixes.svg), -drawn by `scripts/plot-kernel-lane-fixes.py` from -[`docs/assets/results/mage-007/kernel-lane-fixes.json`](../assets/results/mage-007/kernel-lane-fixes.json). +
+
+ + + Two schematics of one thread's memory traffic. Before: one feature per thread, a 32-bit load per edge, one edge in flight. After: four features per thread in a single 128-bit load, two edges unrolled and in flight. + +
+

What one thread reads in the neighbor kernel, before and after. Drawing, not a measurement: the load widths and the edges in flight are read from the two kernels' code.

+
+ Download SVG + PNG + Script ↗ +
+
+
+
-## Values +Four features per thread makes that one 128-bit load, and unrolling the edge walk by two puts two +gathers in flight. Measured 9.45 → 7.06 µs across 100 launches at 4096×64×65536. Triton's kernel, +which loads 32 edges at once into a tile, is still ahead at 5.97 µs; the gathered rows are about +17 MB, which is L2 traffic on this part, so more edges in flight is what is left. + +Layer normalization at 4096-wide rows needed no rewrite. The kernel for these rows splits each row +across two warps, and the two partial sums meet in 64 bytes of shared memory behind one barrier. +It was gated at width 2048, so 4096 fell through to a kernel that walks the row in a single warp. +The gate was stale: the only structural requirement is that each warp's span be a whole number of +32-lane steps, and at 4096 each half is 2048 elements, which is. Removing the gate measured +255.41 → 186.80 µs, against Triton's 181.59. + +Bias + GELU took the same treatment and lost 2.6×. The shape is 4096×768: 3.1M elements, or 12,288 +blocks of 256 threads with one element each, which already fills the device. Four elements per +thread cuts that to 786,432 threads and gives each thread four serial `tanh` evaluations. Measured +27.00 µs against 10.29 µs, and reverted. Vector width pays when a thread is short of work to +issue; it costs when the grid is already saturating the machine. + +| Operation | Before | After | Reference (other session) | +| --- | ---: | ---: | --- | +| Neighbor aggregation 4096×64×65536 | 9.45 | **7.06** | Triton 5.97, PyTorch 120.93 | +| LayerNorm 4096×4096 | 255.41 | **186.80** | Triton 181.59 | +| Bias + GELU 4096×768 | 10.29 (kept) | 27.00 (reverted) | Triton 7.76, cuTile 7.97 | + +*Microseconds of GPU kernel time, mean of 100 launches after 25 warm-up launches. Each before and +after pair was measured in one session on an idle device; the references are quoted from mage-004 +and mage-006 and are drawn hatched in the figure below because they are not pairings.*
Three rows of horizontal bars, GPU kernel time in microseconds, each row before and after one change. LayerNorm at 4096 by 4096: the single-warp kernel 255.41 becomes the two-warp kernel 186.80, with Triton at 181.59 for reference. Neighbor aggregation at 4096 by 64 by 65536: the scalar kernel 9.45 becomes the feature-quad kernel 7.06, with Triton at 5.97 for reference. Bias plus GELU at 4096 by 768: the scalar kernel stays at 10.29 and the feature-quad variant that measured 27.00, 2.6 times slower, was reverted, with Triton at 7.76 and cuTile Rust at 7.97 for reference. + alt="Three rows of horizontal bars of GPU kernel time in microseconds, each row before and after one change. Neighbor aggregation falls from 9.45 to 7.06 with Triton at 5.97 for reference. LayerNorm at 4096 by 4096 falls from 255.41 to 186.80 with Triton at 181.59. Bias plus GELU keeps its scalar kernel at 10.29 while the feature-quad variant that measured 27.00 was reverted, with Triton at 7.76 and cuTile Rust at 7.97.">
-

GPU kernel time before and after each change, one Nsight Systems capture of 100 launches per bar, mean. Each before-and-after pair was measured in one session on an idle device; the Triton and cuTile bars are drawn hatched because they are quoted from their own sessions (mage-006 and mage-004) and are references, not pairings. Lower is better.

+

The same three changes as bars, with the other-session references hatched. Lower is better.

Download SVG PNG @@ -67,11 +72,11 @@ drawn by `scripts/plot-kernel-lane-fixes.py` from
Read the plotted values (µs of GPU kernel time) - + - +
OperationBeforeAfterReference
LayerNorm 4096×4096255.41186.80Triton 181.59
Neighbor aggregation 4096×64×655369.457.06Triton 5.97
LayerNorm 4096×4096255.41186.80Triton 181.59
Bias + GELU 4096×76810.29 (kept)27.00 (reverted)Triton 7.76, cuTile 7.97
@@ -80,58 +85,23 @@ drawn by `scripts/plot-kernel-lane-fixes.py` from
-Kernel time in microseconds, mean of 100 launches. - -| Operation | Before | After | Change | Reference (other session) | -| --- | ---: | ---: | ---: | --- | -| LayerNorm 4096×4096 | 255.41 (`layer_norm_warp`) | **186.80** (`layer_norm_pair`) | 1.37× faster | Triton 181.59 (mage-006) | -| Neighbor aggregation 4096×64×65536 | 9.45 (scalar) | **7.06** (feature quad) | 1.34× faster | Triton 5.97, PyTorch 120.93 (mage-006) | -| Bias + GELU 4096×768 | 10.29 (scalar, kept) | 27.00 (feature quad, **rejected**) | 2.6× slower | Triton 7.76, cuTile 7.97 (mage-004) | - -Worst full-output error: 4.77e-07 (LayerNorm), 3.58e-07 (neighbor), 5.31e-06 (GELU). - -## What each change was +## Limits of these numbers -**LayerNorm at wide rows** (issue #54, PR #69). The two-warp row kernel was gated at -`width <= 2048`, so a 4096-wide row fell back to `layer_norm_warp` — the capture named -the kernel, which is how the fallback was confirmed rather than assumed. The pairwise -kernel's own soundness condition, each warp's span a multiple of 32 lanes, already -held at 4096: the cap was a stale heuristic. Raising it to 4096 is the whole change. +- No hardware counters are available, so "L2 traffic" is arithmetic on the bytes moved, not a + measurement. +- One capture per point: a kernel time carries no interval of its own. +- The references come from other sessions (`scripts/evolve_capture.py --all` runs every + implementation in one session, and has not been run on these shapes). +- The LayerNorm change is only validated where the harness points it: 4096×4096, 4096×512, + 3072×1024, 2048×2048. Widths above 4096 still fall through. +- GELU's rejection is one pair, not a sweep: two features per thread, and quads with other block + sizes, were not tried. -**Neighbor aggregation** (issue #59, PR #71). The scalar kernel gives each thread one -feature and a serial edge loop, so one gather is in flight per thread. The quad form -gives each thread four features — one 128-bit load per edge — and unrolls the edge -walk by two, so two independent gathers overlap. Triton's 5.97 µs remains ahead, so -the gap narrows from 1.7× to 1.2×; the kernel moves about 17 MB of gathered rows in -7 µs, which is L2-bandwidth work on this part, so more edges in flight per thread is -the next lever. - -**Bias + GELU** (issue #59, PR #71). The same vectorization was tried and rejected: -four features per thread measured 27.00 µs against the scalar kernel's 10.29 µs in the -same session. The LayerNorm treatment does not transfer to this elementwise shape — -one element per thread wins, most likely on occupancy. Kept here because it is the -counterexample: "vectorize the elementwise kernel" is not a rule. - -## What the numbers do not establish - -- No hardware counters are available, so the L2-bandwidth reading of the neighbor - kernel is inferred from traffic arithmetic, not measured. -- One capture per point: a kernel time here carries no interval of its own. -- The references are quoted from other sessions. A consolidated pass - (`scripts/evolve_capture.py --all`) runs every implementation in one session, and - running it would make these four arms a single measurement. -- The LayerNorm change is validated at 4096×4096, 4096×512, 3072×1024 and 2048×2048; - widths above 4096 still fall back, and were not measured. -- GELU's rejection is a single before/after pair, not a sweep: other vector widths - (two features, or quads with a different block size) were not tried. - -## Reproduction +## Reproduce ```bash source scripts/oxide-env.sh cd examples/oxide && CARGO_BUILD_JOBS=4 cargo oxide build --arch sm_89 && cd ../.. - -# inputs in the harness layout, at the shape under test .venv/bin/python -c " import sys; sys.path.insert(0, 'examples/oxide') import experiment @@ -139,9 +109,10 @@ from pathlib import Path experiment.generate(Path('artifacts/mage-007/layernorm-4096x4096'), 'layernorm', [4096, 4096], warmup=25, iterations=100) " - -# one operation, one capture; the kernel time is in the report's kernels.json -.venv/bin/python -m mage profile-exec --backend nsys --capture-range cuda --output-dir artifacts/mage-007/layernorm-4096x4096/nsys -- examples/oxide/target/release/mage-oxide artifacts/mage-007/layernorm-4096x4096 --iterations 100 --capture - +.venv/bin/python -m mage profile-exec --backend nsys --capture-range cuda \ + --output-dir artifacts/mage-007/layernorm-4096x4096/nsys -- \ + examples/oxide/target/release/mage-oxide artifacts/mage-007/layernorm-4096x4096 \ + --iterations 100 --capture uv run --script scripts/plot-kernel-lane-fixes.py +uv run --script scripts/plot-neighbor-access.py ``` diff --git a/scripts/check-front-matter.py b/scripts/check-front-matter.py new file mode 100644 index 0000000..171d0f4 --- /dev/null +++ b/scripts/check-front-matter.py @@ -0,0 +1,38 @@ +"""Pre-flight: every note's front matter must parse as YAML before it is pushed. + +A colon inside an unquoted value breaks the site build, and the failure surfaces +only as a failed deploy after the merge. Run this before pushing docs changes. + +Run: uv run --script scripts/check-front-matter.py +""" +import sys +from pathlib import Path + +ROOT = Path(__file__).resolve().parents[1] + +try: + import yaml +except ImportError: # pragma: no cover - PyYAML is a Jekyll-side concern + print("PyYAML is not installed; skipping (uv pip install pyyaml)") + sys.exit(0) + +failures = [] +checked = 0 +for path in sorted((ROOT / "docs").rglob("*.md")): + text = path.read_text(encoding="utf-8") + if not text.startswith("---\n"): + continue + end = text.find("\n---", 4) + if end == -1: + failures.append((path, "front matter is not closed")) + continue + checked += 1 + try: + yaml.safe_load(text[4:end]) + except yaml.YAMLError as error: + failures.append((path, str(error).splitlines()[0])) + +print(f"checked {checked} documents with front matter") +for path, reason in failures: + print(f"FAIL {path.relative_to(ROOT)}: {reason}") +sys.exit(1 if failures else 0) diff --git a/scripts/plot-neighbor-access.py b/scripts/plot-neighbor-access.py new file mode 100644 index 0000000..235aa89 --- /dev/null +++ b/scripts/plot-neighbor-access.py @@ -0,0 +1,100 @@ +"""Schematic: what one thread reads in the neighbor kernel, before and after. + +Not a measurement — a drawing of the access pattern, taken from the two kernels' +code, with the load widths and the traffic they generate annotated. It exists +because the change is easier to see than to read: the same work, addressed +differently. + +Run: uv run --script scripts/plot-neighbor-access.py +""" +from pathlib import Path + +import matplotlib + +matplotlib.use("Agg") +import matplotlib.pyplot as plt # noqa: E402 +from matplotlib.patches import FancyArrowPatch, Rectangle # noqa: E402 + +ROOT = Path(__file__).resolve().parents[1] +OUT = ROOT / "docs/assets/figures/mage-007" + +BG, INK, MUTED, RULE = "#101217", "#edf0f5", "#a0a9b9", "#303641" +ACCENT, WARM = "#91dbba", "#e0a08a" +FLOAT_W, FLOAT_H = 13.0, 11.0 + + +def draw_rows(ax, x0, y0, rows, cols, label, color, quad_span=None): + """Draw a few x[] rows as grids of floats, optionally highlighting a quad.""" + for r in range(rows): + for c in range(cols): + highlighted = quad_span is not None and quad_span[0] <= c < quad_span[1] + ax.add_patch(Rectangle((x0 + c * FLOAT_W, y0 - r * FLOAT_H), FLOAT_W - 1.4, FLOAT_H - 1.4, + facecolor=color if highlighted else RULE, + edgecolor=INK if highlighted else "none", linewidth=.9, zorder=3)) + ax.text(x0 - 6, y0 - r * FLOAT_H + FLOAT_H / 2 - 1, label if r == 0 else "", + color=MUTED, fontsize=8, ha="right", va="center") + + +def panel(ax, title, subtitle, load_bits, per_thread, edges_in_flight, quad_span, color): + ax.set_xlim(0, 150) + ax.set_ylim(0, 100) + ax.axis("off") + ax.text(2, 92, title, color=INK, fontsize=11.5, weight="bold", va="top") + ax.text(2, 83, subtitle, color=MUTED, fontsize=8.6, va="top") + + # the x[] rows this thread reads through its edges + ax.text(2, 66, "x rows named by the edges", color=MUTED, fontsize=8.4, va="top") + draw_rows(ax, 46, 60, rows=3, cols=8, label="row", color=color, quad_span=quad_span) + ax.text(46, 22, f"{per_thread} load{'s' if per_thread > 1 else ''} per thread per edge, " + f"{load_bits} bits each", color=color, fontsize=9, va="top", weight="bold") + ax.text(46, 13, f"{edges_in_flight} edge{'s' if edges_in_flight > 1 else ''} in flight", color=MUTED, + fontsize=8.6, va="top") + + # the arrow that stands for the gather + for i in range(min(edges_in_flight, 2)): + ax.add_patch(FancyArrowPatch((40, 62 - i * 14), (44.5, 58 - i * 14), arrowstyle="-|>", + mutation_scale=9, color=color, linewidth=1.4, zorder=4)) + ax.text(2, 30, "weights, indices\nread per edge", color=MUTED, fontsize=8.4, va="top") + + +def main(): + fig, axes = plt.subplots(1, 2, figsize=(9.6, 3.5)) + fig.subplots_adjust(left=.01, right=.99, top=.80, bottom=.02, wspace=.06) + fig.text(.01, .965, "One thread's memory traffic in the neighbor kernel", color=INK, + fontsize=13, weight="bold", va="top") + fig.text(.01, .915, "Each output row is a sum over its edges; the question is how many bytes one " + "thread asks for at a time", color=MUTED, fontsize=9, va="top") + panel(axes[0], "before", "one feature per thread, one edge per step", 32, 1, 1, None, WARM) + panel(axes[1], "after", "four features per thread, two edges unrolled", 128, 1, 2, (2, 6), ACCENT) + for ax in axes: + ax.set_facecolor(BG) + fig.patch.set_facecolor(BG) + OUT.mkdir(parents=True, exist_ok=True) + metadata = {"Date": None, "Description": + "Schematic of the neighbor kernel's access pattern before and after the change; " + "not a measurement. See docs/experiments/mage-007.md."} + fig.savefig(OUT / "neighbor-access.svg", metadata=metadata) + svg = OUT / "neighbor-access.svg" + svg.write_bytes(b"\n".join(line.rstrip() for line in svg.read_bytes().splitlines()) + b"\n") + fig.savefig(OUT / "neighbor-access.png", dpi=200, metadata=metadata) + plt.close(fig) + + # the stacked variant the notes reference on narrow screens + fig, axes = plt.subplots(2, 1, figsize=(3.6, 5.4)) + fig.subplots_adjust(left=.02, right=.98, top=.90, bottom=.03, hspace=.35) + fig.text(.02, .975, "One thread's memory traffic", color=INK, fontsize=10.5, weight="bold", va="top") + fig.text(.02, .935, "in the neighbor kernel, before and after", color=MUTED, fontsize=8, va="top") + panel(axes[0], "before", "one feature per thread, 32-bit loads", 32, 1, 1, None, WARM) + panel(axes[1], "after", "four features per thread, 128-bit loads", 128, 1, 2, (2, 6), ACCENT) + for ax in axes: + ax.set_facecolor(BG) + fig.patch.set_facecolor(BG) + fig.savefig(OUT / "neighbor-access-mobile.svg", metadata=metadata) + mobile = OUT / "neighbor-access-mobile.svg" + mobile.write_bytes(b"\n".join(line.rstrip() for line in mobile.read_bytes().splitlines()) + b"\n") + plt.close(fig) + print("wrote", OUT / "neighbor-access.svg") + + +if __name__ == "__main__": + main()