1/* @ajax/mesh-rt-block: the block collectives of a multi-warp node in a fixed order (barrier,
2 * broadcast, reductions, scans), the fixed f32 operations, and typed windows
3 * of the node's shared memory.
4 *
5 * Every thread of the node must call a collective, the same ones in the same
6 * order: never after a branch the threads may take differently, never after
7 * some threads returned from the task. Warps with nothing to do still call.
8 *
9 * A fixed order: the float results are bitwise the same on every GPU and in
10 * the CPU reference (mesh_ref::block), because the combination tree is fixed
11 * and every f32 op is correctly rounded, unfused and FTZ:
12 * - a reduction: within each warp a butterfly (lane l combines with lane
13 * l ^ m for m = 16, 8, 4, 2, 1), then the same butterfly over the 32 warp
14 * results, missing warps filled with the identity (sum: -0.0). The tree is
15 * the one over 1024 leaves whatever `warps` is;
16 * - a scan (inclusive): within each warp Hillis-Steele (lane l adds lane
17 * l - d for d = 1, 2, 4, 8, 16, the earlier operand first), the warp totals
18 * scanned the same way, then each warp after the first combines its prefix
19 * with its own scan (prefix + x); exclusive: the thread before's inclusive
20 * value within a warp, a warp's first thread the prefix (warp 0: identity);
21 * - max over f32 is the IEEE total order on the bits: one of the inputs,
22 * unchanged. u32/u64 sums wrap.
23 * The scratch is the runtime's shared area (RT_SHARED_BYTES): each collective
24 * starts with a barrier, so they may follow each other directly.
25 *
26 * Rust's generic `Word` and closures become one function per type (u32, f32,
27 * u64) taking the operation as a function pointer; every call is inlined, so
28 * the pointer is a constant and the operation is inlined too.
29 * Ported from alpha-rust-mesh crates/mesh-rt/src/block.rs. */
30#ifndef MESH_RT_BLOCK_H
31#define MESH_RT_BLOCK_H
3233#include <stdbool.h>
34#include <stdint.h>
3536#include "mesh/abi/base.h"
37#include "mesh/abi/fault.h"
38#include "mesh/rt/prim.h"
39#include "mesh/rt/view.h"
4041/* The runtime's shared area (RT_SHARED_BYTES), byte offsets. The node loop's
42 * ticket descriptor (thread 0 writes it, the block reads it after a barrier). */
43#define MESH_RT_DESC 0x00u
44#define MESH_RT_D_T 0x00u
45#define MESH_RT_D_GO 0x04u
46#define MESH_RT_D_OP 0x08u
47#define MESH_RT_D_ARG 0x0cu
48#define MESH_RT_D_CSLOT 0x10u
49/* One 8-byte word per warp (reductions, scans). */
50#define MESH_RT_RED 0x40u
51/* The broadcast word. */
52#define MESH_RT_BCAST (MESH_RT_RED + 8u * 32u)
53_Static_assert(MESH_RT_DESC + 0x20u <= MESH_RT_RED && MESH_RT_BCAST + 8u <= MESH_RT_SHARED_BYTES,
54 "the runtime's shared area holds the descriptor, the warp words and the broadcast word");
5556/* ---- the fixed operations (mesh_ref::block has the same) */
57/* f32 sum's identity: -0.0 (x + -0.0 == x for every x, +0 and -0 included). */
58#define MESH_RT_SUM_F32_ID (-0.0f)
59/* f32 max's identity: the least value of the total order (a negative NaN, all bits set). */
60#define MESH_RT_MAX_F32_ID_BITS 0xffffffffu
61MESH_RT_INLINE uint32_t mesh_rt_order_key(uint32_t b) { return (b & 0x80000000u) ? ~b : (b | 0x80000000u); }
62/* The larger of a and b in the IEEE total order of their bits (one of them, unchanged). */
63MESH_RT_INLINE float mesh_rt_max_f32(float a, float b) {
64 return mesh_rt_order_key(mesh_f32_bits(b)) > mesh_rt_order_key(mesh_f32_bits(a)) ? b : a;
65}
66MESH_RT_INLINE float mesh_rt_add_f32(float a, float b) { return a + b; }
67MESH_RT_INLINE uint32_t mesh_rt_add_u32(uint32_t a, uint32_t b) { return a + b; }
68MESH_RT_INLINE uint32_t mesh_rt_max_u32(uint32_t a, uint32_t b) { return a > b ? a : b; }
69MESH_RT_INLINE uint32_t mesh_rt_min_u32(uint32_t a, uint32_t b) { return a < b ? a : b; }
70MESH_RT_INLINE uint64_t mesh_rt_add_u64(uint64_t a, uint64_t b) { return a + b; }
71MESH_RT_INLINE uint64_t mesh_rt_max_u64(uint64_t a, uint64_t b) { return a > b ? a : b; }
72MESH_RT_INLINE uint64_t mesh_rt_min_u64(uint64_t a, uint64_t b) { return a < b ? a : b; }
7374typedef uint32_t (*mesh_rt_op_u32)(uint32_t, uint32_t);
75typedef float (*mesh_rt_op_f32)(float, float);
76typedef uint64_t (*mesh_rt_op_u64)(uint64_t, uint64_t);
7778/* ---- moving words between lanes and through shared memory (Rust's Word) */
79MESH_RT_INLINE uint32_t mesh_rt_word_u32_xor(uint32_t v, uint32_t m) { return mesh_rt_shfl_bfly(v, m); }
80MESH_RT_INLINE uint32_t mesh_rt_word_u32_up(uint32_t v, uint32_t d) { return mesh_rt_shfl_up(v, d); }
81MESH_RT_INLINE uint32_t mesh_rt_word_u32_idx(uint32_t v, uint32_t s) { return mesh_rt_shfl_idx(v, s); }
82MESH_RT_INLINE uint32_t mesh_rt_word_u32_load(const uint8_t *p) { return *(const volatile uint32_t *)(const volatile void *)p; }
83MESH_RT_INLINE void mesh_rt_word_u32_store(uint32_t v, uint8_t *p) { *(volatile uint32_t *)(volatile void *)p = v; }
8485MESH_RT_INLINE float mesh_rt_word_f32_xor(float v, uint32_t m) { return mesh_f32_from_bits(mesh_rt_shfl_bfly(mesh_f32_bits(v), m)); }
86MESH_RT_INLINE float mesh_rt_word_f32_up(float v, uint32_t d) { return mesh_f32_from_bits(mesh_rt_shfl_up(mesh_f32_bits(v), d)); }
87MESH_RT_INLINE float mesh_rt_word_f32_idx(float v, uint32_t s) { return mesh_f32_from_bits(mesh_rt_shfl_idx(mesh_f32_bits(v), s)); }
88MESH_RT_INLINE float mesh_rt_word_f32_load(const uint8_t *p) { return mesh_f32_from_bits(mesh_rt_word_u32_load(p)); }
89MESH_RT_INLINE void mesh_rt_word_f32_store(float v, uint8_t *p) { mesh_rt_word_u32_store(mesh_f32_bits(v), p); }
9091MESH_RT_INLINE uint64_t mesh_rt_word_u64_xor(uint64_t v, uint32_t m) {
92 return (uint64_t)mesh_rt_shfl_bfly((uint32_t)v, m) | (uint64_t)mesh_rt_shfl_bfly((uint32_t)(v >> 32), m) << 32;
93}
94MESH_RT_INLINE uint64_t mesh_rt_word_u64_up(uint64_t v, uint32_t d) {
95 return (uint64_t)mesh_rt_shfl_up((uint32_t)v, d) | (uint64_t)mesh_rt_shfl_up((uint32_t)(v >> 32), d) << 32;
96}
97MESH_RT_INLINE uint64_t mesh_rt_word_u64_idx(uint64_t v, uint32_t s) {
98 return (uint64_t)mesh_rt_shfl_idx((uint32_t)v, s) | (uint64_t)mesh_rt_shfl_idx((uint32_t)(v >> 32), s) << 32;
99}
100MESH_RT_INLINE uint64_t mesh_rt_word_u64_load(const uint8_t *p) { return *(const volatile uint64_t *)(const volatile void *)p; }
101MESH_RT_INLINE void mesh_rt_word_u64_store(uint64_t v, uint8_t *p) { *(volatile uint64_t *)(volatile void *)p = v; }
102103/* The collectives for one word type T (suffix S): warp_reduce, warp_scan,
104 * block reduce, scan and broadcast. */
105#define MESH_RT_COLLECTIVES(T, S) \
106 /* Butterfly over the warp's 32 lanes (m = 16, 8, 4, 2, 1); every lane ends with the same value when op is \
107 * commutative. Warp-collective. */ \
108 MESH_RT_INLINE T mesh_rt_warp_reduce_##S(T v, mesh_rt_op_##S op) { \
109 for (uint32_t m = 16; m > 0; m >>= 1) v = op(v, mesh_rt_word_##S##_xor(v, m)); \
110 return v; \
111 } \
112 /* Inclusive Hillis-Steele scan over the warp's lanes (d = 1, 2, 4, 8, 16: lane l >= d computes \
113 * op(lane l - d's value, its own)). Warp-collective. */ \
114 MESH_RT_INLINE T mesh_rt_warp_scan_##S(T v, mesh_rt_op_##S op) { \
115 uint32_t l = mesh_rt_lane(); \
116 for (uint32_t d = 1; d < 32; d <<= 1) { \
117 T y = mesh_rt_word_##S##_up(v, d); \
118 if (l >= d) v = op(y, v); \
119 } \
120 return v; \
121 } \
122 /* Block reduction in the fixed order: every thread gets the result. */ \
123 MESH_RT_INLINE T mesh_rt_block_reduce_##S(T v, T id, mesh_rt_op_##S op) { \
124 uint32_t l = mesh_rt_lane(), w = mesh_rt_tid() >> 5, warps = mesh_rt_ntid() >> 5; \
125 uint8_t *s = mesh_rt_smem_base(); \
126 mesh_rt_bar_sync(); \
127 T r = mesh_rt_warp_reduce_##S(v, op); \
128 if (l == 0) mesh_rt_word_##S##_store(r, s + MESH_RT_RED + 8u * w); \
129 mesh_rt_bar_sync(); \
130 T x = l < warps ? mesh_rt_word_##S##_load(s + MESH_RT_RED + 8u * l) : id; \
131 return mesh_rt_warp_reduce_##S(x, op); \
132 } \
133 /* Block scan in the fixed order: inclusive, or exclusive (thread 0: id). */ \
134 MESH_RT_INLINE T mesh_rt_block_scan_##S(T v, T id, mesh_rt_op_##S op, bool inclusive) { \
135 uint32_t l = mesh_rt_lane(), w = mesh_rt_tid() >> 5, warps = mesh_rt_ntid() >> 5; \
136 uint8_t *s = mesh_rt_smem_base(); \
137 mesh_rt_bar_sync(); \
138 T inc = mesh_rt_warp_scan_##S(v, op); \
139 if (l == 31) mesh_rt_word_##S##_store(inc, s + MESH_RT_RED + 8u * w); \
140 mesh_rt_bar_sync(); \
141 T t = l < warps ? mesh_rt_word_##S##_load(s + MESH_RT_RED + 8u * l) : id; \
142 T ts = mesh_rt_warp_scan_##S(t, op); \
143 /* The total of warps 0..w-1 (lane w - 1's inclusive scan of the totals). */ \
144 T prefix = mesh_rt_word_##S##_idx(ts, (w - 1u) & 31u); \
145 if (w != 0) inc = op(prefix, inc); \
146 if (inclusive) return inc; \
147 T up = mesh_rt_word_##S##_up(inc, 1); \
148 if (l > 0) return up; \
149 return w == 0 ? id : prefix; \
150 } \
151 /* Thread src's v on every thread of the block. */ \
152 MESH_RT_INLINE T mesh_rt_block_broadcast_##S(T v, uint32_t src) { \
153 uint8_t *s = mesh_rt_smem_base(); \
154 mesh_rt_bar_sync(); \
155 if (mesh_rt_tid() == src) mesh_rt_word_##S##_store(v, s + MESH_RT_BCAST); \
156 mesh_rt_bar_sync(); \
157 return mesh_rt_word_##S##_load(s + MESH_RT_BCAST); \
158 }
159160MESH_RT_COLLECTIVES(uint32_t, u32)
161MESH_RT_COLLECTIVES(float, f32)
162MESH_RT_COLLECTIVES(uint64_t, u64)
163164/* ---- a typed window of the node's shared memory (Ctx::shared): len
165 * elements, the same memory for every thread of the node and private to it.
166 * Ordering between threads is the program's: write, barrier, read. get/set
167 * are bounds checked (out of range: 0 / dropped, and FAULT_SHARED_BOUNDS | i << 8);
168 * the _unchecked ones and the raw pointer are not. */
169typedef struct {
170 uint8_t *base;
171 uint32_t len; /* elements */
172 uint8_t *ctl;
173} mesh_shared;
174175#define MESH_RT_SHARED_ACCESS(T, S) \
176 MESH_RT_INLINE T mesh_shared_get_##S(mesh_shared sh, uint32_t i) { \
177 if (i < sh.len) return ((T *)(void *)sh.base)[i]; \
178 mesh_rt_fault(sh.ctl, MESH_FAULT_SHARED_BOUNDS | i << 8); \
179 return (T)0; \
180 } \
181 MESH_RT_INLINE void mesh_shared_set_##S(mesh_shared sh, uint32_t i, T v) { \
182 if (i < sh.len) ((T *)(void *)sh.base)[i] = v; \
183 else mesh_rt_fault(sh.ctl, MESH_FAULT_SHARED_BOUNDS | i << 8); \
184 } \
185 MESH_RT_INLINE T mesh_shared_get_unchecked_##S(mesh_shared sh, uint32_t i) { return ((T *)(void *)sh.base)[i]; } \
186 MESH_RT_INLINE void mesh_shared_set_unchecked_##S(mesh_shared sh, uint32_t i, T v) { ((T *)(void *)sh.base)[i] = v; } \
187 MESH_RT_INLINE T *mesh_shared_ptr_##S(mesh_shared sh) { return (T *)(void *)sh.base; }
188189MESH_RT_SHARED_ACCESS(uint32_t, u32)
190MESH_RT_SHARED_ACCESS(float, f32)
191MESH_RT_SHARED_ACCESS(uint64_t, u64)
192193#endif