// sample.wgsl

/**
 * GPU-Side Sampling Kernel
 *
 * Performs temperature scaling, top-k selection, softmax, and sampling
 * entirely on GPU. Only reads back the single selected token ID.
 *
 * Reduces readback from 1MB (256K vocab × 4 bytes) to 4 bytes.
 *
 * Algorithm:
 * 1. Temperature scaling: logits = logits / temperature
 * 2. Parallel top-k: Each workgroup finds local top-k, then merge
 * 3. Softmax on top-k candidates
 * 4. Multinomial sampling with provided random value
 */

// Configuration
override WORKGROUP_SIZE: u32 = 256u;
const MAX_WORKGROUP_SIZE: u32 = 256u;
const MAX_TOP_K: u32 = 128u;  // Max top-k supported
const NEG_INF: f32 = -3.402823e+38;

struct Uniforms {
    vocab_size: u32,
    top_k: u32,
    temperature: f32,
    random_value: f32,  // Pre-generated random [0,1) for sampling
    pad_token_id: u32,
    logit_softcap: f32,  // Gemma 2: 30.0, 0.0 = disabled
    output_index: u32,   // Index into output token buffer
    pad0: u32,
}

// Apply softcapping: softcap * tanh(x / softcap)
// Returns x unchanged if softcap <= 0
fn apply_softcap(x: f32, softcap: f32) -> f32 {
    if (softcap <= 0.0) {
        return x;
    }
    return softcap * tanh(x / softcap);
}

fn candidate_beats(candidate_value: f32, candidate_index: u32, best_value: f32, best_index: u32) -> bool {
    if (candidate_value > best_value) {
        return true;
    }
    if (candidate_value < best_value) {
        return false;
    }
    return candidate_index < best_index;
}

@group(0) @binding(0) var<uniform> u: Uniforms;
@group(0) @binding(1) var<storage, read> logits: array<f32>;              // [vocabSize]
@group(0) @binding(2) var<storage, read_write> output: array<u32>;         // [N] - selected tokens
@group(0) @binding(3) var<storage, read_write> topk_indices: array<u32>;    // [topK] - intermediate
@group(0) @binding(4) var<storage, read_write> topk_logits: array<f32>;     // [topK] - intermediate

// Shared memory for workgroup-level reduction
var<workgroup> shared_values: array<f32, MAX_WORKGROUP_SIZE>;
var<workgroup> shared_indices: array<u32, MAX_WORKGROUP_SIZE>;

// Phase 1: Find local max in each workgroup for parallel top-k
// Each thread scans a chunk of vocabulary, keeps local top element
@compute @workgroup_size(WORKGROUP_SIZE, 1, 1)
fn find_topk_phase1(
    @builtin(global_invocation_id) gid: vec3<u32>,
    @builtin(local_invocation_id) lid: vec3<u32>,
    @builtin(workgroup_id) wgid: vec3<u32>,
    @builtin(num_workgroups) num_wg: vec3<u32>
) {
    let thread_idx = lid.x;
    let global_idx = gid.x;
    let vocab_size = u.vocab_size;
    let temperature = u.temperature;
    let pad_id = u.pad_token_id;
    let softcap = u.logit_softcap;

    // Single workgroup: write all logits directly for exact top-k
    if (num_wg.x == 1u) {
        var val: f32 = NEG_INF;
        if (global_idx < vocab_size && global_idx != pad_id) {
            val = apply_softcap(logits[global_idx], softcap) / temperature;
        }
        topk_logits[thread_idx] = val;
        topk_indices[thread_idx] = global_idx;
        return;
    }

    // Each thread finds max in its assigned range
    var local_max: f32 = NEG_INF;  // -FLT_MAX
    var local_max_idx: u32 = 0u;

    // Stride through vocabulary
    var idx = global_idx;
    while (idx < vocab_size) {
        if (idx != pad_id) {
            // Apply softcapping before temperature scaling
            let val = apply_softcap(logits[idx], softcap) / temperature;
            if (candidate_beats(val, idx, local_max, local_max_idx)) {
                local_max = val;
                local_max_idx = idx;
            }
        }
        idx = idx + WORKGROUP_SIZE * num_wg.x;
    }

    shared_values[thread_idx] = local_max;
    shared_indices[thread_idx] = local_max_idx;
    workgroupBarrier();

    // Reduce within workgroup to find workgroup's top value
    var stride = WORKGROUP_SIZE / 2u;
    while (stride > 0u) {
        if (thread_idx < stride) {
            if (candidate_beats(
                shared_values[thread_idx + stride],
                shared_indices[thread_idx + stride],
                shared_values[thread_idx],
                shared_indices[thread_idx]
            )) {
                shared_values[thread_idx] = shared_values[thread_idx + stride];
                shared_indices[thread_idx] = shared_indices[thread_idx + stride];
            }
        }
        workgroupBarrier();
        stride = stride / 2u;
    }

    // Thread 0 writes workgroup result
    if (thread_idx == 0u) {
        topk_logits[wgid.x] = shared_values[0];
        topk_indices[wgid.x] = shared_indices[0];
    }
}

