/*
Copyright (c) 2025-2025 Huawei Technologies Co., Ltd.

sysHAX-adapter is licensed under Mulan PSL v2.
You can use this software according to the terms and conditions of the Mulan PSL v2.
You may obtain a copy of Mulan PSL v2 at:
    http://license.coscl.org.cn/MulanPSL2
THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY OR FIT FOR A PARTICULAR
PURPOSE.
See the Mulan PSL v2 for more details.
Created: 2026-1-31
Desc: CPU inference qwen3 moe model
*/

#include "qwen3_moe.h"

namespace cpu_inference {

namespace {
    float16_t exp_table_fp16[1 << 16];
    static bool is_exp_table_fp16_init = false;
}

Qwen3MoeLayer::Qwen3MoeLayer(std::vector<int> numa_list, int layer_id, std::string weight_quant,
        std::string input_quant, int nrc, const Qwen3MoeConfig& config){
    if(!is_exp_table_fp16_init){
        is_exp_table_fp16_init = true;
        for(int i = 0; i < (1 << 16); ++i) {
            float f = (float)(*(float16_t*)(&i));
            exp_table_fp16[i] = (float16_t)(expf(f));
        }
    }
    this->mergeSilu = new MergeSilu(numa_list);
    this->numa_list = numa_list;
    this->layer_id = layer_id;
    this->weight_quant = weight_quant;
    this->input_quant = input_quant;
    this->nrc = nrc;
    this->config = config;
    if(weight_quant == "fp16" && input_quant == "fp16"){
        this->matMul = new MatMulF16(nrc);
    }
    else if (weight_quant == "q8_0" && input_quant == "q8_0"){
        this->matMul = new MatMulQ8(nrc);
    }
    else if  (weight_quant == "q4_0" && input_quant == "q8_0"){
        this->matMul = new MatMulQ4Q8(nrc);
    }
    else if  (weight_quant == "q8align" && input_quant == "q8align"){
        this->matMul = new MatMulQ8Align(nrc);
    }
    else {
        throw std::runtime_error(std::string("not supported matmul type"));
    }
    if(input_quant == "fp16"){
        this->elem_cnt_in_block = QuantizationBase<float16_t>::elem_cnt_in_block;
        this->block_byte_size = QuantizationBase<float16_t>::block_byte_size;
    }
    else if(input_quant == "q8_0"){
        this->elem_cnt_in_block = QuantizationBase<block_q8_0>::elem_cnt_in_block;
        this->block_byte_size = QuantizationBase<block_q8_0>::block_byte_size;
    }
    else if(input_quant == "q8align"){
        this->elem_cnt_in_block = Q8Align::elem_cnt_in_block;
        this->block_byte_size = Q8Align::block_byte_size;
    }
    else {
        throw std::runtime_error(std::string("not supported input quant type") + input_quant);
    }
}

Qwen3MoeLayer::~Qwen3MoeLayer(){
    if(this->matMul){
        delete this->matMul;
    }
    int numa_cnt = numa_list.size();
    if (gate_weight) numa_free(gate_weight, numa_cnt * sizeof(void*));
    if (up_weight)   numa_free(up_weight,   numa_cnt * sizeof(void*));
    if (down_weight) numa_free(down_weight, numa_cnt * sizeof(void*));
}
void Qwen3MoeLayer::load_model(ModelWeightBase* model_weight){
    std::string gate_name = std::to_string(this->layer_id) + "_" + "gate_proj.weight";
    std::string up_name = std::to_string(this->layer_id) + "_" + "up_proj.weight";
    std::string down_name = std::to_string(this->layer_id) + "_" + "down_proj.weight";
    int numa_cnt = numa_list.size();
    this->gate_weight = (void**)numa_alloc_onnode(numa_cnt * sizeof(void*), 0);
    this->up_weight   = (void**)numa_alloc_onnode(numa_cnt * sizeof(void*), 0);
    this->down_weight = (void**)numa_alloc_onnode(numa_cnt * sizeof(void*), 0);
    for (int i = 0; i < numa_cnt; i++){
        this->gate_weight[i] = model_weight->get_weight(gate_name).memory_map.at(i);
        this->up_weight[i] = model_weight->get_weight(up_name).memory_map.at(i);
        this->down_weight[i] = model_weight->get_weight(down_name).memory_map.at(i);
    }
}

void Qwen3MoeLayer::softmax_topk(
    const WorkDivider* work,
    std::vector<std::vector<int>>& expert_tokens,
    std::vector<std::vector<float>>& expert_weights_list,
    std::vector<int>& experts,
    float16_t* router_logits,
    int token_cnt,
    int* max_idx_record, 
    f32* max_val_record){
    MultiNumaWorkRange multiRange;
    divide_work_all_numas(work, token_cnt, &multiRange);
    int num_experts = this->config.num_experts;
    int top_k = this->config.num_experts_per_tok;
    for (int token_idx = multiRange.begin_numa + multiRange.begin_thread; token_idx < multiRange.begin_numa + multiRange.end_thread; token_idx++) {
        int* max_idx_arr = (int*)max_idx_record + token_idx * top_k;
        f32* max_val_arr = (f32*)max_val_record + token_idx * top_k;
        // 获取top-k专家索引
        for (int k = 0; k < top_k; k++) {
            max_idx_arr[k] = -1;
            max_val_arr[k] = -FLT_MAX;
        }
        for (int j = 0; j < num_experts; j++) {
            f32 val = (f32)(router_logits[token_idx * num_experts + j]);
            if(val > max_val_arr[0]){
                max_idx_arr[0] = j;
                max_val_arr[0] = val;
            }
            else {
                continue;
            }
            int heap_idx = 0;
            while(1){
                int left_child_idx = (heap_idx << 1) + 1;
                int right_child_idx = (heap_idx << 1) + 2;
                int swap_idx = -1;
                if(right_child_idx < top_k){
                    if(max_val_arr[left_child_idx] < max_val_arr[right_child_idx]){
                        swap_idx = left_child_idx;
                    }
                    else{
                        swap_idx = right_child_idx;
                    }
                }
                else if(left_child_idx < top_k){
                    swap_idx = left_child_idx;
                }
                else {
                    break;
                }
                if(max_val_arr[swap_idx] < max_val_arr[heap_idx]){
                    std::swap(max_idx_arr[swap_idx], max_idx_arr[heap_idx]);
                    std::swap(max_val_arr[swap_idx], max_val_arr[heap_idx]);
                    heap_idx = swap_idx;
                }
                else{
                    break;
                }
            }
        }
        f32 max_val = max_val_arr[top_k - 1];
        for(int i = 0;i < top_k;i++){
            if(max_val_arr[i] > max_val){
                max_val = max_val_arr[i];
            }
        }
        float max_val_sum = 0;
        for(int i = 0;i < top_k;i++){
            float16_t diff = ((float16_t)(max_val_arr[i] - max_val));
            max_val_arr[i] = (float)exp_table_fp16[*(uint16_t *)&diff];
            max_val_sum += max_val_arr[i];
        }
        for(int i = 0;i < top_k;i++){
            max_val_arr[i] /= max_val_sum;
        }
    }

#pragma omp barrier

    if(work->global_tid == 0){
        for(int token_idx = 0;token_idx < token_cnt;token_idx++){
            for (int k = 0; k < top_k; k++) {
                int exp_idx = max_idx_record[token_idx * top_k + k];
                f32 exp_weight = max_val_record[token_idx * top_k + k];
                expert_tokens[exp_idx].push_back(token_idx);
                expert_weights_list[exp_idx].push_back(exp_weight);
            }
        }
        experts.resize(num_experts);
        for (int exp = 0; exp < num_experts; exp++){
            experts[exp] = exp;
        }
        sort(experts.begin(), experts.end(), [&](int a, int b) {
            return expert_tokens[a].size() > expert_tokens[b].size();
        });
    }
#pragma omp barrier

}

void Qwen3MoeLayer::forward(Qwen3MoeRunState* run_state, torch::Tensor& output, 
        const torch::Tensor& hidden_states, const torch::Tensor& router_logits){
    int64_t token_cnt = hidden_states.size(0);
    int first_numa = numa_list[0];
    int numa_cnt = numa_list.size();
    int top_k = this->config.num_experts_per_tok;
    int num_experts = this->config.num_experts;
    int hidden_size = this->config.hidden_size;
    int moe_intermediate_size = this->config.moe_intermediate_size;

    // 使用预分配的内存
    float16_t* expert_output = (float16_t*)output.data_ptr();
    memset(expert_output, 0, token_cnt * hidden_size * sizeof(float16_t));

    float16_t* hidden_states_ptr = (float16_t*)hidden_states.data_ptr();
    float16_t* router_logits_ptr = (float16_t*)router_logits.data_ptr();

    std::vector<std::vector<int>> expert_tokens(num_experts);
    std::vector<std::vector<f32>> expert_weights_list(num_experts);
    std::vector<int> experts;
    int max_idx_record[token_cnt * top_k];
    f32 max_val_record[token_cnt * top_k];

#pragma omp parallel
{
    WorkDivider work;
    init_work_divider(&work, numa_cnt);

    softmax_topk(&work, expert_tokens, expert_weights_list, experts, router_logits_ptr, token_cnt, max_idx_record, max_val_record);

    init_work_divider(&work, numa_cnt, top_k, true);
    for (int base = 0; base < experts.size(); base += top_k){
        if (expert_tokens[experts[base]].size() == 0) break;

        SingleNumaWorkRange srange;
        int I = work.global_tid / work.threads_per_numa;
        int exp = experts[base + I];
        int num_tokens = expert_tokens[exp].size();
        float16_t *expert_input = static_cast<float16_t*>(run_state->expert_input[I]);
        float16_t *expert_gate = static_cast<float16_t*>(run_state->expert_gate[I]);
        float16_t *expert_up = static_cast<float16_t*>(run_state->expert_up[I]);
        float16_t *expert_down = static_cast<float16_t*>(run_state->expert_down[I]);
        float16_t *expert_fused = static_cast<float16_t*>(run_state->expert_fused[I]);
        void** quant_expert_input = (void**)run_state->quant_expert_input[I];
        void** quant_expert_fused = (void**)run_state->quant_expert_fused[I];
        // 1. 提取专家的token隐藏状态
        divide_all_work(&work, hidden_size, &srange);
        if (num_tokens)
        for (int i = 0; i < num_tokens; i++) {
            int token_idx = expert_tokens[exp][i];
            memcpy(expert_input + i * hidden_size + srange.begin_thread, 
                hidden_states_ptr + token_idx * hidden_size + srange.begin_thread,
                srange.work_per_thread * sizeof(float16_t));
        }
        // 2. 计算专家上的投影
        std::string quant_type = this->input_quant;
        size_t total_elem_size = num_tokens * hidden_size;
        int dstBlockNum = this->elem_cnt_in_block;
        int dstBlockSize = this->block_byte_size;
#pragma omp barrier
        divide_all_work(&work, total_elem_size / dstBlockNum, &srange);
        if (num_tokens){
            if (quant_type == "q8align"){
                Q8Align::q8align_quantize((float16_t*)expert_input,
                    (int8_t*)((void**)quant_expert_input[work.my_numa])[0],
                    (float*)((void**)quant_expert_input[work.my_numa])[1],
                    srange.begin_thread, srange.work_per_thread, hidden_size / dstBlockNum, true);
            }else{
                Quantize((char *)quant_expert_input[work.my_numa] + srange.begin_thread * dstBlockSize,
                                            expert_input + srange.begin_thread * dstBlockNum,
                                            quant_type, srange.work_per_thread * dstBlockNum);
            }
        }

        if (num_tokens) matMul->matmul(&work, expert_gate, this->gate_weight, exp, 
                                        quant_expert_input, 0, num_tokens, hidden_size, moe_intermediate_size);

        if (num_tokens) matMul->matmul(&work, expert_up, this->up_weight, exp, 
                                        quant_expert_input, 0, num_tokens, hidden_size, moe_intermediate_size);

        // 3. silu激活函数 
        this->mergeSilu->forward(work, expert_fused, expert_gate,
             expert_up, num_tokens, moe_intermediate_size);

#pragma omp barrier
        total_elem_size = num_tokens * moe_intermediate_size;
        divide_all_work(&work, total_elem_size / dstBlockNum, &srange);
        if (num_tokens){
            if (quant_type == "q8align"){
                Q8Align::q8align_quantize((float16_t*)expert_fused, 
                    (int8_t*)((void**)quant_expert_fused[work.my_numa])[0], 
                    (float*)((void**)quant_expert_fused[work.my_numa])[1], 
                    srange.begin_thread, srange.work_per_thread, moe_intermediate_size / dstBlockNum, true);
            }else{
                Quantize((char *)quant_expert_fused[work.my_numa] + srange.begin_thread * dstBlockSize, 
                                    expert_fused + srange.begin_thread * dstBlockNum,
                                    quant_type, srange.work_per_thread * dstBlockNum);
            }
        }

        if (num_tokens) matMul->matmul(&work, expert_down, this->down_weight, exp, 
                                        quant_expert_fused, 0, num_tokens, moe_intermediate_size, hidden_size);

        // 5.加权聚合到最终输出
        divide_all_work(&work, num_tokens * hidden_size, &srange);
        if (num_tokens)
        for (int i = srange.begin_thread; i < srange.end_thread; i++){
            float16_t *output_ptr = expert_output + expert_tokens[exp][i / hidden_size] * hidden_size + (i % hidden_size);
            *output_ptr +=  expert_weights_list[exp][i / hidden_size] * expert_down[i];
        }
    }
}
}

Qwen3MoeModel::Qwen3MoeModel(std::vector<int> numa_list, std::string weight_quant, 
    std::string input_quant, int nrc, const Qwen3MoeConfig& config){
    this->numa_list = numa_list;
    this->weight_quant = weight_quant;
    this->input_quant = input_quant;
    this->config = config;
    for(int i = 0;i < config.num_hidden_layers;i++){
        Qwen3MoeLayer* layer = new Qwen3MoeLayer(numa_list, i, weight_quant, input_quant, 
            nrc, config);
        layers.push_back(layer);
    }
}

Qwen3MoeModel::~Qwen3MoeModel(){
    for(int i = 0;i < layers.size();i++){
        if(layers[i]){
            delete layers[i];
            layers[i] = nullptr;
        }
    }
}
void Qwen3MoeModel::load_model(ModelWeightBase* model_weight){
    for(const auto& layer_ptr: layers){
        layer_ptr->load_model(model_weight);
    }
}
Qwen3MoeRunState::Qwen3MoeRunState(
        const Qwen3MoeConfig& config, 
        const std::vector<int>& numa_list,
        const std::string& quant_type, 
        int token_cnt){
    if (token_cnt&1) token_cnt++;
    this->config = config;
    this->token_cnt = token_cnt;
    int hidden_size = config.hidden_size;
    int num_experts = config.num_experts;
    int moe_intermediate_size = config.moe_intermediate_size;
    int top_k = config.num_experts_per_tok;

    this->numa_cnt = numa_list.size();
    this->quant_type = quant_type;
    if (num_experts > 0) {
        std::vector<int>expert_numa_id;
        // 每个专家使用的numa编号，目前支持有1,2,4个numa的情况，由于目前是均分专家的，因此每个numa上的线程数量最好相同。
        for (int i = 0; i < top_k; i++){
            expert_numa_id.push_back(i / (top_k / numa_cnt));
        }
        // 专家计算缓冲区（按最大可能的专家token数分配）
        int max_expert_tokens = token_cnt; // 最坏情况下所有token都分配给一个专家
        this->expert_input = (void**)numa_alloc_onnode(top_k * sizeof(void**), 0);
        this->expert_gate = (void**)numa_alloc_onnode(top_k * sizeof(void**), 0);
        this->expert_up = (void**)numa_alloc_onnode(top_k * sizeof(void**), 0);
        this->expert_fused = (void**)numa_alloc_onnode(top_k * sizeof(void**), 0);
        this->expert_down = (void**)numa_alloc_onnode(top_k * sizeof(void**), 0);
        for (int i = 0; i < top_k; i++){
            this->expert_input[i] = (float16_t*)numa_alloc_onnode(max_expert_tokens * hidden_size * sizeof(float16_t), expert_numa_id[i]);
            this->expert_gate[i] = (float16_t*)numa_alloc_onnode(max_expert_tokens * moe_intermediate_size * sizeof(float16_t), expert_numa_id[i]);
            this->expert_up[i] = (float16_t*)numa_alloc_onnode(max_expert_tokens * moe_intermediate_size * sizeof(float16_t), expert_numa_id[i]);
            this->expert_fused[i] = (float16_t*)numa_alloc_onnode(max_expert_tokens * moe_intermediate_size * sizeof(float16_t), expert_numa_id[i]);
            this->expert_down[i] = (float16_t*)numa_alloc_onnode(max_expert_tokens * hidden_size * sizeof(float16_t), expert_numa_id[i]);
        }

        // 专家量化缓冲区
        if (quant_type == "q8align"){
            void**** quant_expert_input_ptr = (void****)numa_alloc_onnode(top_k * sizeof(void***), 0);
            void**** quant_expert_fused_ptr = (void****)numa_alloc_onnode(top_k * sizeof(void***), 0);
            for (int j = 0; j < top_k; j++){
                quant_expert_input_ptr[j] = (void***)numa_alloc_onnode(numa_cnt * sizeof(void**), expert_numa_id[j]);
                quant_expert_fused_ptr[j] = (void***)numa_alloc_onnode(numa_cnt * sizeof(void**), expert_numa_id[j]);
                
                for (int i = 0; i < numa_cnt; i++) {
                    quant_expert_input_ptr[j][i] = (void**)numa_alloc_onnode(2 * sizeof(void*), i);
                    quant_expert_input_ptr[j][i][0] = (void*)numa_alloc_onnode(max_expert_tokens * hidden_size * sizeof(int8_t), i);
                    quant_expert_input_ptr[j][i][1] = (void*)numa_alloc_onnode(max_expert_tokens * hidden_size / QK8_ALIGN * 2 * sizeof(float), i);

                    quant_expert_fused_ptr[j][i] = (void**)numa_alloc_onnode(2 * sizeof(void*), i);
                    quant_expert_fused_ptr[j][i][0] = (void*)numa_alloc_onnode(max_expert_tokens * moe_intermediate_size * sizeof(int8_t), i);
                    quant_expert_fused_ptr[j][i][1] = (void*)numa_alloc_onnode(max_expert_tokens * moe_intermediate_size / QK8_ALIGN * 2 * sizeof(float), i);
                }
            }
            this->quant_expert_input = (void***)quant_expert_input_ptr;
            this->quant_expert_fused = (void***)quant_expert_fused_ptr;
        }else{
            this->quant_expert_input = (void***)numa_alloc_onnode(top_k * sizeof(void***), 0);
            this->quant_expert_fused = (void***)numa_alloc_onnode(top_k * sizeof(void***), 0);
            for (int j = 0; j < top_k; j++){
                this->quant_expert_input[j] = (void**)numa_alloc_onnode(numa_cnt * sizeof(void**), expert_numa_id[j]);
                this->quant_expert_fused[j] = (void**)numa_alloc_onnode(numa_cnt * sizeof(void**), expert_numa_id[j]);
                for (int i = 0; i < numa_cnt; i++) {
                    this->quant_expert_input[j][i] = (void*)numa_alloc_onnode(max_expert_tokens * hidden_size * sizeof(f32), i);
                    this->quant_expert_fused[j][i] = (void*)numa_alloc_onnode(max_expert_tokens * moe_intermediate_size * sizeof(f32), i);
                }
            }
        }
    } 
}

Qwen3MoeRunState::~Qwen3MoeRunState(){
    // ========== 释放MoE缓冲区 ==========
    int hidden_size = config.hidden_size;
    int num_experts = config.num_experts;
    int moe_intermediate_size = config.moe_intermediate_size;
    int top_k = config.num_experts_per_tok;
    if (num_experts > 0) {
        for (int j = 0; j < top_k; j++){
            if (expert_input[j]) numa_free(expert_input[j], token_cnt * hidden_size * sizeof(float16_t));
            if (expert_gate[j]) numa_free(expert_gate[j], token_cnt * moe_intermediate_size* sizeof(float16_t)); 
            if (expert_up[j]) numa_free(expert_up[j], token_cnt * moe_intermediate_size * sizeof(float16_t)); 
            if (expert_fused[j]) numa_free(expert_fused[j], token_cnt * moe_intermediate_size * sizeof(float16_t)); 
            if (expert_down[j]) numa_free(expert_down[j], token_cnt * hidden_size * sizeof(float16_t));
            if(quant_type == "q8align"){
                void**** quant_expert_input_ptr = (void****)quant_expert_input;
                void**** quant_expert_fused_ptr = (void****)quant_expert_fused;
                if (quant_expert_input_ptr[j]) {
                    for (int i = 0; i < numa_cnt; i++) {
                        if (quant_expert_input_ptr[j][i]) {
                            numa_free(quant_expert_input_ptr[j][i][0], token_cnt * hidden_size * sizeof(int8_t));
                            numa_free(quant_expert_input_ptr[j][i][1], token_cnt * hidden_size / QK8_ALIGN * 2 * sizeof(float));
                        }
                        numa_free(quant_expert_input_ptr[j][i], 2 * sizeof(void**));
                    }
                    numa_free(quant_expert_input_ptr[j], numa_cnt * sizeof(void***));
                }
                
                if (quant_expert_fused_ptr[j]) {
                    for (int i = 0; i < numa_cnt; i++) {
                        if (quant_expert_fused_ptr[j][i]) {
                            numa_free(quant_expert_fused_ptr[j][i][0], token_cnt * moe_intermediate_size * sizeof(int8_t));
                            numa_free(quant_expert_fused_ptr[j][i][1], token_cnt * moe_intermediate_size / QK8_ALIGN * 2 * sizeof(float));
                        }
                        numa_free(quant_expert_fused_ptr[j][i], 2 * sizeof(void**));
                    }
                    numa_free(quant_expert_fused_ptr[j], numa_cnt * sizeof(void***));
                }
            }else{
                if (quant_expert_input[j]) {
                    for (int i = 0; i < numa_cnt; i++) {
                        if (quant_expert_input[j][i]) {
                            numa_free(quant_expert_input[j][i], token_cnt * hidden_size * sizeof(f32));
                        }
                    }
                    numa_free(quant_expert_input[j], numa_cnt * sizeof(void**));
                }
                if (quant_expert_fused[j]) {
                    for (int i = 0; i < numa_cnt; i++) {
                        if (quant_expert_fused[j][i]) {
                            numa_free(quant_expert_fused[j][i], token_cnt * moe_intermediate_size * sizeof(f32));
                        }
                    }
                    numa_free(quant_expert_fused[j], numa_cnt * sizeof(void**));
                }
            }
        }
        numa_free(quant_expert_input, top_k * sizeof(void***));
        numa_free(quant_expert_fused, top_k * sizeof(void***));
        numa_free(expert_input, top_k * sizeof(void**));
        numa_free(expert_gate, top_k * sizeof(void**));
        numa_free(expert_up, top_k * sizeof(void**));
        numa_free(expert_fused, top_k * sizeof(void**));
        numa_free(expert_down, top_k * sizeof(void**));
    }
}   

void Qwen3MoeModel::forward(torch::Tensor& expert_output, const torch::Tensor& hidden_states, 
    const torch::Tensor& router_logits, int64_t layer_id){
    if(layer_id == 0){
        if(this->run_state){
            delete this->run_state;
            this->run_state = nullptr;
        }
        int token_cnt = hidden_states.size(0);
        this->run_state = new Qwen3MoeRunState(this->config, this->numa_list, 
            this->input_quant, token_cnt);
    }
    this->layers[layer_id]->forward(this->run_state, expert_output, hidden_states, router_logits);
    if(layer_id + 1 == this->config.num_hidden_layers){
        delete this->run_state;
        this->run_state = nullptr;
    }
}

};