如何在C++中实现全局可访问的线程专属数据存储？

问题描述 #

需要为每个线程存储专属数据，支持多位置全局访问，但无法直接使用extern thread_local，希望实现类似以下语法的访问方式：

auto data = std::this_thread::my_data;

需求背景 #

基于LibTorch（PyTorch C++ API）设计调度器，多线程并行运行时每个线程绑定独立的CUDA Context。目前通过cuCtxGetCurrent()获取当前Context，但该操作存在不可忽略的性能开销，且多线程调用时会引发同步问题。

更新说明 #

无法直接使用thread_local的核心原因：调度器中调用的forward是LibTorch已编译完成的后端函数，内部会调用大量PyTorch原生函数。需要在这些后端代码逻辑中获取当前线程分配的CUDA Context和Stream，但直接使用thread_local变量难以让后端代码直接访问，因此需要一种全局可访问的线程本地存储方案。

可行实现方案 #

方案一：基于C++ thread_local的全局封装 #

利用thread_local静态变量的线程隔离特性，封装成全局可访问的接口，避免直接使用extern thread_local的限制：

#include <thread>

// 根据业务需求定义线程本地数据结构
struct ThreadCUDAData {
    void* cuda_context;
    void* cuda_stream;
    // 可添加其他线程专属字段
};

namespace ThreadLocalCUDA {
    // 线程本地静态变量，每个线程拥有独立副本
    static thread_local ThreadCUDAData* tls_data = nullptr;

    // 设置当前线程的CUDA数据
    void set_cuda_data(ThreadCUDAData* data) {
        tls_data = data;
    }

    // 获取当前线程的CUDA数据，模拟目标语法的全局访问
    ThreadCUDAData& get_cuda_data() {
        // 按需初始化，实际场景可根据业务逻辑调整（比如调度器分配时初始化）
        if (!tls_data) {
            tls_data = new ThreadCUDAData{};
        }
        return *tls_data;
    }
}

// 使用示例
auto& current_thread_cuda = ThreadLocalCUDA::get_cuda_data();
// 等价于期望的std::this_thread::my_data

优势：

• 访问无锁开销，性能接近原生thread_local
• 线程生命周期自动绑定，若用动态分配，可结合线程清理回调管理内存
• 接口简洁，全局任意位置可调用

方案二：基于系统原生TLS API（跨平台兼容） #

若需兼容C++11之前的环境，或需要更底层的控制，可使用系统原生线程本地存储API：

POSIX平台（Linux/macOS）

#include <pthread.h>

struct ThreadCUDAData {
    void* cuda_context;
    void* cuda_stream;
};

static pthread_key_t cuda_tls_key;

// 线程退出时自动清理数据的回调函数
void cuda_tls_destructor(void* data) {
    delete static_cast<ThreadCUDAData*>(data);
}

// 程序启动时初始化TLS键（需调用一次）
void init_cuda_tls() {
    pthread_key_create(&cuda_tls_key, cuda_tls_destructor);
}

// 全局获取当前线程的CUDA数据
ThreadCUDAData& get_cuda_data() {
    ThreadCUDAData* data = static_cast<ThreadCUDAData*>(pthread_getspecific(cuda_tls_key));
    if (!data) {
        data = new ThreadCUDAData{};
        pthread_setspecific(cuda_tls_key, data);
    }
    return *data;
}

Windows平台

#include <windows.h>

struct ThreadCUDAData {
    void* cuda_context;
    void* cuda_stream;
};

static DWORD cuda_tls_index;

// 程序启动时初始化TLS索引（需调用一次）
void init_cuda_tls() {
    cuda_tls_index = TlsAlloc();
}

// 全局获取当前线程的CUDA数据
ThreadCUDAData& get_cuda_data() {
    ThreadCUDAData* data = static_cast<ThreadCUDAData*>(TlsGetValue(cuda_tls_index));
    if (!data) {
        data = new ThreadCUDAData{};
        TlsSetValue(cuda_tls_index, data);
    }
    return *data;
}

// 线程退出时手动清理（或结合DLL_THREAD_DETACH事件）
void cleanup_cuda_tls() {
    ThreadCUDAData* data = static_cast<ThreadCUDAData*>(TlsGetValue(cuda_tls_index));
    delete data;
    TlsSetValue(cuda_tls_index, nullptr);
}

优势：

• 兼容更低版本的C++标准和系统环境
• 底层可控，可自定义线程退出时的清理逻辑

方案适配LibTorch场景 #

在调度器分配CUDA Context和Stream给线程时，调用set_cuda_data（或对应系统API的设置函数）完成绑定。之后无论在自定义代码还是LibTorch后端逻辑中，只要包含该封装的头文件（或通过C接口导出给后端库），即可直接调用get_cuda_data()获取当前线程的专属CUDA资源，完全替代cuCtxGetCurrent()，消除其带来的开销和同步问题。

内容的提问来源于stack exchange，提问作者Amir Fakhim Babaei