// Phase 2: Merge workgroup results and select final top-k
// Single workgroup sorts and selects top-k from workgroup results
@compute @workgroup_size(WORKGROUP_SIZE, 1, 1)
fn find_topk_phase2(
    @builtin(local_invocation_id) lid: vec3<u32>
) {
    let thread_idx = lid.x;
    let top_k = u.top_k;
    let num_groups = min(WORKGROUP_SIZE, (u.vocab_size + WORKGROUP_SIZE - 1u) / WORKGROUP_SIZE);
    let num_candidates = select(num_groups, min(u.vocab_size, WORKGROUP_SIZE), num_groups == 1u);

    // Load workgroup results into shared memory
    // Assume <= WORKGROUP_SIZE workgroups from phase 1
    if (thread_idx < WORKGROUP_SIZE) {
        if (thread_idx < num_candidates) {
            shared_values[thread_idx] = topk_logits[thread_idx];
            shared_indices[thread_idx] = topk_indices[thread_idx];
        } else {
            shared_values[thread_idx] = NEG_INF;
            shared_indices[thread_idx] = 0u;
        }
    }
    workgroupBarrier();

    // Thread 0 does partial selection sort for top-k
    if (thread_idx == 0u) {
        for (var k: u32 = 0u; k < top_k && k < num_candidates; k = k + 1u) {
            var max_idx = k;
            var max_val = shared_values[k];

            for (var i: u32 = k + 1u; i < num_candidates; i = i + 1u) {
                if (candidate_beats(shared_values[i], shared_indices[i], max_val, shared_indices[max_idx])) {
                    max_val = shared_values[i];
                    max_idx = i;
                }
            }

            if (max_idx != k) {
                let tmp_val = shared_values[k];
                let tmp_idx = shared_indices[k];
                shared_values[k] = shared_values[max_idx];
                shared_indices[k] = shared_indices[max_idx];
                shared_values[max_idx] = tmp_val;
                shared_indices[max_idx] = tmp_idx;
            }
        }

        // Write sorted top-k back
        for (var k: u32 = 0u; k < top_k; k = k + 1u) {
            topk_logits[k] = shared_values[k];
            topk_indices[k] = shared_indices[k];
        }
    }
}

// Phase 3: Softmax on top-k and sample
@compute @workgroup_size(WORKGROUP_SIZE, 1, 1)
fn softmax_and_sample(
    @builtin(local_invocation_id) lid: vec3<u32>
) {
    let thread_idx = lid.x;
    let top_k = u.top_k;
    let random_val = u.random_value;

    // Load top-k logits
    if (thread_idx < top_k) {
        shared_values[thread_idx] = topk_logits[thread_idx];
        shared_indices[thread_idx] = topk_indices[thread_idx];
    }
    workgroupBarrier();

    // Thread 0 does softmax and sampling
    if (thread_idx == 0u) {
        // Find max for numerical stability
        var max_val: f32 = shared_values[0];
        for (var i: u32 = 1u; i < top_k; i = i + 1u) {
            max_val = max(max_val, shared_values[i]);
        }

        // Compute exp and sum
        var exp_sum: f32 = 0.0;
        for (var i: u32 = 0u; i < top_k; i = i + 1u) {
            let exp_val = exp(shared_values[i] - max_val);
            shared_values[i] = exp_val;
            exp_sum = exp_sum + exp_val;
        }

        // Normalize to probabilities and sample
        let inv_sum = 1.0 / exp_sum;
        var cum_prob: f32 = 0.0;
        var selected_token: u32 = shared_indices[top_k - 1u];  // Default to last

        for (var i: u32 = 0u; i < top_k; i = i + 1u) {
            let prob = shared_values[i] * inv_sum;
            cum_prob = cum_prob + prob;
            if (cum_prob >= random_val) {
                selected_token = shared_indices[i];
                break;
            }
        }

        output[u.output_index] = selected_token;
    }
}

