1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
// Shared threadgroup reduction primitives.
//
// Canonical owner of the per-row reduction. `reduce.metal` and `rms_norm.metal`
// both reduce a row inside one threadgroup, and each .metal file is a separate
// translation unit, so without a header the second one to need this copies
// it. Two copies of a reduction that must agree on lane assignment and
// barrier placement is a defect waiting for one of them to be edited.
//
// `build.rs` compiles only `*.metal`, so this file is included, never compiled
// on its own — and `track_kernel_sources` emits `rerun-if-changed` for `.h` as
// well, without which editing this header would leave every dependent kernel
// silently stale in the metallib while the tests reported a pass.
using namespace metal;
/// Largest threadgroup a row reduction is launched with; the host caps lanes
/// per row well below it (see `reduce_tptg` in src/nn.rs).
constant uint REDUCE_MAX_TG = 1024u;
/// Simdgroups a launch can hold: `REDUCE_MAX_TG` lanes over 32-wide simdgroups.
constant uint REDUCE_MAX_SIMDGROUPS = REDUCE_MAX_TG / 32u;
/// Sum `v` across the threadgroup, returned in every lane.
///
/// Simdgroup-first, two levels: `simd_sum` folds the 32 lanes with no memory
/// traffic, one lane per simdgroup publishes its partial, one barrier, and
/// then every simdgroup folds the partials with a second `simd_sum` (lane
/// `l` takes partial `l`), so the total is uniform in every lane with no
/// serial loop and no second barrier. A launch that fits one simdgroup never
/// touches `scratch` or a barrier.
///
/// The tree this replaced (`REDUCE_TREE`, retired 2026-09-05) halved the whole threadgroup
/// through threadgroup memory with a barrier per round — ten rounds and
/// eleven barriers at 1024 lanes. A first cut folded the partials with a
/// serial loop per thread instead of the second shuffle, and lost to the
/// tree from 32 simdgroups up: 32 dependent threadgroup-memory loads per
/// thread cost more than ten barriers. Measured 2026-09-05, see
/// `bench/results/row_reduction_simd_m5pro.txt`.
///
/// `scratch` must hold `REDUCE_MAX_SIMDGROUPS` floats. The reads of `scratch`
/// happen after the barrier here, so a second reduction must not write the
/// same region until another barrier separates them: give each reduction its
/// own region (`scratch + REDUCE_MAX_SIMDGROUPS`) rather than reusing one.
inline float
/// Maximum of `v` across the threadgroup, returned in every lane. Same
/// structure and scratch contract as [`reduce_row_add`].
inline float