numr 0.5.2

High-performance numerical computing with multi-backend GPU acceleration (CPU/CUDA/WebGPU)
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41

// Fused elementwise WGSL shaders (F32 only)
// fused_mul_add: out = a * b + c
// fused_add_mul: out = (a + b) * c
// fused_mul_add_scalar: out = a * scale + bias

struct TernaryParams {
    numel: u32,
}

struct ScalarFmaParams {
    numel: u32,
    scale: f32,
    bias: f32,
    _pad: u32,
}

// ============================================================================
// Ternary ops: 3 inputs (a, b, c), 1 output
// ============================================================================

@group(0) @binding(0) var<storage, read_write> tern_a: array<f32>;
@group(0) @binding(1) var<storage, read_write> tern_b: array<f32>;
@group(0) @binding(2) var<storage, read_write> tern_c: array<f32>;
@group(0) @binding(3) var<storage, read_write> tern_out: array<f32>;
@group(0) @binding(4) var<uniform> tern_params: TernaryParams;

@compute @workgroup_size(256)
fn fused_mul_add_f32(@builtin(global_invocation_id) gid: vec3<u32>) {
    let idx = gid.x;
    if (idx < tern_params.numel) {
        tern_out[idx] = fma(tern_a[idx], tern_b[idx], tern_c[idx]);
    }
}

@compute @workgroup_size(256)
fn fused_add_mul_f32(@builtin(global_invocation_id) gid: vec3<u32>) {
    let idx = gid.x;
    if (idx < tern_params.numel) {
        tern_out[idx] = (tern_a[idx] + tern_b[idx]) * tern_c[idx];
    }
}