/*
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 cpu utils
*/

#ifndef CPU_UTILS_H
#define CPU_UTILS_H

#include <vector>
#include <string>
#include <omp.h>

// 工作分配结构体
typedef struct WorkDivider {
    int num_threads;        // 当前并行组内的线程数量（= 全局总线程数 / para）
    int para;               // 并行任务组的数量（各组互不干扰，独立执行相同或不同任务）
    int global_tid;         // 当前线程的全局唯一 ID（范围：0 ～ total_threads - 1）
    int tid;                // 当前线程在其所属组内的局部 ID（范围：0 ～ num_threads - 1）
    int num_numas;          // 参与计算的 NUMA 节点总数
    int threads_per_numa;   // 每个 NUMA 节点上分配的线程数（通常 = num_threads / num_numas）
    int my_numa;            // 当前线程所在的 NUMA 节点编号（范围：0 ～ num_numas - 1）
    int tid_in_numa;        // 当前线程在其所在 NUMA 节点内的局部 ID（范围：0 ～ threads_per_numa - 1）
} WorkDivider;

// 工作范围结构体:单numa和多numa
typedef struct SingleNumaWorkRange {
    int begin_thread;
    int end_thread;
    int work_per_thread;
} SingleNumaWorkRange;

typedef struct MultiNumaWorkRange {
    int begin_numa;
    int work_per_numa;
    int begin_thread;
    int end_thread;
    int work_per_thread;
} MultiNumaWorkRange;

// 工作分配函数
void divide_work_single_numa(const WorkDivider *divider, int total_workitems, SingleNumaWorkRange *pstSingleRange);
void init_work_divider(WorkDivider* divider, int numas, int para = 1, bool exp = false);
// exp = true 对应一个专家加载到一个numa节点上， false 对应一个专家均分多个numa节点上
void divide_all_work(const WorkDivider* divider, int total_workitems, SingleNumaWorkRange* pstSingleRange);
// divide_all_work: 将总工作量均匀分配给当前线程所在 NUMA 节点内的所有线程，threads_per_numa 个线程平分 total_workitems。
void divide_work_first_numa(const WorkDivider* divider, int total_workitems, SingleNumaWorkRange* pstSingleRange);
// divide_work_first_numa: 将总工作量仅分配给当前 NUMA 节点中“第一个”线程组（通常是 my_numa == 0 的线程），其余线程不分配工作。
void divide_work_all_numas(const WorkDivider* divider, int total_workitems, MultiNumaWorkRange* pstNulRange); 
// divide_work_all_numas: 将总工作量按 NUMA 节点和线程两级粒度进行划分。
// 首先将工作均分到 [begin_numa, end_numa) 的每个 NUMA 节点（每 NUMA 分得 work_per_numa 份），
// 然后在每个 NUMA 节点内部，再将分配到的工作均分给其所属的 [begin_thread, end_thread) 线程。

#endif // CPU_UTILS_H