f64 ln/exp emit invalid GLSL.std.450 OpExtInst (Log/Exp) on SPIR-V backend

[shimwell]

Many thanks for making this software, it looks like the best option for a project I am planning and I am super keen to give it a go but hit a bit of a road block for my specific project

Calling x.ln() or x.exp() on f64 inside a #[cube] function compiles to OpExtInst Log / OpExtInst Exp with f64 operands. That's invalid per the GLSL.std.450 extended-instruction spec:

The operand x must be a scalar or vector whose component type is 16-bit or 32-bit floating-point.

NVIDIA's Vulkan driver happens to accept it and returns roughly the right values. Mesa RADV+ACO correctly rejects it the pipeline "compiles" but emits

ACO ERROR:
    Unimplemented NIR instr bit size: div 64    %27 = fexp2 %26
ACO ERROR:
    Unimplemented NIR instr bit size: div 64    %31 = flog2 %24

and the kernel returns garbage (exp(1e-6) → NaN, ln(e) → 0.6534…).

Confirmed with the Mesa team in https://gitlab.freedesktop.org/mesa/mesa/-/work_items/15359 they consider their behaviour correct, and the fix has to live wherever the SPIR-V is produced.

f64::sqrt is fine — no such restriction in the spec, and it works at 0 ULP drift on the same driver.

Same shape of codegen bug as #1200 (Metal MSL, bf16 transcendentals invalid). Different backend, different type, same underlying issue: codegen ignores per-backend operand-type restrictions on transcendentals.

Repro

```rust
#[cube(launch_unchecked)]
fn ln_kernel(input: &Array, output: &mut Array) {
if ABSOLUTE_POS >= input.len() { terminate!(); }
output[ABSOLUTE_POS] = input[ABSOLUTE_POS].ln();
}
```

Run on a Vulkan adapter exposing SHADER_F64 with cubecl-wgpu's spirv feature. AMD gfx1151 + Mesa 25.2.8 reproduces 100%. Full repro at https://github.com/fusion-neutronics/yamc, branch `develop-vulkan`, `crates/yamc-gpu/src/kernels/f64_transcendentals.rs`.

Suggested fix

When the SPIR-V backend lowers Float::ln / Float::exp for an f64 type, either:

• emit a polyfill (extract exponent + mantissa, polynomial approximant) so the resulting SPIR-V uses only spec-valid ops, or
• error at compile time so callers polyfill in their own kernels rather than getting silent NVIDIA-only correctness.

Probably worth auditing the other GLSL.std.450 ops with the same operand restriction (Sin, Cos, Pow, Tan, Asin, etc.) for the same lowering bug.

Versions: `cubecl 0.10.0-pre.4`, `cubecl-wgpu 0.10.0-pre.4` with `spirv` feature.

Can attach the SPIR-V dump if useful.

[shimwell]

another data point for this issue on the same backend.

f64 sin and cos have the same issue. Adding x.sin() or x.cos() to a #[cube] function on the SPIR-V backend produces:

Unimplemented NIR instr bit size: div 64    %26 = fsin %24
Unimplemented NIR instr bit size: div 64    %31 = fcos_amd %26

and the kernel returns garbage (~-2.4e+220 instead of sin(0) = 0).

GLSL.std.450 has the same f16/f32-only operand restriction on Sin, Cos, Tan, Asin, Acos, Atan, Sinh, Cosh, Tanh, Pow etc, so whatever the fix shape ends up being for Log/Exp it'd cover all of them.

Worked around this in our project by using Marsaglia's rejection method on the unit disk to get (cos φ, sin φ) without ever calling sin/cos. For ln/exp our existing polyfills cover us. So we're unblocked on AMD/RADV, just leaving a note in case it's helpful when someone gets to this.

Repro: same yamc-gpu branch, but the f64 sin/cos probe:

cargo test -p yamc-gpu kernels::f64_sin_cos -- --include-ignored --nocapture

Source at https://github.com/fusion-neutronics/yamc/blob/develop-vulkan/crates/yamc-gpu/src/kernels/f64_sin_cos.rs.

[shimwell]

Follow-up: f64 sin / cos are also broken on NVIDIA Vulkan, not just AMD/RADV. The previous comment flagged that all GLSL.std.450 transcendentals with the f16/f32 operand restriction are likely affected, confirming Sin/Cos from a different driver in case it helps scope the fix.

Adapter: NVIDIA RTX A500 Laptop GPU (DiscreteGpu, vendor 0x10de), Vulkan, cubecl 0.10.0 + cubecl-wgpu spirv feature. Same garbage value as the AMD/RADV case — every input maps to ~1.4e-289:

cos: max ULP drift = u64::MAX (input[5] = π, GPU ≈ 1.36e-289, CPU = -1)
sin: max ULP drift = u64::MAX (input[6] = -π/2, GPU ≈ 1.36e-289, CPU = -1)

So unlike Log/Exp where NVIDIA "happens to accept it and returns roughly the right values," for Sin/Cos NVIDIA produces the same ~1.4e-289 as RADV. That pattern hints the bug is on the SPIR-V emission side regardless of driver, consistent with the suggested fix shape.

Worth noting because the Marsaglia rejection-sampling workaround in the earlier comment isn't always enough: in our case the Watt fission spectrum sampler called .cos() directly to construct cos²(π/2 · ξ), and with cos ≈ 0 the Watt distribution collapses to a single-exponential. On actinide verification spheres this produced 3–10× reaction-rate inflation on every nontrivial-σ inelastic MT, plus i64 fixed-point tally accumulator overflow that read back as negative tallies on (n,3n) / fission / heating. Non-fissile nuclides (Fe56 / Cu63 etc) didn't trip the bug because their inelastic-only code paths didn't enter the cos call.

Workaround. Taylor-series cos_f64 / sin_f64 polyfills with k·(π/2) range reduction, in the same shape as the existing ln_f64 / exp_f64: . ~1 ulp accurate on inputs not at a zero crossing

Versions: cubecl 0.10.0, cubecl-wgpu 0.10.0 with spirv feature.