1/* @ajax/mesh-rt-ptx: the device primitives of mesh-rt on NVIDIA GPUs, as inline PTX.
2 *
3 * Special registers, system-scope memory ordering, atomics, warp primitives
4 * and the block barrier: everything the mesh shares with the host goes
5 * through these (or the volatile payload accessors, ordered by an
6 * acquire/release pair around them). Each asm has a "memory" clobber, as
7 * Rust's asm! without `nomem` is a compiler barrier.
8 *
9 * Only for nvptx64 (clang --target=nvptx64-nvidia-cuda); elsewhere this header
10 * is empty, and @ajax/mesh-rt-sim provides the same functions on the CPU.
11 * Ported from alpha-rust-mesh crates/mesh-rt/src/ptx.rs. */
12#ifndef MESH_RT_PTX_H
13#define MESH_RT_PTX_H
14
15#if defined(__NVPTX__)
16
17#include <stdbool.h>
18#include <stdint.h>
19
20#ifndef MESH_RT_INLINE
21#define MESH_RT_INLINE static inline __attribute__((always_inline))
22#endif
23
24/* The node's dynamic shared memory: RT_SHARED_BYTES for the runtime, then the
25 * program's `shared_bytes` (meshd launches every block with both). */
26extern __attribute__((address_space(3), aligned(16))) unsigned char mesh_smem[];
27
28MESH_RT_INLINE uint32_t mesh_rt_smid(void) { uint32_t r; __asm__ __volatile__("mov.u32 %0, %%smid;" : "=r"(r)::"memory"); return r; }
29/* The thread's index in its block (node): 0..32 x warps. */
30MESH_RT_INLINE uint32_t mesh_rt_tid(void) { uint32_t r; __asm__ __volatile__("mov.u32 %0, %%tid.x;" : "=r"(r)::"memory"); return r; }
31/* The thread's lane in its warp (0..32). */
32MESH_RT_INLINE uint32_t mesh_rt_lane(void) { return mesh_rt_tid() & 31u; }
33/* The launch's dynamic shared memory bytes (RT_SHARED_BYTES + the manifest's `shared_bytes`). */
34MESH_RT_INLINE uint32_t mesh_rt_dyn_smem(void) { uint32_t r; __asm__ __volatile__("mov.u32 %0, %%dynamic_smem_size;" : "=r"(r)::"memory"); return r; }
35/* Threads per block (32 x the manifest's `warps`). */
36MESH_RT_INLINE uint32_t mesh_rt_ntid(void) { uint32_t r; __asm__ __volatile__("mov.u32 %0, %%ntid.x;" : "=r"(r)::"memory"); return r; }
37MESH_RT_INLINE uint32_t mesh_rt_ctaid(void) { uint32_t r; __asm__ __volatile__("mov.u32 %0, %%ctaid.x;" : "=r"(r)::"memory"); return r; }
38MESH_RT_INLINE uint64_t mesh_rt_gtimer(void) { uint64_t r; __asm__ __volatile__("mov.u64 %0, %%globaltimer;" : "=l"(r)::"memory"); return r; }
39
40MESH_RT_INLINE uint32_t mesh_rt_ld_acquire(const uint8_t *p) {
41    uint32_t r; __asm__ __volatile__("ld.acquire.sys.global.u32 %0, [%1];" : "=r"(r) : "l"(p) : "memory"); return r;
42}
43MESH_RT_INLINE uint64_t mesh_rt_ld_acquire64(const uint8_t *p) {
44    uint64_t r; __asm__ __volatile__("ld.acquire.sys.global.u64 %0, [%1];" : "=l"(r) : "l"(p) : "memory"); return r;
45}
46MESH_RT_INLINE uint32_t mesh_rt_ld_relaxed(const uint8_t *p) {
47    uint32_t r; __asm__ __volatile__("ld.relaxed.sys.global.u32 %0, [%1];" : "=r"(r) : "l"(p) : "memory"); return r;
48}
49MESH_RT_INLINE uint64_t mesh_rt_ld_relaxed64(const uint8_t *p) {
50    uint64_t r; __asm__ __volatile__("ld.relaxed.sys.global.u64 %0, [%1];" : "=l"(r) : "l"(p) : "memory"); return r;
51}
52MESH_RT_INLINE void mesh_rt_st_release(uint8_t *p, uint32_t v) {
53    __asm__ __volatile__("st.release.sys.global.u32 [%0], %1;" ::"l"(p), "r"(v) : "memory");
54}
55MESH_RT_INLINE void mesh_rt_st_release64(uint8_t *p, uint64_t v) {
56    __asm__ __volatile__("st.release.sys.global.u64 [%0], %1;" ::"l"(p), "l"(v) : "memory");
57}
58MESH_RT_INLINE void mesh_rt_st_relaxed(uint8_t *p, uint32_t v) {
59    __asm__ __volatile__("st.relaxed.sys.global.u32 [%0], %1;" ::"l"(p), "r"(v) : "memory");
60}
61MESH_RT_INLINE uint32_t mesh_rt_atom_add(uint8_t *p, uint32_t v) {
62    uint32_t r; __asm__ __volatile__("atom.add.relaxed.sys.global.u32 %0, [%1], %2;" : "=r"(r) : "l"(p), "r"(v) : "memory"); return r;
63}
64MESH_RT_INLINE uint32_t mesh_rt_atom_or(uint8_t *p, uint32_t v) {
65    uint32_t r; __asm__ __volatile__("atom.or.relaxed.sys.global.b32 %0, [%1], %2;" : "=r"(r) : "l"(p), "r"(v) : "memory"); return r;
66}
67MESH_RT_INLINE uint32_t mesh_rt_atom_and(uint8_t *p, uint32_t v) {
68    uint32_t r; __asm__ __volatile__("atom.and.relaxed.sys.global.b32 %0, [%1], %2;" : "=r"(r) : "l"(p), "r"(v) : "memory"); return r;
69}
70MESH_RT_INLINE uint32_t mesh_rt_atom_cas(uint8_t *p, uint32_t cmp, uint32_t desired) {
71    uint32_t r;
72    __asm__ __volatile__("atom.cas.relaxed.sys.global.b32 %0, [%1], %2, %3;" : "=r"(r) : "l"(p), "r"(cmp), "r"(desired) : "memory");
73    return r;
74}
75MESH_RT_INLINE void mesh_rt_red_add(uint8_t *p, uint32_t v) {
76    __asm__ __volatile__("red.relaxed.sys.global.add.u32 [%0], %1;" ::"l"(p), "r"(v) : "memory");
77}
78MESH_RT_INLINE void mesh_rt_fence_sc(void) { __asm__ __volatile__("fence.sc.sys;" ::: "memory"); }
79MESH_RT_INLINE void mesh_rt_fence_acq_rel(void) { __asm__ __volatile__("fence.acq_rel.sys;" ::: "memory"); }
80MESH_RT_INLINE void mesh_rt_nanosleep(uint32_t ns) { __asm__ __volatile__("nanosleep.u32 %0;" ::"r"(ns) : "memory"); }
81MESH_RT_INLINE void mesh_rt_warp_sync(void) { __asm__ __volatile__("bar.warp.sync -1;" ::: "memory"); }
82/* The block barrier (`barrier.sync 0`): every thread of the node waits for
83 * every other, and memory operations before it are ordered before those after
84 * it for the whole block. Every thread must reach it. Not `bar.sync`
85 * (= barrier.sync.aligned, which needs every thread of a warp at the same
86 * instruction): inline asm is not convergent to LLVM, which may duplicate it
87 * into both arms of a branch, and an aligned barrier reached that way hangs. */
88MESH_RT_INLINE void mesh_rt_bar_sync(void) { __asm__ __volatile__("barrier.sync 0;" ::: "memory"); }
89/* The generic address of the node's dynamic shared memory. */
90MESH_RT_INLINE uint8_t *mesh_rt_smem_base(void) { return (uint8_t *)mesh_smem; }
91
92MESH_RT_INLINE uint32_t mesh_rt_shfl_idx(uint32_t v, uint32_t src) {
93    uint32_t r; __asm__ __volatile__("shfl.sync.idx.b32 %0, %1, %2, 31, -1;" : "=r"(r) : "r"(v), "r"(src) : "memory"); return r;
94}
95MESH_RT_INLINE uint32_t mesh_rt_shfl_bfly(uint32_t v, uint32_t m) {
96    uint32_t r; __asm__ __volatile__("shfl.sync.bfly.b32 %0, %1, %2, 31, -1;" : "=r"(r) : "r"(v), "r"(m) : "memory"); return r;
97}
98/* Lane l gets lane l - d's value (lanes below d keep their own). */
99MESH_RT_INLINE uint32_t mesh_rt_shfl_up(uint32_t v, uint32_t d) {
100    uint32_t r; __asm__ __volatile__("shfl.sync.up.b32 %0, %1, %2, 0, -1;" : "=r"(r) : "r"(v), "r"(d) : "memory"); return r;
101}
102/* Bits of the converged lanes whose predicate is true. */
103MESH_RT_INLINE uint32_t mesh_rt_ballot(bool p) {
104    uint32_t r;
105    __asm__ __volatile__("{ .reg .pred q; setp.ne.u32 q, %1, 0; vote.sync.ballot.b32 %0, q, -1; }" : "=r"(r) : "r"((uint32_t)p) : "memory");
106    return r;
107}
108
109/* Payload access, ordered by the acquire/release around it. */
110MESH_RT_INLINE uint32_t mesh_rt_rd(const uint8_t *p) { return *(const volatile uint32_t *)p; }
111MESH_RT_INLINE uint64_t mesh_rt_rd64(const uint8_t *p) { return *(const volatile uint64_t *)p; }
112MESH_RT_INLINE void mesh_rt_wr(uint8_t *p, uint32_t v) { *(volatile uint32_t *)p = v; }
113MESH_RT_INLINE void mesh_rt_wrb(uint8_t *p, uint8_t v) { *(volatile uint8_t *)p = v; }
114MESH_RT_INLINE void mesh_rt_wr64(uint8_t *p, uint64_t v) { *(volatile uint64_t *)p = v; }
115/* A panic, or a broken runtime invariant: faults the context (meshd sees the error). */
116MESH_RT_INLINE _Noreturn void mesh_rt_trap(void) {
117    __asm__ __volatile__("trap;" ::: "memory");
118    __builtin_unreachable();
119}
120
121#endif /* __NVPTX__ */
122#endif