enable f16;

#define DECLARE_BYTE_LOADERS_SRC0
#include "common_decls.tmpl"

#include "mul_mat_decls.tmpl"

#ifdef VEC
fn store_val(acc: array<array<f16, TILE_M>, TILE_N>, tn: u32, tm: u32) -> vec4<f32> {
    return vec4<f32>(f32(acc[tn][tm]), f32(acc[tn][tm + 1]), f32(acc[tn][tm + 2]), f32(acc[tn][tm + 3]));
}
#endif

#ifdef SCALAR
fn store_val(acc: array<array<f16, TILE_M>, TILE_N>, tn: u32, tm: u32) -> f32 {
    return f32(acc[tn][tm]);
}
#endif

struct MulMatIdParams {
    offset_src0: u32,
    offset_src1: u32,
    offset_dst: u32,

    k: u32,
    m: u32,
    n_expert: u32,
    n_expert_used: u32,
    n_tokens: u32,
    b_ne1: u32,

    stride_01: u32,
    stride_11: u32,
    stride_02: u32,
    stride_12: u32,
};

@group(0) @binding(0) var<storage, read_write> src0: array<SRC0_TYPE>; // [cols, rows, n_expert]
@group(0) @binding(1) var<storage, read_write> src1: array<SRC1_TYPE>; // [cols, b_ne1, n_tokens]
@group(0) @binding(2) var<storage, read_write> dst: array<DST_TYPE>;   // [rows, n_expert_used, n_tokens]
@group(0) @binding(3) var<storage, read_write> global_gathered_expert_used: array<u32>; // [n_expert][n_tokens]
@group(0) @binding(4) var<storage, read_write> global_gathered_tokens: array<u32>; // [n_expert][n_tokens]
@group(0) @binding(5) var<storage, read_write> gathered_count_ids: array<u32>; // [n_expert]

@group(0) @binding(6) var<uniform> params: MulMatIdParams;

fn get_local_n(thread_id: u32) -> u32 {
    return thread_id / WORKGROUP_SIZE_M;
}
fn get_local_m(thread_id: u32) -> u32 {
    return thread_id % WORKGROUP_SIZE_M;
}

const TOTAL_WORKGROUP_SIZE = WORKGROUP_SIZE_M * WORKGROUP_SIZE_N;
const TILE_SRC0_SHMEM = TILE_K * WORKGROUP_SIZE_M * TILE_M;
const TILE_SRC1_SHMEM = TILE_K * WORKGROUP_SIZE_N * TILE_N;

var<workgroup> shmem: array<f16, TILE_SRC0_SHMEM + TILE_SRC1_SHMEM>;
var<workgroup> gathered_expert_used: array<u32, TILE_N * WORKGROUP_SIZE_N>;
var<workgroup> gathered_tokens: array<u32, TILE_N * WORKGROUP_SIZE_N>;

#ifdef INIT_SRC1_SHMEM_FLOAT
fn init_shmem_id_src1(thread_id: u32, offset_src1: u32, rest_token_n: u32, k_outer: u32) {
    for (var elem_idx = thread_id * VEC_SIZE; elem_idx < TILE_SRC1_SHMEM; elem_idx += TOTAL_WORKGROUP_SIZE * VEC_SIZE) {
        let tile_n = elem_idx / TILE_K;
        let tile_k = elem_idx % TILE_K;
        if (tile_n < rest_token_n) {
            let global_src10 = k_outer + tile_k;
            let expert_used_idx = gathered_expert_used[tile_n] % params.b_ne1;
            let token_idx = gathered_tokens[tile_n];
            let src1_idx = offset_src1 + token_idx * params.stride_12 + expert_used_idx * params.stride_11 + global_src10;
            let src1_val = select(
                SRC1_TYPE(0.0),
                src1[src1_idx/VEC_SIZE],
                global_src10 < params.k);
            store_shmem(SHMEM_TYPE(src1_val), TILE_SRC0_SHMEM + elem_idx);
        } else {
            store_shmem(SHMEM_TYPE(0.0), TILE_SRC0_SHMEM + elem_idx);
        }
    }
}
#endif // INIT_SRC1_SHMEM_FLOAT

