HAMi-core 深度解析 - 显存隔离机制

uint64_t get_limit_from_env(const char *env_str) {
    // 支持的单位后缀
    // G / g → × 1024 × 1024 × 1024
    // M / m → × 1024 × 1024
    // K / k → × 1024
    // 无后缀 → 字节

    char *endptr;
    uint64_t value = strtoull(env_str, &endptr, 10);

    switch (*endptr) {
        case 'G': case 'g': value *= 1024 * 1024 * 1024; break;
        case 'M': case 'm': value *= 1024 * 1024; break;
        case 'K': case 'k': value *= 1024; break;
    }
    return value;
}

void lock_shrreg() {
    struct timespec timeout;
    clock_gettime(CLOCK_REALTIME, &timeout);
    timeout.tv_sec += LOCK_TIMEOUT_SECONDS;  // 通常 5 秒超时

    int ret = sem_timedwait(&shared_region->sem, &timeout);
    if (ret == -1 && errno == ETIMEDOUT) {
        // 超时 - 尝试修复死锁
        fix_lock_shrreg();
    }
    shared_region->owner_pid = getpid();
}

void unlock_shrreg() {
    shared_region->owner_pid = 0;
    sem_post(&shared_region->sem);
}

void init_proc_slot_withlock() {
    lock_shrreg();

    pid_t my_pid = getpid();

    // 查找空闲槽位
    for (int i = 0; i < MAX_PROCS; i++) {
        if (shared_region->procs[i].pid == 0) {
            // 找到空闲槽位
            shared_region->procs[i].pid = my_pid;
            shared_region->procs[i].status = ACTIVE;
            shared_region->procs[i].used_memory = 0;
            shared_region->proc_num++;
            break;
        }
    }

    unlock_shrreg();
}

void add_gpu_device_memory_usage(pid_t pid, int dev, uint64_t usage, int type) {
    lock_shrreg();

    shrreg_proc_slot_t *slot = find_proc_slot(pid);
    switch (type) {
        case MEMORY_TYPE_CONTEXT:
            slot->devices[dev].context_memory += usage;
            break;
        case MEMORY_TYPE_MODULE:
            slot->devices[dev].module_memory += usage;
            break;
        case MEMORY_TYPE_DATA:
            slot->devices[dev].data_memory += usage;
            break;
    }
    slot->devices[dev].total_used =
        slot->devices[dev].context_memory +
        slot->devices[dev].module_memory +
        slot->devices[dev].data_memory;

    unlock_shrreg();
}

int oom_check(uint64_t addon) {
    // 获取当前设备的全局使用量（所有进程之和）
    uint64_t current_usage = get_gpu_memory_usage(current_device);

    // 获取当前设备的显存限制
    uint64_t limit = get_device_memory_limit(current_device);

    // 判断是否超额
    if (current_usage + addon > limit) {
        return 1;  // OOM - 拒绝分配
    }
    return 0;  // OK - 允许分配
}

CUresult hook_cuMemFree_v2(CUdeviceptr dptr) {
    // 1. 在 allocated_list 中查找该地址
    allocated_list_entry *entry = find_chunk(dptr);

    if (entry != NULL) {
        uint64_t size = entry->size;
        int type = entry->type;

        // 2. 调用真实的 cuMemFree
        CUresult result = real_cuMemFree_v2(dptr);

        if (result == CUDA_SUCCESS) {
            // 3. 从 allocated_list 中移除
            remove_chunk(dptr);

            // 4. 更新共享区域使用量
            rm_gpu_device_memory_usage(getpid(), current_device, size, type);
        }

        return result;
    }

    // 未在追踪列表中 - 直接调用真实函数
    return real_cuMemFree_v2(dptr);
}

CUresult hook_cuMemGetInfo_v2(size_t *free, size_t *total) {
    // 获取配额限制
    uint64_t limit = get_device_memory_limit(current_device);

    // 获取当前使用量
    uint64_t used = get_gpu_memory_usage(current_device);

    // 返回虚拟化后的数据
    *total = limit;                                     // 总量 = 配额
    *free  = (limit > used) ? (limit - used) : 0;     // 可用 = 配额 - 已用

    return CUDA_SUCCESS;
}

CUresult hook_cuMemAllocAsync(CUdeviceptr *dptr, size_t bytesize,
                               CUstream hStream) {
    // OOM 检查（与同步路径相同）
    if (oom_check(bytesize)) {
        return CUDA_ERROR_OUT_OF_MEMORY;
    }

    // 调用真实的异步分配
    CUresult result = real_cuMemAllocAsync(dptr, bytesize, hStream);

    if (result == CUDA_SUCCESS) {
        // 记录分配（注意：异步分配的实际时机可能延迟）
        add_chunk(*dptr, bytesize, MEMORY_TYPE_DATA);
        add_gpu_device_memory_usage(getpid(), current_device,
                                     bytesize, MEMORY_TYPE_DATA);
    }

    return result;
}

// 在初始化时注册 fork 处理器
pthread_atfork(NULL, NULL, child_reinit_flag);

void child_reinit_flag() {
    // 在子进程中：重置初始化状态
    // 这样子进程的第一次 CUDA 调用会触发完整的初始化流程
    init_status = UNINITIALIZED;

    // 注意：不清理父进程的共享区域槽位
    // 父进程的槽位由父进程自身管理
}

void exit_handler() {
    struct timespec timeout;
    clock_gettime(CLOCK_REALTIME, &timeout);
    timeout.tv_sec += EXIT_LOCK_TIMEOUT;  // 退出时的锁超时

    int ret = sem_timedwait(&shared_region->sem, &timeout);
    if (ret == 0) {
        // 成功获取锁
        pid_t my_pid = getpid();

        // 查找并清理自己的槽位
        for (int i = 0; i < MAX_PROCS; i++) {
            if (shared_region->procs[i].pid == my_pid) {
                // 释放该进程在所有设备上的使用量
                memset(&shared_region->procs[i], 0,
                       sizeof(shrreg_proc_slot_t));
                shared_region->proc_num--;
                break;
            }
        }

        sem_post(&shared_region->sem);
    }
    // 如果获取锁超时，放弃清理
    // 后续其他进程会通过 fix_lock_shrreg 回收
}

uint64_t get_gpu_memory_usage(int device_id) {
    uint64_t total_used = 0;

    lock_shrreg();

    for (int i = 0; i < MAX_PROCS; i++) {
        if (shared_region->procs[i].pid != 0 &&
            shared_region->procs[i].status == ACTIVE) {
            total_used += shared_region->procs[i]
                          .devices[device_id].total_used;
        }
    }

    unlock_shrreg();
    return total_used;
}