ROCm HSAKMT pthread_atfork 回调函数技术分析-优快云博客

本文链接：https://blog.youkuaiyun.com/shenjunpeng/article/details/154647351

1. 概述

本文档深入分析 ROCm HSAKMT 库中的三个 pthread_atfork 回调函数：prepare_fork_handler、parent_fork_handler 和 child_fork_handler。这些函数是多进程环境下保证线程安全和资源一致性的关键组件，特别是在 GPU 计算环境中处理进程分叉时的复杂场景。

2. 函数功能分析

2.1 prepare_fork_handler 函数

static void prepare_fork_handler(void)
{
    pthread_mutex_lock(&hsakmt_mutex);
}

功能说明：

执行时机：在 fork() 调用之前，在父进程中执行
主要作用：获取全局互斥锁 hsakmt_mutex，确保在 fork 操作期间没有其他线程能够修改 HSAKMT 的全局状态
安全保障：防止在 fork 过程中出现竞态条件，确保父进程和子进程都能看到一致的内存状态

设计意图：
该函数实现了 fork 操作的原子性保护。在多线程环境中，如果不加锁保护，可能会出现以下问题：

线程 A 正在修改 HSAKMT 的全局状态（如设备列表、内存映射等）
线程 B 同时调用 fork()
子进程可能看到不一致的状态，导致后续操作失败

2.2 parent_fork_handler 函数

static void parent_fork_handler(void)
{
    pthread_mutex_unlock(&hsakmt_mutex);
}

功能说明：

执行时机：在 fork() 调用之后，在父进程中执行
主要作用：释放之前在 prepare_fork_handler 中获取的互斥锁
状态恢复：使父进程恢复到 fork 之前的正常状态，允许其他线程继续访问 HSAKMT 资源

设计考虑：
父进程在 fork 后需要继续正常运行，因此必须释放锁以避免死锁。父进程的所有资源（文件描述符、内存映射、设备句柄等）都保持有效，不需要重新初始化。

2.3 child_fork_handler 函数

static void child_fork_handler(void)
{
    pthread_mutex_init(&hsakmt_mutex, NULL);
    hsakmt_forked = true;
}

功能说明：

执行时机：在 fork() 调用之后，在子进程中执行
主要作用：
1. 重新初始化互斥锁 hsakmt_mutex
2. 设置 hsakmt_forked 标志为 true

关键设计决策：

互斥锁重新初始化：
- 子进程继承了父进程的内存空间，包括被锁定的 mutex
- 由于子进程中只有一个线程（调用 fork 的线程），直接解锁可能导致未定义行为
- 重新初始化 mutex 确保子进程有一个干净的、未锁定的 mutex 状态
forked 标志设置：
- hsakmt_forked = true 标记子进程已经从父进程分叉出来
- 这个标志被 hsakmt_is_forked_child() 函数使用，用于检测进程是否为子进程
- 触发子进程中的资源清理和重新初始化逻辑

3. 系统架构分析

3.1 全局状态管理

HSAKMT 库维护以下关键的全局状态：

// 来自 globals.c
extern int hsakmt_kfd_fd;		//全局fd
extern pthread_mutex_t hsakmt_mutex;      // 全局互斥锁
extern bool hsakmt_forked;               // fork 状态标志
extern unsigned long hsakmt_kfd_open_count; // KFD 打开计数

3.2 Fork 检测机制

系统使用双重检测机制来识别子进程：

bool hsakmt_is_forked_child(void)
{
    pid_t cur_pid;
    
    if (hsakmt_forked)          // 1. 检查 atfork 标志
        return true;
    
    cur_pid = getpid();
    
    if (parent_pid == -1) {     // 2. 首次调用，记录父进程 PID
        parent_pid = cur_pid;
        return false;
    }
    
    if (parent_pid != cur_pid) { // 3. PID 检测机制
        hsakmt_forked = true;
        return true;
    }
    
    return false;
}

设计巧思：

pthread_atfork 机制处理标准 fork() 调用
PID 检测机制处理直接系统调用或 clone() 等绕过 libc 的情况

资源清理策略

子进程中的资源清理通过 clear_after_fork() 函数实现：

static void clear_after_fork()
{
    hsakmt_clear_process_doorbells();    // 清理门铃页
    hsakmt_clear_events_page();          // 清理事件页
    hsakmt_fmm_clear_all_mem();          // 清理内存映射
    hsakmt_destroy_device_debugging_memory(); // 清理调试内存
 
    if (hsakmt_kfd_fd >= 0) {
        close(fd);                          // 关闭 KFD 文件描述符
    }
    
    hsakmt_kfd_open_count = 0;              // 重置计数器
    parent_pid = -1;                        // 重置父进程 PID
    hsakmt_forked = false;                  // 重置 fork 标志
}