HAMi-core 深度解析 - 算力隔离机制

void setspec() {
    // 通过 NVML 获取 GPU 的 SM 数量
    nvmlDeviceGetNumGpuCores(&g_sm_num);

    // 通过 CUDA API 获取每个 SM 的最大线程数
    cuDeviceGetAttribute(&g_max_thread_per_sm,
                         CU_DEVICE_ATTRIBUTE_MAX_THREADS_PER_MULTIPROCESSOR,
                         device);

    // 计算总 CUDA Core 数量（令牌桶容量基准）
    g_total_cuda_cores = g_max_thread_per_sm * g_sm_num * FACTOR;

    // 初始化令牌桶
    available_tokens = g_total_cuda_cores;
}

void rate_limiter(int grid_size) {
    // 如果未设置 SM 限制，直接返回
    if (target_sm_utilization == 0) return;

    // 如果策略为 disable，直接返回
    if (utilization_policy == POLICY_DISABLE) return;

    while (1) {
        int64_t current = atomic_load(&available_tokens);

        if (current >= grid_size) {
            // 尝试原子减少令牌
            if (atomic_compare_exchange_strong(&available_tokens,
                                               &current,
                                               current - grid_size)) {
                return;  // 成功获取令牌，允许 Kernel 启动
            }
            // CAS 失败，重试
            continue;
        }

        // 令牌不足，等待补充
        usleep(RATE_LIMITER_SLEEP_US);  // 短暂休眠
    }
}

int64_t delta(int actual_utilization, int target_utilization) {
    int diff = target_utilization - actual_utilization;

    // 基础补充量 = 总核心数的一个比例
    int64_t base = g_total_cuda_cores / REPLENISH_DIVISOR;

    if (diff > 0) {
        // 实际利用率 < 目标 → 增加令牌（加速）
        return base * (1 + diff / SCALE_FACTOR);
    } else if (diff < 0) {
        // 实际利用率 > 目标 → 减少令牌甚至不补充（减速）
        return base * diff / SCALE_FACTOR;  // 可能为负数
    } else {
        // 恰好达标 → 正常补充
        return base;
    }
}

void change_token(int64_t adjustment) {
    while (1) {
        int64_t current = atomic_load(&available_tokens);
        int64_t new_value = current + adjustment;

        // 限制令牌上限，防止累积过多
        if (new_value > g_total_cuda_cores) {
            new_value = g_total_cuda_cores;
        }
        // 限制令牌下限
        if (new_value < 0) {
            new_value = 0;
        }

        if (atomic_compare_exchange_strong(&available_tokens,
                                           &current, new_value)) {
            return;  // CAS 成功
        }
        // CAS 失败，重试
    }
}

CUresult hook_cuLaunchKernel(
    CUfunction f,
    unsigned int gridDimX, unsigned int gridDimY, unsigned int gridDimZ,
    unsigned int blockDimX, unsigned int blockDimY, unsigned int blockDimZ,
    unsigned int sharedMemBytes,
    CUstream hStream,
    void **kernelParams,
    void **extra)
{
    // 1. 计算 grid 大小
    int grid_size = gridDimX * gridDimY * gridDimZ;

    // 2. 调用限流器（可能会阻塞等待令牌）
    rate_limiter(grid_size);

    // 3. 记录 Kernel 启动时间（用于利用率计算）
    record_kernel_launch_time();

    // 4. 调用真实的 cuLaunchKernel
    return real_cuLaunchKernel(f,
        gridDimX, gridDimY, gridDimZ,
        blockDimX, blockDimY, blockDimZ,
        sharedMemBytes, hStream,
        kernelParams, extra);
}

// 获取进程级别的利用率信息
nvmlReturn_t nvmlDeviceGetProcessUtilization(
    nvmlDevice_t device,
    nvmlProcessUtilizationSample_t *utilization,
    unsigned int *processSamplesCount,
    unsigned long long lastSeenTimeStamp
);

// 返回结构体
typedef struct {
    unsigned int pid;           // 进程 PID
    unsigned long long timeStamp;
    unsigned int smUtil;        // SM 利用率百分比 (0-100)
    unsigned int memUtil;       // 显存带宽利用率
    unsigned int encUtil;       // 编码器利用率
    unsigned int decUtil;       // 解码器利用率
} nvmlProcessUtilizationSample_t;

// 在共享区域的进程槽位中记录利用率
typedef struct {
    pid_t pid;
    int target_sm_util;      // 目标利用率（配置值）
    int actual_sm_util;       // 实际利用率（NVML 采集）
    int64_t available_tokens; // 当前可用令牌
    uint64_t kernel_count;    // 累计 Kernel 启动次数
    uint64_t throttle_count;  // 累计被限流次数
} utilization_info_t;