cuda driver-上下文管理

前言

CUcontext是CUDA（Compute Unified Device Architecture）中的一个概念，用于管理和跟踪与GPU设备的交互。它代表了一个CUDA上下文，即一个应用程序与GPU之间的通信环境。

1.CUcontext作用

在CUDA中，一个应用程序可以与多个GPU设备进行交互。每个GPU设备都有一个唯一的设备ID来标识它们。当应用程序需要与某个特定GPU设备进行通信时，它需要先创建一个与该设备关联的上下文（CUcontext）。
CUcontext提供了一种隔离和管理与GPU设备的交互的方式。通过创建和销毁上下文，应用程序可以切换和管理不同的GPU设备。上下文还提供了一些重要的功能，如内存管理、核函数的执行等。
在CUDA编程中，应用程序需要使用一些API函数来创建和操作上下文，如cuCtxCreate()用于创建上下文，cuCtxDestroy()用于销毁上下文等。在创建上下文时，还可以指定一些参数来配置上下文的属性，如与上下文关联的设备ID等。
总之，CUcontext是CUDA中用于管理和跟踪与GPU设备的交互的重要概念，它提供了一种隔离和管理GPU设备的方式，使应用程序能够有效地利用GPU的计算资源。

2.CUcontext常用API

API	Description
cuCtxCreate	用于创建上下文
cuCtxDestroy	用于销毁上下文
cuDevicePrimaryCtxRetain	自动管理上下文而无需显式地设置当前上下文
cuCtxGetCurrent	获取当前上下文
cuCtxPushCurrent	压入
cuCtxPopCurrent	弹出

实例代码如下：

#include <cuda.h>
#include <stdio.h>
#include <string.h>

#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)
    {
        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初始化
    if (!checkDriver(cuInit(0)))
    {
        return -1;
    }

    // 获取 cuda 驱动版本
    int driver_version = 0;
    if (!checkDriver(cuDriverGetVersion(&driver_version)))
    {
        return -1;
    }
    printf("CUDA Driver version is %d\n", driver_version);

    // 获取当前设备信息
    char device_name[100];
    CUdevice device = 0;
    if (!checkDriver(cuDeviceGetName(device_name, sizeof(device_name), device)))
    {
        return -1;
    }
    printf("Device %d name is %s\n", device, device_name);

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

    // 可以使用上下文堆栈对设备管理多个上下文
    // 压入当前context
    // 将这个 ctxA 压入CPU调用的thread上。专门用一个thread以栈的方式来管理多个contexts的切换
    checkDriver(cuCtxPushCurrent(ctxA));
    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));
    // 获取current_context(栈顶的)
    checkDriver(cuCtxGetCurrent(&current_context));
    // 弹出的是ctxA
    printf("after poping, popped_ctx = %p\n", popped_ctx);
    // current_context是ctxB
    printf("after poping, current_context = %p\n", current_context);

    checkDriver(cuCtxDestroy(ctxA));
    checkDriver(cuCtxDestroy(ctxB));
    // 更推荐使用cuDevicePrimaryCtxRetain获取与设备关联的context
    // 注意这个重点，以后的runtime也是基于此, 自动为设备只关联一个context
    checkDriver(cuDevicePrimaryCtxRetain(&ctxA, device));
    printf("ctxA = %p\n", ctxA);
    checkDriver(cuDevicePrimaryCtxRelease(device));
    return 0;
}

代码开始创建了两个上下文 ctxA 和 ctxB。通过调用 cuCtexCreate 函数来为特定设备（使用设备标识符 device）创建上下文。然后，代码使用 cuCtxGetCurrent 函数获取当前上下文，并打印其地址。可以看到，在刚创建上下文后，当前上下文与 ctxB 的地址相同。
接下来，代码通过 cuCtxPushCurrent 函数将 ctxA 压入上下文栈，成为当前上下文。然后使用 cuCtxPopCurrent 函数将当前上下文弹出，并用 popped_ctx 变量接收被弹出的上下文。再次调用 cuCtxGetCurrent 函数可以看到当前上下文变成了 ctxB，而 popped_ctx 中保存了被弹出的 ctxA。
最后，代码也演示了使用 cuDevicePrimaryCtxRetain 函数来在 device 上指定一个新地址对ctxA进行管理，自动管理上下文。
运行结果如下:

CUDA Driver version is 11040
Device 0 name is NVIDIA GeForce RTX 3060 Laptop GPU
ctxA = 0x5620c3abc020
ctxB = 0x5620c3f88a20
current_context = 0x5620c3f88a20
after pushing, current_context = 0x5620c3abc020
after poping, popped_ctx = 0x5620c3abc020
after poping, current_context = 0x5620c3f88a20
ctxA = 0x5620c3ac6f20