// fused_swiglu.wgsl

/**
 * SwiGLU Activation Kernel (Row-Split Fused Gate/Up)
 *
 * Computes: output = SiLU(gate + gate_bias) * (up + up_bias)
 *
 * input layout: [num_tokens, 2 * dim] flattened, row-major
 *   row = [gate[0..dim), up[0..dim)]
 * bias layout: [2 * dim]
 *   bias = [gate_bias[0..dim), up_bias[0..dim)]
 * output layout: [num_tokens, dim]
 */

override WORKGROUP_SIZE: u32 = 256u;

struct Uniforms {
    num_tokens: u32,
    dim: u32,
    bias_offset: u32,  // byte offset into bias buffer (divide by 4 for F32 index)
    clamp_max: f32,    // SwiGLU clamp (0 = disabled)
}

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

fn sigmoid(x: f32) -> f32 {
    return 1.0 / (1.0 + exp(-x));
}

fn silu(x: f32) -> f32 {
    return x * sigmoid(x);
}

fn clamp_swiglu(x: f32) -> f32 {
    if (u.clamp_max <= 0.0) {
        return x;
    }
    return clamp(x, -u.clamp_max, u.clamp_max);
}

@compute @workgroup_size(WORKGROUP_SIZE, 1, 1)
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
    let idx = gid.x;
    let total = u.num_tokens * u.dim;
    if (idx >= total) {
        return;
    }

    let token_idx = idx / u.dim;
    let dim_idx = idx % u.dim;

    // bias_offset is in bytes, convert to F32 index
    let bias_base = u.bias_offset / 4u;

    let row_base = token_idx * u.dim * 2u;
    let gate = input[row_base + dim_idx] + bias[bias_base + dim_idx];
    let up = input[row_base + u.dim + dim_idx] + bias[bias_base + u.dim + dim_idx];

    output[idx] = clamp_swiglu(silu(gate) * up);
}