1/* blocks: the block collectives of a multi-warp node, for kcheck (program.toml).
2 * Every thread computes the same collectives in the same order; thread 0
3 * writes the node-wide results, every thread its own scans and shared word.
4 * Ported from alpha-rust-mesh programs/blocks/src/lib.rs. */
5#include "mesh/rt.h"
6
7#define BLOCKS_MAGIC 0x4b434c42u /* "BLCK" */
8#define BLOCKS_V_MAGIC 0x0u
9#define BLOCKS_V_TASKS 0x4u
10#define BLOCKS_R0 0x1000u
11#define BLOCKS_REGION 0x10000u
12#define BLOCKS_HDR 0x40u
13#define BLOCKS_PER 48u
14#define BLOCKS_OP_EARLY 2u
15
16/* mesh_ref::block::mix (lowbias32). */
17static uint32_t blocks_mix(uint32_t x) {
18    x ^= x >> 16;
19    x *= 0x7feb352du;
20    x ^= x >> 15;
21    x *= 0x846ca68bu;
22    return x ^ (x >> 16);
23}
24
25/* mesh_ref::block::test_values. */
26static void blocks_test_values(uint32_t seed, uint32_t t, float *f, uint32_t *u, uint64_t *l) {
27    uint32_t h = blocks_mix((seed * 0x9e3779b9u) ^ t);
28    uint32_t h2 = blocks_mix(h ^ 0x85ebca6bu);
29    uint32_t h3 = blocks_mix(h2 ^ 0xc2b2ae35u);
30    uint32_t bits;
31    switch (h & 15u) {
32    case 0: bits = (h & 16u) ? 0x80000000u : 0u; break;
33    case 1: bits = (h & 0x80000000u) | (h2 & 0x007fffffu) | 1u; break;
34    case 2: bits = (h & 0x80000000u) | (200u << 23) | (h2 & 0x007fffffu); break;
35    default: bits = (h & 0x80000000u) | ((112u + (h2 >> 28)) << 23) | (h3 & 0x007fffffu); break;
36    }
37    *f = mesh_f32_from_bits(bits);
38    *u = h2;
39    *l = (uint64_t)h3 << 32 | h;
40}
41
42static void blocks_init(const mesh_ctx *ctx) {
43    if (mesh_ctx_leader(ctx)) mesh_ctx_set_u32(ctx, BLOCKS_V_MAGIC, BLOCKS_MAGIC);
44}
45
46static uint32_t blocks_task(const mesh_ctx *ctx, uint32_t op, uint32_t arg) {
47    mesh_view v = mesh_ctx_view(ctx);
48    if (op == MESH_OP_INIT) {
49        blocks_init(ctx);
50        return MESH_R_NONE;
51    }
52    if (op == MESH_OP_TASK) {
53        uint32_t t = mesh_ctx_thread(ctx), n = mesh_ctx_threads(ctx);
54        float f;
55        uint32_t u;
56        uint64_t l;
57        blocks_test_values(arg, t, &f, &u, &l);
58        float sf = mesh_ctx_block_sum_f32(ctx, f), mf = mesh_ctx_block_max_f32(ctx, f);
59        uint32_t su = mesh_ctx_block_sum_u32(ctx, u), mu = mesh_ctx_block_max_u32(ctx, u);
60        uint64_t sl = mesh_ctx_block_sum_u64(ctx, l), ml = mesh_ctx_block_max_u64(ctx, l);
61        float fi = mesh_ctx_block_scan_f32(ctx, f, true), fe = mesh_ctx_block_scan_f32(ctx, f, false);
62        uint32_t ui = mesh_ctx_block_scan_u32(ctx, u, true), ue = mesh_ctx_block_scan_u32(ctx, u, false);
63        uint64_t li = mesh_ctx_block_scan_u64(ctx, l, true), le = mesh_ctx_block_scan_u64(ctx, l, false);
64        uint32_t b = mesh_ctx_block_broadcast_u32(ctx, u, arg % n);
65        /* Shared memory: every thread's f32 bits in, the mirrored thread's out. */
66        uint32_t back = 0;
67        mesh_shared sh;
68        if (mesh_ctx_shared(ctx, 0, n, 4, &sh)) {
69            mesh_shared_set_u32(sh, t, mesh_f32_bits(f));
70            mesh_ctx_block_barrier(ctx);
71            back = mesh_shared_get_u32(sh, n - 1 - t);
72        }
73        uint32_t r = BLOCKS_R0 + BLOCKS_REGION * (mesh_ctx_ticket(ctx) % 4);
74        uint32_t p = r + BLOCKS_HDR + BLOCKS_PER * t;
75        uint32_t mine[9] = {mesh_f32_bits(fi), mesh_f32_bits(fe), ui, ue, (uint32_t)li, (uint32_t)(li >> 32),
76                            (uint32_t)le, (uint32_t)(le >> 32), back};
77        for (uint32_t i = 0; i < 9; i++) mesh_view_set_u32(v, p + 4 * i, mine[i]);
78        if (mesh_ctx_leader(ctx)) {
79            uint32_t hdr[12] = {n, mesh_f32_bits(sf), mesh_f32_bits(mf), su, mu, (uint32_t)sl, (uint32_t)(sl >> 32),
80                                (uint32_t)ml, (uint32_t)(ml >> 32), b, mesh_ctx_node(ctx), mesh_ctx_warps(ctx)};
81            for (uint32_t i = 0; i < 12; i++) mesh_view_set_u32(v, r + 4 * i, hdr[i]);
82            mesh_view_set_u32(v, BLOCKS_V_TASKS, mesh_view_u32(v, BLOCKS_V_TASKS) + 1);
83            mesh_ctx_count(ctx, 0, 1);
84        }
85        return su;
86    }
87    if (op == BLOCKS_OP_EARLY) {
88        /* Warps 1.. leave the task at once; warp 0 sums its lanes (a warp collective only). */
89        if (mesh_ctx_warp(ctx) != 0) return 0;
90        float s = mesh_rt_warp_sum((float)mesh_ctx_lane(ctx));
91        if (mesh_ctx_leader(ctx)) mesh_ctx_count(ctx, 1, 1);
92        return (uint32_t)s;
93    }
94    return MESH_R_BAD_OP;
95}
96
97MESH_PROGRAM(blocks_init, blocks_task);