@compute @workgroup_size(TOTAL_WORKGROUP_SIZE)
fn main(@builtin(workgroup_id) wg_id: vec3<u32>,
        @builtin(local_invocation_id) local_id: vec3<u32>,
        @builtin(num_workgroups) num_wg: vec3<u32>) {

    let thread_id = local_id.x;
    let local_m = get_local_m(thread_id);
    let local_n = get_local_n(thread_id);

    var expert_idx:u32 = 0xFFFFFFFFu;
    var wg_in_batch:u32 = 0;
    var wg_sum:u32 = 0;
    let wg_m_count = (params.m + WORKGROUP_SIZE_M * TILE_M - 1u) / (WORKGROUP_SIZE_M * TILE_M);
    let wg_linear = wg_id.y * num_wg.x + wg_id.x;

    for (var i = 0u;i < params.n_expert;i += 1) {
        let wg_n_count = (gathered_count_ids[i] + WORKGROUP_SIZE_N * TILE_N - 1u) / (WORKGROUP_SIZE_N * TILE_N);
        let wg_per_matrix = wg_m_count * wg_n_count;
        if (wg_sum <= wg_linear && wg_linear < wg_sum + wg_per_matrix) {
            expert_idx = i;
            wg_in_batch = wg_linear - wg_sum;
            break;
        }
        wg_sum += wg_per_matrix;
    }

    let is_valid = expert_idx != 0xFFFFFFFFu;

    var wg_m: u32 = 0;
    var wg_n: u32 = 0;
    var offset_wg_m: u32 = 0;
    var offset_wg_n: u32 = 0;
    var rest_token_n: u32 = 0;
    var src0_batch_offset: u32 = 0;

    wg_m = wg_in_batch % wg_m_count;
    wg_n = wg_in_batch / wg_m_count;

    offset_wg_m = wg_m * WORKGROUP_SIZE_M * TILE_M;
    offset_wg_n = wg_n * WORKGROUP_SIZE_N * TILE_N;

    if (is_valid) {
        rest_token_n = gathered_count_ids[expert_idx] - offset_wg_n;
        let global_gathered_base = expert_idx * params.n_tokens + offset_wg_n;
        for (var i = thread_id; i < TILE_N * WORKGROUP_SIZE_N && offset_wg_n + i < gathered_count_ids[expert_idx]; i += TOTAL_WORKGROUP_SIZE) {
            gathered_expert_used[i] = global_gathered_expert_used[global_gathered_base + i];
            gathered_tokens[i] = global_gathered_tokens[global_gathered_base + i];
        }
        src0_batch_offset = params.offset_src0 + expert_idx * params.stride_02;
    }

    workgroupBarrier();

    let output_row_base = offset_wg_m + local_m * TILE_M;
    let output_col_base = offset_wg_n + local_n * TILE_N;

    let dst2_stride = params.m * params.n_expert_used;
    let dst1_stride = params.m;

    var acc: array<array<f16, TILE_M>, TILE_N>;

    for (var k_outer = 0u; k_outer < params.k; k_outer += TILE_K) {

        if (is_valid) {
            init_shmem_src0(thread_id, src0_batch_offset, offset_wg_m, k_outer);
            init_shmem_id_src1(thread_id, params.offset_src1, rest_token_n, k_outer);
        }

        workgroupBarrier();

        if (is_valid) {
            let k_end = min(TILE_K, params.k - k_outer);

            for (var k_inner = 0u; k_inner < k_end; k_inner++) {
                var src0_tile: array<f16, TILE_M>;
                for (var tm = 0u; tm < TILE_M; tm++) {
                    let src0_m = local_m * TILE_M + tm;
                    let src0_idx = k_inner + src0_m * TILE_K;
                    src0_tile[tm] = shmem[src0_idx];
                }
                for (var tn = 0u; tn < TILE_N; tn++) {
                    let src1_n = local_n * TILE_N + tn;
                    let src1_idx = src1_n * TILE_K + k_inner;
                    let src1_val = shmem[TILE_SRC0_SHMEM + src1_idx];
                    for (var tm = 0u; tm < TILE_M; tm++) {
                        acc[tn][tm] += src0_tile[tm] * src1_val;
                    }
                }
            }
        }

        workgroupBarrier();
    }

    if (is_valid) {
        for (var tn = 0u; tn < TILE_N; tn++) {
            let n_idx = output_col_base + tn;
            if (n_idx < gathered_count_ids[expert_idx]) {
                let dst1_idx = gathered_expert_used[n_idx - offset_wg_n];
                let dst2_idx = gathered_tokens[n_idx - offset_wg_n];
                let dst12_offset = params.offset_dst + dst2_idx * dst2_stride + dst1_idx * dst1_stride;
                for (var tm = 0u; tm < TILE_M; tm += VEC_SIZE) {
                    let global_row = output_row_base + tm;
                    if (global_row < params.m) {
                        let dst_idx = dst12_offset + global_row;
                        dst[dst_idx/VEC_SIZE] = store_val(acc, tn, tm);
                    }
                }
            }
        }
    }
}