2.3.cuda驱动API-上下文管理设置及其作用

前言

杜老师推出的 tensorRT从零起步高性能部署 课程，之前有看过一遍，但是没有做笔记，很多东西也忘了。这次重新撸一遍，顺便记记笔记

本次课程学习精简 CUDA 教程-Driver API 上下文管理设置及其作用

课程大纲可看下面的思维导图

1. CUcontext

对于 context 你需要知道的是：

1. context 是一种上下文，可以关联对 GPU 的所有操作
2. context 与一块显关联，一个显卡可以被多个 context 关联
3. 每个线程都有一个栈结构储存 context，栈顶是当前使用的 context，对应有 push、pop 函数操作 context 的栈，所有 api 都以当前 context 为操作目标

试想一下，如果执行任何操作你都需要传递一个 device 才决定送到哪个设备执行，得多麻烦

下图对比了没有 context 和有 context 的代码

图1-1 没有context和有context代码对比

从图中可以看出，在没有使用 context 的情况下，我们直接调用 CUDA 的内存分配、释放和数据拷贝函数，这些函数的参数中需要传入设备标识符来指定操作的设备。每个函数调用都是独立的，没有建立起与设备的关联

在有 context 的情况下，我们首先使用 cuCreateContext 函数创建了一个 context，并将其与特定设备关联起来。然后使用 cuPushCurrent 函数将该 context 设置为当前 context。接下来的内存分配、释放和数据拷贝函数调用将自动使用当前 context 进行操作，而不需要显式地指定设备标识符。

在完成上下文相关的操作后，我们使用 cuPopCurrent 函数将 context 从当前 contex 栈中弹出，恢复之前的上下文设置。

使用 context 的好处是可以将一系列相关的 CUDA 操作关联到一个 context 中，简化代码，并提提高执行效率。

上面提到的都是关于手动管理 context，而关于自动管理 context 有以下几点需要说明：

1. 由于是高频操作，是一个线程基本固定访问一个显卡不变，且只使用一个 context，很少会用到多 context
2. CreateContext、PushCurrent、PopCurrent 这种多 context 管理就显得麻烦，还得再简单
3. 因此推出了 cuDevicePrimaryCtxRetain，为设备关联主 context，分配、释放、设置、栈都不用你管
4. primaryContext：给我设备 id，给你 context 并设置好，此时一个显卡对应一个 primary context
5. 不用线程，只要设备 id 一样，primary context 就一样，context 是线程安全的

下图对比了手动管理 context 和自动管理 context 的代码

图1-2 有contex和不用管理context栈代码对比

在上图中，我们使用 cuDevicePrimaryCtxRetain 函数将设备的主要上下文（primary context）与特定设备关联起来。接下来的内存分配、释放和数据拷贝函数调用将使用该主要上下文进行操作，而无需显式地设置当前上下文。这种情况下，不需要显式地管理上下文栈，代码更加简洁。

context 案例的示例代码如下：

// CUDA驱动头文件cuda.h
#include <cuda.h>   // include <> 和 "" 的区别    
#include <stdio.h>  // include <> : 标准库文件 
#include <string.h> // include "" : 自定义文件  详细情况请查看 readme.md -> 5

#define checkDriver(op)  __check_cuda_driver((op), #op, __FILE__, __LINE__)

bool __check_cuda_driver(CUresult code, const char* op, const char* file, int line){
    if(code != CUresult::CUDA_SUCCESS){    // 如果 成功获取CUDA情况下的返回值 与我们给定的值(0)不相等， 即条件成立， 返回值为flase
        const char* err_name = nullptr;    // 定义了一个字符串常量的空指针
        const char* err_message = nullptr;  
        cuGetErrorName(code, &err_name);    
        cuGetErrorString(code, &err_message);   
        printf("%s:%d  %s failed. \n  code = %s, message = %s\n", file, line, op, err_name, err_message); //打印错误信息
        return false;
    }
    return true;
}

