1/* @ajax/mesh-rt-prim: mesh-rt's device primitives for whichever target is being compiled:
2 * inline PTX on nvptx64 (@ajax/mesh-rt-ptx), the CPU simulator elsewhere
3 * (@ajax/mesh-rt-sim), plus the f32 helpers built on them.
4 *
5 * Ported from alpha-rust-mesh crates/mesh-rt/src/ptx.rs. */
6#ifndef MESH_RT_PRIM_H
7#define MESH_RT_PRIM_H
8
9#include <stdbool.h>
10#include <stdint.h>
11
12#ifndef MESH_RT_INLINE
13#define MESH_RT_INLINE static inline __attribute__((always_inline))
14#endif
15
16#if defined(__NVPTX__)
17#include "mesh/rt/ptx.h"
18#else
19#include "mesh/rt/sim.h"
20#endif
21
22/* f32 <-> its bits (Rust's to_bits / from_bits). */
23MESH_RT_INLINE uint32_t mesh_f32_bits(float f) {
24    union { float f; uint32_t u; } x;
25    x.f = f;
26    return x.u;
27}
28MESH_RT_INLINE float mesh_f32_from_bits(uint32_t u) {
29    union { uint32_t u; float f; } x;
30    x.u = u;
31    return x.f;
32}
33
34MESH_RT_INLINE float mesh_rt_rdf(const uint8_t *p) { return mesh_f32_from_bits(mesh_rt_rd(p)); }
35MESH_RT_INLINE void mesh_rt_wrf(uint8_t *p, float v) { mesh_rt_wr(p, mesh_f32_bits(v)); }
36MESH_RT_INLINE float mesh_rt_shfl_idx_f(float v, uint32_t src) { return mesh_f32_from_bits(mesh_rt_shfl_idx(mesh_f32_bits(v), src)); }
37
38/* Butterfly sum: lane l adds lane l ^ m for m = 16, 8, 4, 2, 1. The pair sums
39 * are commutative, so every lane ends with the bitwise-same value. */
40MESH_RT_INLINE float mesh_rt_warp_sum(float v) {
41    for (uint32_t m = 16; m > 0; m >>= 1) v += mesh_f32_from_bits(mesh_rt_shfl_bfly(mesh_f32_bits(v), m));
42    return v;
43}
44
45#endif