/*
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 merge silu operator
*/

#include "merge_silu.h"

namespace cpu_inference {

namespace {
    float16_t silu_table_fp16[1 << 16];
    static bool is_silu_table_fp16_init = false;
}

MergeSilu::MergeSilu(const std::vector<int>& numa_list){
    if(!is_silu_table_fp16_init){
        is_silu_table_fp16_init = true;
        for(int i = 0; i < (1 << 16); ++i) {
            float f = fp16_to_fp32(*(float16_t*)(&i));
            silu_table_fp16[i] = f / (1.0 + expf(-1*f));
        }
    }
    this->numa_list = numa_list;
}

void MergeSilu::forward(const WorkDivider& work, float16_t* output, void* input1, void* input2, 
        int token_cnt, int intermediate_size){
    int first_numa = numa_list[0];
    SingleNumaWorkRange srange;
    if (work.num_numas == 1) {
        // 此次为专家并行计算情况，每个numa都有对应的专家需要计算。
        divide_all_work(&work, token_cnt * intermediate_size, &srange);
    } else {
        // 这里是一般情况，只使用0号numa进行计算。
        divide_work_first_numa(&work, token_cnt * intermediate_size, &srange);
    }
    // 只在0号numa计算
    float16_t* w1_ptr = (float16_t*)input1;
    float16_t* w3_ptr = (float16_t*)input2;
    float16_t* out_ptr = (float16_t*)output;
    // work.my_numa == first_numa 是指一般情况，不涉及专家并行计算，只使用0号numa计算以及将结果存入0号numa。
    // work.num_numas == 1 为专家并行计算的情况，对于每个专家来说自己所处的numa都相当于在第一个numa上。
    if (work.my_numa == first_numa || work.num_numas == 1) {
        for (int item = srange.begin_thread; item < srange.end_thread; item++) {
            float16_t w1 = w1_ptr[item];       // w1部分
            float16_t w3 = w3_ptr[item];       // w3部分
            // SiLU(w1) * w3
            out_ptr[item] = silu_table_fp16[*(uint16_t *)&w1] * w3;
        }
    }
}

};