// gather_split4_f16_vec4_f16_out.wgsl

/**
 * Gather Kernel (split4 F16 input -> F16 output, vec4)
 *
 * Row-major embedding lookup for logical embedding tables split across up to
 * four GPU buffers to stay under maxStorageBufferBindingSize.
 */

enable f16;

override WORKGROUP_SIZE_VEC4: u32 = 64u;

struct Uniforms {
    num_tokens: u32,
    hidden_size: u32,
    vocab_size: u32,
    index_offset: u32,
    input_hidden_size: u32,
    hidden_offset: u32,
    section0_rows: u32,
    section1_rows: u32,
    section2_rows: u32,
    section3_rows: u32,
    _pad0: u32,
    _pad1: u32,
}

@group(0) @binding(0) var<uniform> u: Uniforms;
@group(0) @binding(1) var<storage, read> indices: array<u32>;
@group(0) @binding(2) var<storage, read> embeddings0: array<f16>;
@group(0) @binding(3) var<storage, read> embeddings1: array<f16>;
@group(0) @binding(4) var<storage, read> embeddings2: array<f16>;
@group(0) @binding(5) var<storage, read> embeddings3: array<f16>;
@group(0) @binding(6) var<storage, read_write> output_f16: array<f16>;

fn write_zero(out_base: u32) {
    output_f16[out_base] = f16(0.0);
    output_f16[out_base + 1u] = f16(0.0);
    output_f16[out_base + 2u] = f16(0.0);
    output_f16[out_base + 3u] = f16(0.0);
}

@compute @workgroup_size(WORKGROUP_SIZE_VEC4, 1, 1)
fn gather_vec4_f16_out(@builtin(global_invocation_id) gid: vec3<u32>) {
    let tid = gid.x;
    let vec4_per_row = u.hidden_size / 4u;
    let total_vec4s = u.num_tokens * vec4_per_row;

    if (tid >= total_vec4s) {
        return;
    }

    let token_idx = tid / vec4_per_row;
    let vec4_idx = tid % vec4_per_row;
    let token_id = indices[token_idx + u.index_offset];
    let out_base = tid * 4u;

    if (token_id >= u.vocab_size) {
        write_zero(out_base);
        return;
    }

    let section1_start = u.section0_rows;
    let section2_start = section1_start + u.section1_rows;
    let section3_start = section2_start + u.section2_rows;
    let section_end = section3_start + u.section3_rows;
    let dim_base = u.hidden_offset + vec4_idx * 4u;

    var local_token: u32;
    var embed_base: u32;
    if (token_id < section1_start) {
        local_token = token_id;
        embed_base = local_token * u.input_hidden_size + dim_base;
        output_f16[out_base] = embeddings0[embed_base];
        output_f16[out_base + 1u] = embeddings0[embed_base + 1u];
        output_f16[out_base + 2u] = embeddings0[embed_base + 2u];
        output_f16[out_base + 3u] = embeddings0[embed_base + 3u];
    } else if (token_id < section2_start) {
        local_token = token_id - section1_start;
        embed_base = local_token * u.input_hidden_size + dim_base;
        output_f16[out_base] = embeddings1[embed_base];
        output_f16[out_base + 1u] = embeddings1[embed_base + 1u];
        output_f16[out_base + 2u] = embeddings1[embed_base + 2u];
        output_f16[out_base + 3u] = embeddings1[embed_base + 3u];
    } else if (token_id < section3_start) {
        local_token = token_id - section2_start;
        embed_base = local_token * u.input_hidden_size + dim_base;
        output_f16[out_base] = embeddings2[embed_base];
        output_f16[out_base + 1u] = embeddings2[embed_base + 1u];
        output_f16[out_base + 2u] = embeddings2[embed_base + 2u];
        output_f16[out_base + 3u] = embeddings2[embed_base + 3u];
    } else if (token_id < section_end) {
        local_token = token_id - section3_start;
        embed_base = local_token * u.input_hidden_size + dim_base;
        output_f16[out_base] = embeddings3[embed_base];
        output_f16[out_base + 1u] = embeddings3[embed_base + 1u];
        output_f16[out_base + 2u] = embeddings3[embed_base + 2u];
        output_f16[out_base + 3u] = embeddings3[embed_base + 3u];
    } else {
        write_zero(out_base);
    }
}