// Combined single-pass version for smaller vocabularies (<= 65536)
// Uses hierarchical reduction within single kernel
@compute @workgroup_size(WORKGROUP_SIZE, 1, 1)
fn sample_single_pass(
    @builtin(global_invocation_id) gid: vec3<u32>,
    @builtin(local_invocation_id) lid: vec3<u32>,
    @builtin(num_workgroups) num_wg: vec3<u32>
) {
    let thread_idx = lid.x;
    let vocab_size = u.vocab_size;
    let top_k = min(u.top_k, MAX_TOP_K);
    let temperature = u.temperature;
    let random_val = u.random_value;
    let pad_id = u.pad_token_id;
    let softcap = u.logit_softcap;

    // Phase 1: Find global max
    var local_max: f32 = NEG_INF;
    var local_max_idx: u32 = 0u;

    var idx = gid.x;
    while (idx < vocab_size) {
        if (idx != pad_id) {
            // Apply softcapping before temperature scaling
            let val = apply_softcap(logits[idx], softcap) / temperature;
            if (candidate_beats(val, idx, local_max, local_max_idx)) {
                local_max = val;
                local_max_idx = idx;
            }
        }
        idx = idx + num_wg.x * WORKGROUP_SIZE;
    }

    shared_values[thread_idx] = local_max;
    shared_indices[thread_idx] = local_max_idx;
    workgroupBarrier();

    // Reduce to find workgroup max
    var stride = WORKGROUP_SIZE / 2u;
    while (stride > 0u) {
        if (thread_idx < stride) {
            if (candidate_beats(
                shared_values[thread_idx + stride],
                shared_indices[thread_idx + stride],
                shared_values[thread_idx],
                shared_indices[thread_idx]
            )) {
                shared_values[thread_idx] = shared_values[thread_idx + stride];
                shared_indices[thread_idx] = shared_indices[thread_idx + stride];
            }
        }
        workgroupBarrier();
        stride = stride / 2u;
    }

    // For single workgroup, thread 0 can do everything
    if (thread_idx == 0u && num_wg.x == 1u) {
        // We have top-1, but need top-k
        // For small vocab, just do the full selection
        // This simplified version selects top-1 only (greedy)
        // Full top-k sampling requires multi-pass for large vocab

        output[u.output_index] = shared_indices[0];
    }
}

// Greedy argmax for deterministic decoding (temperature=0 equivalent)
@compute @workgroup_size(WORKGROUP_SIZE, 1, 1)
fn argmax(
    @builtin(global_invocation_id) gid: vec3<u32>,
    @builtin(local_invocation_id) lid: vec3<u32>,
    @builtin(workgroup_id) wgid: vec3<u32>,
    @builtin(num_workgroups) num_wg: vec3<u32>
) {
    let thread_idx = lid.x;
    let global_idx = gid.x;
    let vocab_size = u.vocab_size;
    let pad_id = u.pad_token_id;
    let softcap = u.logit_softcap;

    // Each thread finds max in its chunk
    var local_max: f32 = NEG_INF;
    var local_max_idx: u32 = 0u;

    var idx = global_idx;
    while (idx < vocab_size) {
        if (idx != pad_id) {
            // Apply softcapping (argmax is greedy, no temperature)
            let val = apply_softcap(logits[idx], softcap);
            if (candidate_beats(val, idx, local_max, local_max_idx)) {
                local_max = val;
                local_max_idx = idx;
            }
        }
        idx = idx + num_wg.x * WORKGROUP_SIZE;
    }

    shared_values[thread_idx] = local_max;
    shared_indices[thread_idx] = local_max_idx;
    workgroupBarrier();

    // Reduce within workgroup
    var stride = WORKGROUP_SIZE / 2u;
    while (stride > 0u) {
        if (thread_idx < stride) {
            if (candidate_beats(
                shared_values[thread_idx + stride],
                shared_indices[thread_idx + stride],
                shared_values[thread_idx],
                shared_indices[thread_idx]
            )) {
                shared_values[thread_idx] = shared_values[thread_idx + stride];
                shared_indices[thread_idx] = shared_indices[thread_idx + stride];
            }
        }
        workgroupBarrier();
        stride = stride / 2u;
    }

    // Write workgroup result to global memory
    if (thread_idx == 0u) {
        topk_logits[wgid.x] = shared_values[0];
        topk_indices[wgid.x] = shared_indices[0];
    }
}

// Final reduction for argmax across workgroups
@compute @workgroup_size(WORKGROUP_SIZE, 1, 1)
fn argmax_reduce(
    @builtin(local_invocation_id) lid: vec3<u32>
) {
    let thread_idx = lid.x;
    let num_groups = min(WORKGROUP_SIZE, (u.vocab_size + WORKGROUP_SIZE - 1u) / WORKGROUP_SIZE);

    // Load workgroup maxes (up to WORKGROUP_SIZE)
    if (thread_idx < num_groups) {
        shared_values[thread_idx] = topk_logits[thread_idx];
        shared_indices[thread_idx] = topk_indices[thread_idx];
    } else {
        shared_values[thread_idx] = NEG_INF;
        shared_indices[thread_idx] = 0u;
    }
    workgroupBarrier();

    // Reduce
    var stride = WORKGROUP_SIZE / 2u;
    while (stride > 0u) {
        if (thread_idx < stride) {
            if (candidate_beats(
                shared_values[thread_idx + stride],
                shared_indices[thread_idx + stride],
                shared_values[thread_idx],
                shared_indices[thread_idx]
            )) {
                shared_values[thread_idx] = shared_values[thread_idx + stride];
                shared_indices[thread_idx] = shared_indices[thread_idx + stride];
            }
        }
        workgroupBarrier();
        stride = stride / 2u;
    }

    if (thread_idx == 0u) {
        output[u.output_index] = shared_indices[0];
    }
}