int main(){

    // 检查cuda driver的初始化
    checkDriver(cuInit(0));

    // 为设备创建上下文
    CUcontext ctxA = nullptr;                                   // CUcontext 其实是 struct CUctx_st*（是一个指向结构体CUctx_st的指针）
    CUcontext ctxB = nullptr;
    CUdevice device = 0;
    checkDriver(cuCtxCreate(&ctxA, CU_CTX_SCHED_AUTO, device)); // 这一步相当于告知要某一块设备上的某块地方创建 ctxA 管理数据。输入参数 参考 https://www.cs.cmu.edu/afs/cs/academic/class/15668-s11/www/cuda-doc/html/group__CUDA__CTX_g65dc0012348bc84810e2103a40d8e2cf.html
    checkDriver(cuCtxCreate(&ctxB, CU_CTX_SCHED_AUTO, device)); // 参考 1.ctx-stack.jpg
    printf("ctxA = %p\n", ctxA);
    printf("ctxB = %p\n", ctxB);
    /* 
        contexts 栈：
            ctxB -- top <--- current_context
            ctxA 
            ...
     */

    // 获取当前上下文信息
    CUcontext current_context = nullptr;
    checkDriver(cuCtxGetCurrent(&current_context));             // 这个时候current_context 就是上面创建的context
    printf("current_context = %p\n", current_context);

    // 可以使用上下文堆栈对设备管理多个上下文
    // 压入当前context
    checkDriver(cuCtxPushCurrent(ctxA));                        // 将这个 ctxA 压入CPU调用的thread上。专门用一个thread以栈的方式来管理多个contexts的切换
    checkDriver(cuCtxGetCurrent(&current_context));             // 获取current_context (即栈顶的context)
    printf("after pushing, current_context = %p\n", current_context);
    /* 
        contexts 栈：
            ctxA -- top <--- current_context
            ctxB
            ...
    */
    

    // 弹出当前context
    CUcontext popped_ctx = nullptr;
    checkDriver(cuCtxPopCurrent(&popped_ctx));                   // 将当前的context pop掉，并用popped_ctx承接它pop出来的context
    checkDriver(cuCtxGetCurrent(&current_context));              // 获取current_context(栈顶的)
    printf("after poping, popped_ctx = %p\n", popped_ctx);       // 弹出的是ctxA
    printf("after poping, current_context = %p\n", current_context); // current_context是ctxB

    checkDriver(cuCtxDestroy(ctxA));
    checkDriver(cuCtxDestroy(ctxB));

    // 更推荐使用cuDevicePrimaryCtxRetain获取与设备关联的context
    // 注意这个重点，以后的runtime也是基于此, 自动为设备只关联一个context
    checkDriver(cuDevicePrimaryCtxRetain(&ctxA, device));       // 在 device 上指定一个新地址对ctxA进行管理
    printf("ctxA = %p\n", ctxA);
    checkDriver(cuDevicePrimaryCtxRelease(device));
    return 0;
}

运行效果如下：

图1-3 context案例运行效果

代码开始创建了两个上下文 ctxA 和 ctxB。通过调用 cuCtexCreate 函数来为特定设备（使用设备标识符 device）创建上下文。然后，代码使用 cuCtxGetCurrent 函数获取当前上下文，并打印其地址。可以看到，在刚创建上下文后，当前上下文与 ctxB 的地址相同。

接下来，代码通过 cuCtxPushCurrent 函数将 ctxA 压入上下文栈，成为当前上下文。然后使用 cuCtxPopCurrent 函数将当前上下文弹出，并用 popped_ctx 变量接收被弹出的上下文。再次调用 cuCtxGetCurrent 函数可以看到当前上下文变成了 ctxB，而 popped_ctx 中保存了被弹出的 ctxA。

最后，代码也演示了使用 cuDevicePrimaryCtxRetain 函数来自动管理 context

2. 补充知识

我们需要补充下 context 上下文的相关知识：(from 杜老师)

1. 设备与特定进程相关连的所以状态。比如，你写的一段 kernel code 对 GPU 的使用会造成不同状态（内存映射、分配、加载的code），Context 则保存着所有的管理数据来控制和使用设备

• 一个类比的例子就像你与小明和小红分别对话，你与小明聊 China（中国），你与小红也聊 china（瓷器），但是你们聊的可能都不是一个东西，我们管理这样的数据，所以我们创造了 context
• gpu 的 context 相当于 cpu 的 program，一块 gpu 上可以有个 contexts，但是它们之间是相互隔离的。建议一块设备就一个 context
• 参考：https://aiprobook.com/deep-learning-for-programmers/

2. 上下文管理可以干的事儿：

• 持有分配的内存列表
• 持有加载进该设备的 kernel code
• cpu 与 gpu 之间的 unified memory
• …

3. 如何管理上下文：

• 在 cuda driver 同样需要显示管理上下文
  • 开始时 cuCtxCreate() 创建上下文，结束时 cuCtxDestroy 销毁上下文。像文件管理一样需手动开关。
  • 用 cuDevicePrimaryCtxRetain() 创建上下文更好！
  • cuCtxGetCurrent() 获取当前上下文
  • 可以使用堆栈管理多个上下文 cuCtxPushCurrent() 压入，cuCtxPopCurrent() 推出
  • 对 ctxA 使用 cuCtxPushCurrent() 和 cuCtxCreate() 都相当于将 ctxA 放到栈顶（让它成为 current context）
• cuda runtime 可以自动创建，是基于 cuDevicePrimaryCtxRetain() 创建的

总结

本次课程学习了 Diver API 的上下文管理设置，通过手动创建管理 context 可以将一系列相关的 CUDA 操作关联到一个 context 中，简化代码并提高执行效率。当然更推荐使用 cuDevicePrimaryCtxRetain 获取与设备关联的 context，它不需要显式地管理上下文栈，代码更加简洁。