==> src/block.c <==1/* The block collectives on the CPU (mesh-ref block.rs), generated for f32, u32 and u64. */
2#include "mesh/ref/block.h"
3
4#include "mesh/ftz.h"
5
6float mesh_block_sum_f32_id(void) { return mesh_ftz_from_bits(MESH_BLOCK_SUM_F32_ID_BITS); }
7float mesh_block_max_f32_id(void) { return mesh_ftz_from_bits(MESH_BLOCK_MAX_F32_ID_BITS); }
8
9static uint32_t order_key(uint32_t b) { return (b & 0x80000000u) ? ~b : (b | 0x80000000u); }
10
11float mesh_block_max_f32(float a, float b) {
12 return order_key(mesh_ftz_bits(b)) > order_key(mesh_ftz_bits(a)) ? b : a;
13}
14
15float mesh_block_add_f32(float a, float b) { return mesh_ftz_add(a, b); }
16uint32_t mesh_block_add_u32(uint32_t a, uint32_t b) { return a + b; }
17uint64_t mesh_block_add_u64(uint64_t a, uint64_t b) { return a + b; }
18
19uint32_t mesh_block_mix(uint32_t x) {
20 x ^= x >> 16;
21 x *= 0x7feb352du;
22 x ^= x >> 15;
23 x *= 0x846ca68bu;
24 return x ^ (x >> 16);
25}
26
27void mesh_block_test_values(uint32_t seed, uint32_t t, float *f, uint32_t *u, uint64_t *w) {
28 uint32_t h = mesh_block_mix((seed * 0x9e3779b9u) ^ t);
29 uint32_t h2 = mesh_block_mix(h ^ 0x85ebca6bu);
30 uint32_t h3 = mesh_block_mix(h2 ^ 0xc2b2ae35u);
31 uint32_t bits;
32 switch (h & 15) {
33 case 0: bits = (h & 16) ? 0x80000000u : 0; break; /* a zero */
34 case 1: bits = (h & 0x80000000u) | (h2 & 0x007fffffu) | 1; break; /* subnormal */
35 case 2: bits = (h & 0x80000000u) | (200u << 23) | (h2 & 0x007fffffu); break; /* ~2^73 */
36 default: bits = (h & 0x80000000u) | ((112 + (h2 >> 28)) << 23) | (h3 & 0x007fffffu); /* 2^-15..2^0 */
37 }
38 *f = mesh_ftz_from_bits(bits);
39 *u = h2;
40 *w = (uint64_t)h3 << 32 | h;
41}
42
43#define MESH_BLOCK_DEFINE(T, S) \
44 void mesh_block_warp_reduce_##S(T v[32], mesh_block_op_##S op) { \
45 for (unsigned m = 16; m >= 1; m >>= 1) { \
46 T o[32]; \
47 for (unsigned l = 0; l < 32; l++) o[l] = v[l]; \
48 for (unsigned l = 0; l < 32; l++) v[l] = op(o[l], o[l ^ m]); \
49 } \
50 } \
51 void mesh_block_warp_scan_##S(T v[32], mesh_block_op_##S op) { \
52 for (unsigned d = 1; d <= 16; d <<= 1) { \
53 T o[32]; \
54 for (unsigned l = 0; l < 32; l++) o[l] = v[l]; \
55 for (unsigned l = d; l < 32; l++) v[l] = op(o[l - d], o[l]); \
56 } \
57 } \
58 T mesh_block_reduce_##S(const T *vals, size_t n, T id, mesh_block_op_##S op) { \
59 if (n == 0 || n % 32 != 0 || n > 1024) return id; \
60 size_t warps = n / 32; \
61 T tot[32]; \
62 for (size_t w = 0; w < 32; w++) tot[w] = id; \
63 for (size_t w = 0; w < warps; w++) { \
64 T v[32]; \
65 for (size_t l = 0; l < 32; l++) v[l] = vals[32 * w + l]; \
66 mesh_block_warp_reduce_##S(v, op); \
67 tot[w] = v[0]; \
68 } \
69 mesh_block_warp_reduce_##S(tot, op); \
70 return tot[0]; \
71 } \
72 bool mesh_block_scan_##S(const T *vals, size_t n, T id, mesh_block_op_##S op, bool inclusive, \
73 T *out) { \
74 if (n == 0 || n % 32 != 0 || n > 1024) return false; \
75 size_t warps = n / 32; \
76 T local[32][32], tot[32]; \
77 for (size_t w = 0; w < 32; w++) tot[w] = id; \
78 for (size_t w = 0; w < warps; w++) { \
79 for (size_t l = 0; l < 32; l++) local[w][l] = vals[32 * w + l]; \
80 mesh_block_warp_scan_##S(local[w], op); \
81 tot[w] = local[w][31]; \
82 } \
83 mesh_block_warp_scan_##S(tot, op); \
84 for (size_t w = 0; w < warps; w++) { \
85 T prefix = tot[(w + 31) % 32], inc[32]; \
86 for (size_t l = 0; l < 32; l++) inc[l] = w == 0 ? local[w][l] : op(prefix, local[w][l]); \
87 for (size_t l = 0; l < 32; l++) \
88 out[32 * w + l] = inclusive ? inc[l] : l > 0 ? inc[l - 1] : w == 0 ? id : prefix; \
89 } \
90 return true; \
91 }
92
93MESH_BLOCK_DEFINE(float, f32)
94MESH_BLOCK_DEFINE(uint32_t, u32)
95MESH_BLOCK_DEFINE(uint64_t, u64)