1/* @ajax/mesh-ref-block: the block collectives of a multi-warp node (crates/mesh-rt/src/block.rs,
2 * mesh_rt.h) on the CPU, step for step: the same trees, the same operand order, every f32 op in Z.
3 * A GPU's result for thread values `vals` (thread t's value at index t, a multiple of 32 of them,
4 * at most 1024) is bitwise the one these return.
5 *
6 * Each collective comes for f32, u32 and u64 values with the operation as a function pointer.
7 * Freestanding (no libc). */
8#ifndef MESH_REF_BLOCK_H
9#define MESH_REF_BLOCK_H
10
11#include <stdbool.h>
12#include <stddef.h>
13#include <stdint.h>
14
15/* f32 sum's identity, -0.0, and max's, all bits set (the least value of the total order), as bits. */
16#define MESH_BLOCK_SUM_F32_ID_BITS 0x80000000u
17#define MESH_BLOCK_MAX_F32_ID_BITS 0xffffffffu
18float mesh_block_sum_f32_id(void);
19float mesh_block_max_f32_id(void);
20
21/* The larger in the IEEE total order of the bits (one of them, unchanged): mesh-rt `max_f32`. */
22float mesh_block_max_f32(float a, float b);
23/* add.rn.ftz.f32. */
24float mesh_block_add_f32(float a, float b);
25uint32_t mesh_block_add_u32(uint32_t a, uint32_t b); /* wrapping */
26uint64_t mesh_block_add_u64(uint64_t a, uint64_t b); /* wrapping */
27
28/* A 32-bit integer hash (lowbias32): the same on every backend. */
29uint32_t mesh_block_mix(uint32_t x);
30
31/* The test values programs/blocks gives thread t for seed s: an f32 of mixed signs and magnitudes
32 * (a few zeros, subnormals and large values), a u32 and a u64. */
33void mesh_block_test_values(uint32_t seed, uint32_t t, float *f, uint32_t *u, uint64_t *w);
34
35#define MESH_BLOCK_DECLARE(T, S)                                                                   \
36    typedef T (*mesh_block_op_##S)(T, T);                                                          \
37    /* The warp butterfly, in place: lane l := op(own, lane l ^ m) for m = 16, 8, 4, 2, 1. */       \
38    void mesh_block_warp_reduce_##S(T v[32], mesh_block_op_##S op);                                \
39    /* The warp's inclusive Hillis-Steele scan, in place: lane l >= d := op(lane l - d, own). */    \
40    void mesh_block_warp_scan_##S(T v[32], mesh_block_op_##S op);                                  \
41    /* A block reduction: each warp's butterfly, then the butterfly over the 32 warp results      \
42     * (`id` for the warps the node lacks). n: a multiple of 32, 32..1024 (else `id`). */          \
43    T mesh_block_reduce_##S(const T *vals, size_t n, T id, mesh_block_op_##S op);                  \
44    /* A block scan into out[0..n): inclusive, or exclusive (a warp's first thread gets `id` in    \
45     * warp 0, else the combined prefix). False (out untouched) unless n is valid as above. */     \
46    bool mesh_block_scan_##S(const T *vals, size_t n, T id, mesh_block_op_##S op, bool inclusive,  \
47                             T *out);
48
49MESH_BLOCK_DECLARE(float, f32)
50MESH_BLOCK_DECLARE(uint32_t, u32)
51MESH_BLOCK_DECLARE(uint64_t, u64)
52
53#endif