// AUTO-GENERATED from src/gpu/kernels/gelu.wgsl.
// Edit the source kernel and src/gpu/kernels/codegen/wgsl-variants.js, then run `npm run kernels:codegen:sync`.
// GeLU Activation Kernel with F16 Input/Output
//
// F16 variant for reduced memory bandwidth when using F16 activations.

enable f16;

override WORKGROUP_SIZE: u32 = 256u;
override HAS_GATE: bool = false;
override USE_ROWSPLIT: bool = false;

struct Uniforms {
    size: u32,          // Total output elements
    rowsplit_dim: u32,  // Dim for rowsplit variants (0 when unused)
    _pad0: u32,
    _pad1: u32,
}

@group(0) @binding(0) var<uniform> u: Uniforms;
@group(0) @binding(1) var<storage, read> input: array<f16>;
@group(0) @binding(2) var<storage, read_write> output: array<f16>;
@group(0) @binding(3) var<storage, read> gate: array<f16>;

fn gelu(x: f16) -> f16 {
    let sqrt_2_over_pi = f16(0.7978845608);
    let c = f16(0.044715);
    let inner = sqrt_2_over_pi * (x + c * x * x * x);
    let inner_clamped = clamp(inner, f16(-15.0), f16(15.0));
    return f16(0.5) * x * (f16(1.0) + tanh(inner_clamped));
}

@compute @workgroup_size(WORKGROUP_SIZE, 1, 1)
fn main(
    @builtin(global_invocation_id) global_id: vec3<u32>
) {
    if (USE_ROWSPLIT) {
        if (u.rowsplit_dim == 0u) {
            return;
        }
        let dim = u.rowsplit_dim;
        let token_idx = global_id.y;
        let dim_idx = global_id.x;
        let idx = token_idx * dim + dim_idx;
        if (idx >= u.size || dim_idx >= dim) {
            return;
        }
        let row_base = token_idx * dim * 2u;
        let g = input[row_base + dim_idx];
        let up = input[row_base + dim + dim_idx];
        output[idx] = gelu(g) * up;
        return;
    }

    let idx = global_id.x;
    if (idx >= u.size) {
        return;
    }

    if (HAS_GATE) {
        let up = input[idx];
        let g = gate[idx];
        output[idx] = gelu(g) * up;
        return;
    }

    let x = input[idx];
    output[idx] = gelu(x);
}