/*
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 tp store method
*/

#ifndef TP_METHOD_H
#define TP_METHOD_H

#include <vector>
#include <string>
#include <regex>
#include <unordered_map>

namespace cpu_inference {

enum class StoreType {
    FULL_IN_FIRST_NUMA = 0,
    FULL_IN_ALL_NUMA = 1,
    ROW_SPLIT = 2,
};

inline const std::unordered_map<std::string, StoreType> qwen2_split_method = {
    {"mlp.down_proj.weight", StoreType::ROW_SPLIT},
    {"mlp.gate_proj.weight", StoreType::ROW_SPLIT},
    {"mlp.up_proj.weight", StoreType::ROW_SPLIT},
    {"self_attn.k_proj.weight", StoreType::ROW_SPLIT},
    {"self_attn.q_proj.weight", StoreType::ROW_SPLIT},
    {"self_attn.v_proj.weight", StoreType::ROW_SPLIT},
    {"self_attn.o_proj.weight", StoreType::ROW_SPLIT},
};

inline const std::unordered_map<std::string, StoreType> qwen3_moe_name_store = {
    { "gate_proj.weight", StoreType::FULL_IN_ALL_NUMA},
    { "up_proj.weight", StoreType::FULL_IN_ALL_NUMA},
    { "down_proj.weight", StoreType::FULL_IN_ALL_NUMA}
};

};  // namespace cpu_inference

#endif