// Batched forward evaluation of a BytecodeTape on GPU.
//
// One compute thread per batch element. Each thread walks the tape sequentially,
// maintaining a private section of the values buffer.
// ── OpCode constants (must match OpCode repr(u8) discriminants) ──
const OP_INPUT: u32 = 0u;
const OP_CONST: u32 = 1u;
const OP_ADD: u32 = 2u;
const OP_SUB: u32 = 3u;
const OP_MUL: u32 = 4u;
const OP_DIV: u32 = 5u;
const OP_REM: u32 = 6u;
const OP_POWF: u32 = 7u;
const OP_ATAN2: u32 = 8u;
const OP_HYPOT: u32 = 9u;
const OP_MAX: u32 = 10u;
const OP_MIN: u32 = 11u;
const OP_NEG: u32 = 12u;
const OP_RECIP: u32 = 13u;
const OP_SQRT: u32 = 14u;
const OP_CBRT: u32 = 15u;
const OP_POWI: u32 = 16u;
const OP_EXP: u32 = 17u;
const OP_EXP2: u32 = 18u;
const OP_EXPM1: u32 = 19u;
const OP_LN: u32 = 20u;
const OP_LOG2: u32 = 21u;
const OP_LOG10: u32 = 22u;
const OP_LN1P: u32 = 23u;
const OP_SIN: u32 = 24u;
const OP_COS: u32 = 25u;
const OP_TAN: u32 = 26u;
const OP_ASIN: u32 = 27u;
const OP_ACOS: u32 = 28u;
const OP_ATAN: u32 = 29u;
const OP_SINH: u32 = 30u;
const OP_COSH: u32 = 31u;
const OP_TANH: u32 = 32u;
const OP_ASINH: u32 = 33u;
const OP_ACOSH: u32 = 34u;
const OP_ATANH: u32 = 35u;
const OP_ABS: u32 = 36u;
const OP_SIGNUM: u32 = 37u;
const OP_FLOOR: u32 = 38u;
const OP_CEIL: u32 = 39u;
const OP_ROUND: u32 = 40u;
const OP_TRUNC: u32 = 41u;
const OP_FRACT: u32 = 42u;
const UNUSED: u32 = 0xFFFFFFFFu;
// ── Tape data (bind group 0) ──
struct TapeMeta {
num_ops: u32,
num_inputs: u32,
num_variables: u32,
num_outputs: u32,
batch_size: u32,
_pad0: u32,
_pad1: u32,
_pad2: u32,
}
@group(0) @binding(0) var<storage, read> opcodes: array<u32>;
@group(0) @binding(1) var<storage, read> arg0: array<u32>;
@group(0) @binding(2) var<storage, read> arg1: array<u32>;
@group(0) @binding(3) var<storage, read> constants: array<f32>;
@group(0) @binding(4) var<uniform> tape_meta: TapeMeta;
@group(0) @binding(5) var<storage, read> output_indices: array<u32>;
// ── I/O buffers (bind group 1) ──
@group(1) @binding(0) var<storage, read> inputs: array<f32>;
@group(1) @binding(1) var<storage, read_write> values: array<f32>;
@group(1) @binding(2) var<storage, read_write> outputs: array<f32>;
// ── Manual implementations for functions not in WGSL ──
fn cbrt_f32(x: f32) -> f32 {
// cbrt(x) = sign(x) * |x|^(1/3)
let s = sign(x);
return s * pow(abs(x), 1.0 / 3.0);
}
fn expm1_f32(x: f32) -> f32 {
// Avoid catastrophic cancellation for small |x|
if abs(x) < 1e-4 {
return x + 0.5 * x * x;
}
return exp(x) - 1.0;
}
fn ln1p_f32(x: f32) -> f32 {
// Avoid catastrophic cancellation for small |x|
if abs(x) < 1e-4 {
return x - 0.5 * x * x;
}
return log(1.0 + x);
}
fn sinh_f32(x: f32) -> f32 {
return (exp(x) - exp(-x)) * 0.5;
}
fn cosh_f32(x: f32) -> f32 {
return (exp(x) + exp(-x)) * 0.5;
}
fn asinh_f32(x: f32) -> f32 {
// Use |x| to avoid catastrophic cancellation for large negative x:
// log(x + sqrt(x²+1)) ≈ log(0) when x << 0, but log(|x| + sqrt(x²+1)) is stable.
let a = abs(x);
// Large |x|: x² overflows f32 (|x| > ~1.8e19) → sqrt(Inf)=Inf → log(Inf)=Inf,
// whereas the true value is finite. Use the asymptotic form that never
// forms x²: asinh(x) = sign(x)·(log|x| + log(1 + sqrt(1 + 1/x²))). 1/x² → 0
// as x² overflows, so the result stays finite. (Accurate only for large a;
// small a keeps the direct form to avoid log|x| cancellation.)
if a > 1e9 {
let r = log(a) + log(1.0 + sqrt(1.0 + 1.0 / (a * a)));
return select(-r, r, x >= 0.0);
}
let r = log(a + sqrt(a * a + 1.0));
return select(-r, r, x >= 0.0);
}
fn acosh_f32(x: f32) -> f32 {
// acosh(x) = ln(x + sqrt(x² - 1)). Use factored (x-1)(x+1) under the
// sqrt to retain precision near x=1 (`x*x - 1` rounds away the ε²
// term in f32 for x = 1+ε). Matches kernels::acosh_deriv convention.
// Large x: (x-1)(x+1) overflows f32; use acosh(x) = log(x)+log(1+sqrt(1-1/x²)).
if x > 1e9 {
return log(x) + log(1.0 + sqrt(1.0 - 1.0 / (x * x)));
}
return log(x + sqrt((x - 1.0) * (x + 1.0)));
}
fn atanh_f32(x: f32) -> f32 {
// atanh(x) = 0.5 * ln((1+x)/(1-x))
return 0.5 * log((1.0 + x) / (1.0 - x));
}
fn hypot_f32(a: f32, b: f32) -> f32 {
// Factor out max magnitude to avoid overflow for large inputs.
let ax = abs(a);
let ay = abs(b);
let inf = bitcast<f32>(0x7f800000u);
// IEEE: hypot(±Inf, x) = +Inf for any x (including NaN). The
// rescaled formula would otherwise compute `Inf/Inf = NaN` when
// both operands are Inf, diverging from CPU f64::hypot and the
// CUDA `hypot` builtin.
if ax == inf || ay == inf { return inf; }
let mx = max(ax, ay);
let mn = min(ax, ay);
if mx == 0.0 { return 0.0; }
let r = mn / mx;
return mx * sqrt(1.0 + r * r);
}
fn rem_f32(a: f32, b: f32) -> f32 {
// Rust's % is remainder (truncated), matching: a - trunc(a/b) * b.
// WGSL has no exact fmod, so this is exact only while the quotient is
// exactly representable: |a/b| < 2^24 (f32 mantissa). Beyond that,
// trunc cannot recover the integer quotient and the result diverges
// from CPU/CUDA fmod (e.g. rem(1e8, 3) -> 0 instead of 1).
return a - trunc(a / b) * b;
}
fn recip_f32(x: f32) -> f32 {
return 1.0 / x;
}
fn log10_f32(x: f32) -> f32 {
return log(x) / log(10.0);
}
fn signum_f32(x: f32) -> f32 {
// Rust f32::signum: -1 for -0.0 (sign bit), +1 for +0.0/positive, NaN at NaN.
// `x >= 0.0` wrongly maps -0.0 to +1; inspect the sign bit. Bitcast NaN test
// since `x != x` is unreliable under Metal fast-math.
let b = bitcast<u32>(x);
if ((b & 0x7fffffffu) > 0x7f800000u) { return x; }
return select(1.0, -1.0, (b & 0x80000000u) != 0u);
}
fn is_nan_f32(x: f32) -> bool {
// NaN iff exponent is all-ones and mantissa is non-zero. Inspect the bits
// directly — `x != x` can be folded away by Metal's fast-math.
return (bitcast<u32>(x) & 0x7fffffffu) > 0x7f800000u;
}
fn powf_real(base: f32, b: f32) -> f32 {
// WGSL `pow(x, y)` is undefined for x < 0 (naga lowers it to
// `exp2(y*log2(x))`, and `log2(negative) = NaN`). Rust/C `powf` define
// x^y for x < 0 only when y is an integer: sign(x)^y * |x|^y. A
// non-integer exponent at a negative base is NaN — the same as on CPU.
// 0^0 = 1 (matches CPU/C `powf`); naga lowers `pow(0,0)` to
// `exp2(0*log2(0)) = exp2(NaN) = NaN`, so guard it explicitly.
if base == 0.0 && b == 0.0 { return 1.0; }
if base >= 0.0 { return pow(base, b); }
let rb = round(b);
if rb != b { return bitcast<f32>(0x7fc00000u); }
let mag = pow(abs(base), b);
if (i32(rb) & 1) != 0 { return -mag; }
return mag;
}
fn powi_f32(base: f32, exp_bits: u32) -> f32 {
// The exponent is stored as i32 reinterpreted as u32.
let n = bitcast<i32>(exp_bits);
return powf_real(base, f32(n));
}
// ── Main kernel ──
@compute @workgroup_size(256)
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
let batch_id = gid.x;
// Guard: skip threads beyond the batch size (last workgroup padding).
if batch_id >= tape_meta.batch_size {
return;
}
let num_vars = tape_meta.num_variables;
let num_in = tape_meta.num_inputs;
let num_ops = tape_meta.num_ops;
let n_out = tape_meta.num_outputs;
// Base offset into the per-thread values section.
let base = batch_id * num_vars;
// Initialize values: copy constants, then overwrite input slots from inputs buffer.
for (var i = 0u; i < num_vars; i = i + 1u) {
values[base + i] = constants[i];
}
// Overwrite input slots with this batch element's inputs.
let input_base = batch_id * num_in;
for (var i = 0u; i < num_in; i = i + 1u) {
values[base + i] = inputs[input_base + i];
}
// Walk the tape.
for (var i = num_in; i < num_ops; i = i + 1u) {
let op = opcodes[i];
// Skip Const entries — already initialized from constants buffer.
if op == OP_CONST {
continue;
}
let a_idx = arg0[i];
let b_idx = arg1[i];
let a = values[base + a_idx];
var b = 0.0f;
if b_idx != UNUSED {
b = values[base + b_idx];
}
var r = 0.0f;
switch op {
case 2u /* ADD */: { r = a + b; }
case 3u /* SUB */: { r = a - b; }
case 4u /* MUL */: { r = a * b; }
case 5u /* DIV */: { r = a / b; }
case 6u /* REM */: { r = rem_f32(a, b); }
case 7u /* POWF */: { r = powf_real(a, b); }
case 8u /* ATAN2 */: { r = atan2(a, b); }
case 9u /* HYPOT */: { r = hypot_f32(a, b); }
// Hand-roll the NaN tie-break to return the non-NaN operand,
// matching CPU `f32::max`/`min` and the tangent/Taylor kernels.
// The raw `max`/`min` builtin's NaN behaviour is backend-defined.
case 10u /* MAX */: { if (a >= b || is_nan_f32(b)) { r = a; } else { r = b; } }
case 11u /* MIN */: { if (a <= b || is_nan_f32(b)) { r = a; } else { r = b; } }
case 12u /* NEG */: { r = -a; }
case 13u /* RECIP */: { r = recip_f32(a); }
case 14u /* SQRT */: { r = sqrt(a); }
case 15u /* CBRT */: { r = cbrt_f32(a); }
case 16u /* POWI */: { r = powi_f32(a, b_idx); }
case 17u /* EXP */: { r = exp(a); }
case 18u /* EXP2 */: { r = exp2(a); }
case 19u /* EXPM1 */: { r = expm1_f32(a); }
case 20u /* LN */: { r = log(a); }
case 21u /* LOG2 */: { r = log2(a); }
case 22u /* LOG10 */: { r = log10_f32(a); }
case 23u /* LN1P */: { r = ln1p_f32(a); }
case 24u /* SIN */: { r = sin(a); }
case 25u /* COS */: { r = cos(a); }
case 26u /* TAN */: { r = tan(a); }
case 27u /* ASIN */: { r = asin(a); }
case 28u /* ACOS */: { r = acos(a); }
case 29u /* ATAN */: { r = atan(a); }
case 30u /* SINH */: { r = sinh_f32(a); }
case 31u /* COSH */: { r = cosh_f32(a); }
case 32u /* TANH */: { r = tanh(a); }
case 33u /* ASINH */: { r = asinh_f32(a); }
case 34u /* ACOSH */: { r = acosh_f32(a); }
case 35u /* ATANH */: { r = atanh_f32(a); }
case 36u /* ABS */: { r = abs(a); }
case 37u /* SIGNUM */: { r = signum_f32(a); }
case 38u /* FLOOR */: { r = floor(a); }
case 39u /* CEIL */: { r = ceil(a); }
case 40u /* ROUND */: { let t = trunc(a); r = select(t, t + select(-1.0, 1.0, a >= 0.0), abs(a - t) >= 0.5); }
case 41u /* TRUNC */: { r = trunc(a); }
// WGSL `fract` is floor-based; CPU `f32::fract()` is truncation-based.
case 42u /* FRACT */: { r = a - trunc(a); }
default: { r = 0.0; }
}
values[base + i] = r;
}
// Write outputs.
let out_base = batch_id * n_out;
for (var j = 0u; j < n_out; j = j + 1u) {
let oi = output_indices[j];
outputs[out_base + j] = values[base + oi];
}
}