// residual_vec4.wgsl

/**
 * Residual Add Kernel (vec4)
 *
 * Vectorized element-wise addition for residual connections.
 */

struct Uniforms {
    size: u32,     // Total number of elements
    scale: f32,    // Output scale
    _pad1: u32,
    _pad2: u32,
}

override WORKGROUP_SIZE_VEC4: u32 = 64u;

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

// Vectorized version for better throughput
@compute @workgroup_size(WORKGROUP_SIZE_VEC4, 1, 1)
fn add_vec4(@builtin(global_invocation_id) gid: vec3<u32>) {
    let dispatch_stride = max(u._pad1, 4u);
    let idx = gid.y * dispatch_stride + gid.x * 4u;
    let size = u.size;

    if (idx >= size) {
        return;
    }

    // Handle up to 4 elements at a time
    let remaining = min(4u, size - idx);

    if (remaining >= 4u) {
        output[idx] = (a[idx] + b[idx]) * u.scale;
        output[idx + 1u] = (a[idx + 1u] + b[idx + 1u]) * u.scale;
        output[idx + 2u] = (a[idx + 2u] + b[idx + 2u]) * u.scale;
        output[idx + 3u] = (a[idx + 3u] + b[idx + 3u]) * u.scale;
    } else {
        for (var i = 0u; i < remaining; i = i + 1u) {
            output[idx + i] = (a[idx + i] + b[idx + i]) * u.scale;
        }
    }
}