1/* The CPU implementation of mesh-rt's device primitives, and the launcher. */
2#if !defined(__NVPTX__)
3
4#include "mesh/rt/sim.h"
5#include "mesh/abi/base.h"
6
7#include <errno.h>
8#include <pthread.h>
9#include <stdatomic.h>
10#include <stdio.h>
11#include <stdlib.h>
12#include <string.h>
13#include <time.h>
14
15typedef struct {
16    pthread_barrier_t bar;
17    uint32_t xch[32];
18} sim_warp;
19
20typedef struct {
21    uint32_t ctaid;
22    uint32_t ntid;
23    uint32_t smem_bytes;
24    uint8_t *smem;
25    pthread_barrier_t bar;
26    sim_warp *warps;
27} sim_node;
28
29typedef struct {
30    sim_node *node;
31    uint32_t tid;
32    mesh_sim_kernel kernel;
33    uint8_t *mb;
34    uint32_t slot;
35    pthread_t thread;
36    bool started;
37} sim_thread;
38
39struct mesh_sim_launch {
40    uint32_t nodes;
41    uint32_t threads; /* per node */
42    sim_node *node;
43    sim_thread *thr;
44};
45
46static _Thread_local sim_thread *self;
47
48static sim_thread *me(void) {
49    if (!self) {
50        fprintf(stderr, "mesh-rt-sim: a device primitive was called outside a simulated thread\n");
51        abort();
52    }
53    return self;
54}
55
56static _Atomic uint32_t *a32(const uint8_t *p) { return (_Atomic uint32_t *)(void *)(uintptr_t)p; }
57static _Atomic uint64_t *a64(const uint8_t *p) { return (_Atomic uint64_t *)(void *)(uintptr_t)p; }
58
59uint32_t mesh_rt_smid(void) { return me()->node->ctaid; }
60uint32_t mesh_rt_tid(void) { return me()->tid; }
61uint32_t mesh_rt_lane(void) { return me()->tid & 31u; }
62uint32_t mesh_rt_dyn_smem(void) { return me()->node->smem_bytes; }
63uint32_t mesh_rt_ntid(void) { return me()->node->ntid; }
64uint32_t mesh_rt_ctaid(void) { return me()->node->ctaid; }
65
66uint64_t mesh_rt_gtimer(void) {
67    struct timespec ts;
68    clock_gettime(CLOCK_MONOTONIC, &ts);
69    return (uint64_t)ts.tv_sec * 1000000000u + (uint64_t)ts.tv_nsec;
70}
71
72uint32_t mesh_rt_ld_acquire(const uint8_t *p) { return atomic_load_explicit(a32(p), memory_order_acquire); }
73uint64_t mesh_rt_ld_acquire64(const uint8_t *p) { return atomic_load_explicit(a64(p), memory_order_acquire); }
74uint32_t mesh_rt_ld_relaxed(const uint8_t *p) { return atomic_load_explicit(a32(p), memory_order_relaxed); }
75uint64_t mesh_rt_ld_relaxed64(const uint8_t *p) { return atomic_load_explicit(a64(p), memory_order_relaxed); }
76void mesh_rt_st_release(uint8_t *p, uint32_t v) { atomic_store_explicit(a32(p), v, memory_order_release); }
77void mesh_rt_st_release64(uint8_t *p, uint64_t v) { atomic_store_explicit(a64(p), v, memory_order_release); }
78void mesh_rt_st_relaxed(uint8_t *p, uint32_t v) { atomic_store_explicit(a32(p), v, memory_order_relaxed); }
79uint32_t mesh_rt_atom_add(uint8_t *p, uint32_t v) { return atomic_fetch_add_explicit(a32(p), v, memory_order_relaxed); }
80uint32_t mesh_rt_atom_or(uint8_t *p, uint32_t v) { return atomic_fetch_or_explicit(a32(p), v, memory_order_relaxed); }
81uint32_t mesh_rt_atom_and(uint8_t *p, uint32_t v) { return atomic_fetch_and_explicit(a32(p), v, memory_order_relaxed); }
82uint32_t mesh_rt_atom_cas(uint8_t *p, uint32_t cmp, uint32_t desired) {
83    atomic_compare_exchange_strong_explicit(a32(p), &cmp, desired, memory_order_relaxed, memory_order_relaxed);
84    return cmp;
85}
86void mesh_rt_red_add(uint8_t *p, uint32_t v) { atomic_fetch_add_explicit(a32(p), v, memory_order_relaxed); }
87void mesh_rt_fence_sc(void) { atomic_thread_fence(memory_order_seq_cst); }
88void mesh_rt_fence_acq_rel(void) { atomic_thread_fence(memory_order_acq_rel); }
89
90void mesh_rt_nanosleep(uint32_t ns) {
91    struct timespec ts = {0, (long)ns};
92    while (nanosleep(&ts, &ts) != 0 && errno == EINTR) {
93    }
94}
95
96static void wait_bar(pthread_barrier_t *b) {
97    int r = pthread_barrier_wait(b);
98    if (r != 0 && r != PTHREAD_BARRIER_SERIAL_THREAD) {
99        fprintf(stderr, "mesh-rt-sim: pthread_barrier_wait: %s\n", strerror(r));
100        abort();
101    }
102}
103
104static sim_warp *my_warp(void) {
105    sim_thread *t = me();
106    return &t->node->warps[t->tid >> 5];
107}
108
109void mesh_rt_warp_sync(void) { wait_bar(&my_warp()->bar); }
110void mesh_rt_bar_sync(void) {
111    atomic_thread_fence(memory_order_seq_cst);
112    wait_bar(&me()->node->bar);
113    atomic_thread_fence(memory_order_seq_cst);
114}
115uint8_t *mesh_rt_smem_base(void) { return me()->node->smem; }
116
117/* Every lane deposits v; after the barrier each reads what it wants; the second
118 * barrier keeps the next exchange from overwriting a value not yet read. */
119static uint32_t exchange(uint32_t v, uint32_t (*pick)(const uint32_t *xch, uint32_t lane, uint32_t arg), uint32_t arg) {
120    sim_warp *w = my_warp();
121    uint32_t lane = me()->tid & 31u;
122    w->xch[lane] = v;
123    wait_bar(&w->bar);
124    uint32_t r = pick(w->xch, lane, arg);
125    wait_bar(&w->bar);
126    return r;
127}
128static uint32_t pick_idx(const uint32_t *x, uint32_t lane, uint32_t src) { (void)lane; return x[src & 31u]; }
129static uint32_t pick_bfly(const uint32_t *x, uint32_t lane, uint32_t m) { return x[(lane ^ m) & 31u]; }
130static uint32_t pick_up(const uint32_t *x, uint32_t lane, uint32_t d) { return lane >= d ? x[lane - d] : x[lane]; }
131static uint32_t pick_ballot(const uint32_t *x, uint32_t lane, uint32_t unused) {
132    (void)lane, (void)unused;
133    uint32_t r = 0;
134    for (uint32_t i = 0; i < 32; i++)
135        if (x[i]) r |= 1u << i;
136    return r;
137}
138
139uint32_t mesh_rt_shfl_idx(uint32_t v, uint32_t src) { return exchange(v, pick_idx, src); }
140uint32_t mesh_rt_shfl_bfly(uint32_t v, uint32_t m) { return exchange(v, pick_bfly, m); }
141uint32_t mesh_rt_shfl_up(uint32_t v, uint32_t d) { return exchange(v, pick_up, d); }
142uint32_t mesh_rt_ballot(bool p) { return exchange(p ? 1u : 0u, pick_ballot, 0); }
143
144_Noreturn void mesh_rt_trap(void) {
145    sim_thread *t = self;
146    fprintf(stderr, "mesh-rt-sim: trap (node %u, thread %u)\n", t ? t->node->ctaid : 0u, t ? t->tid : 0u);
147    abort();
148}
149
150/* f32 subnormal inputs and results flush to sign-preserving zero, as the
151 * cluster float mode's .ftz instructions do (README "Cluster float mode"). */
152static void ftz_on(void) {
153#if defined(__x86_64__) || defined(__i386__)
154    uint32_t csr;
155    __asm__ __volatile__("stmxcsr %0" : "=m"(csr));
156    csr |= 0x8040u; /* FTZ (bit 15) | DAZ (bit 6) */
157    __asm__ __volatile__("ldmxcsr %0" ::"m"(csr));
158#elif defined(__aarch64__)
159    uint64_t fpcr;
160    __asm__ __volatile__("mrs %0, fpcr" : "=r"(fpcr));
161    fpcr |= (uint64_t)1 << 24; /* FZ */
162    __asm__ __volatile__("msr fpcr, %0" ::"r"(fpcr));
163#endif
164}
165
166static void *run(void *arg) {
167    sim_thread *t = arg;
168    self = t;
169    ftz_on();
170    t->kernel(t->mb, t->slot);
171    self = NULL;
172    return NULL;
173}
174
175static void free_launch(mesh_sim_launch *l) {
176    for (uint32_t n = 0; n < l->nodes; n++) {
177        sim_node *nd = &l->node[n];
178        if (nd->warps) {
179            for (uint32_t w = 0; w < l->threads / 32; w++) pthread_barrier_destroy(&nd->warps[w].bar);
180            free(nd->warps);
181        }
182        pthread_barrier_destroy(&nd->bar);
183        free(nd->smem);
184    }
185    free(l->node);
186    free(l->thr);
187    free(l);
188}
189
190mesh_sim_launch *mesh_sim_start(mesh_sim_kernel kernel, uint8_t *mb, uint32_t slot, uint32_t nodes, uint32_t warps,
191                                uint32_t shared_bytes) {
192    if (nodes == 0 || warps == 0 || warps > MESH_MAX_WARPS) return NULL;
193    mesh_sim_launch *l = calloc(1, sizeof *l);
194    if (!l) return NULL;
195    l->nodes = nodes;
196    l->threads = 32 * warps;
197    l->node = calloc(nodes, sizeof *l->node);
198    l->thr = calloc((size_t)nodes * l->threads, sizeof *l->thr);
199    if (!l->node || !l->thr) { free(l->node); free(l->thr); free(l); return NULL; }
200    for (uint32_t n = 0; n < nodes; n++) {
201        sim_node *nd = &l->node[n];
202        nd->ctaid = n;
203        nd->ntid = l->threads;
204        nd->smem_bytes = MESH_RT_SHARED_BYTES + shared_bytes;
205        nd->smem = aligned_alloc(16, ((size_t)nd->smem_bytes + 15) & ~(size_t)15);
206        nd->warps = calloc(warps, sizeof *nd->warps);
207        if (!nd->smem || !nd->warps) { fprintf(stderr, "mesh-rt-sim: out of memory\n"); abort(); }
208        memset(nd->smem, 0, nd->smem_bytes);
209        pthread_barrier_init(&nd->bar, NULL, l->threads);
210        for (uint32_t w = 0; w < warps; w++) pthread_barrier_init(&nd->warps[w].bar, NULL, 32);
211    }
212    pthread_attr_t attr;
213    pthread_attr_init(&attr);
214    pthread_attr_setstacksize(&attr, (size_t)1 << 20);
215    for (uint32_t n = 0; n < nodes; n++) {
216        for (uint32_t i = 0; i < l->threads; i++) {
217            sim_thread *t = &l->thr[(size_t)n * l->threads + i];
218            *t = (sim_thread){.node = &l->node[n], .tid = i, .kernel = kernel, .mb = mb, .slot = slot};
219            int r = pthread_create(&t->thread, &attr, run, t);
220            if (r != 0) {
221                /* Threads already started would wait at a barrier forever. */
222                fprintf(stderr, "mesh-rt-sim: pthread_create: %s\n", strerror(r));
223                abort();
224            }
225            t->started = true;
226        }
227    }
228    pthread_attr_destroy(&attr);
229    return l;
230}
231
232void mesh_sim_join(mesh_sim_launch *l) {
233    if (!l) return;
234    for (size_t i = 0; i < (size_t)l->nodes * l->threads; i++)
235        if (l->thr[i].started) pthread_join(l->thr[i].thread, NULL);
236    free_launch(l);
237}
238
239#else
240/* nvptx64: nothing here (an empty translation unit is not ISO C). */
241typedef int mesh_rt_sim_unused;
242#endif