CUDA独立上下文模块加载

在这里插入图片描述

大多数 CUDA 开发人员都熟悉 cuModuleLoad API 及其对应项，用于将包含设备代码的模块加载到 CUDA 上下文中。在大多数情况下，您希望在所有设备上加载相同的设备代码。这需要将设备代码显式加载到每个 CUDA 上下文中。此外，不控制上下文创建和销毁的库和框架必须跟踪它们以显式加载和卸载模块。

这篇文章讨论了 CUDA 12.0 中引入的上下文独立加载，它解决了这些问题。

上下文相关的加载(Context-dependent loading)

传统上，模块加载总是与 CUDA 上下文相关联。下面的代码示例显示了将相同的设备代码加载到两个设备然后在它们上启动内核的传统方法。

// Device 0
cuDeviceGet(&device0, 0);
cuDevicePrimaryCtxRetain(&ctx0, device0);
cuModuleLoad(&module0, “myModule.cubin”);
// Device 1
cuDeviceGet(&device1, 1);
cuDevicePrimaryCtxRetain(&ctx1, device1);
cuModuleLoad(&module1, “myModule.cubin”);

在每个设备上启动内核需要您检索每个模块的 CU 函数，如以下代码示例所示：

// Device 0
cuModuleGetFuntion(&function0, module0, “myKernel”);
cuLaunchKernel(function0, …);
// Device 1
cuModuleGetFuntion(&function1, module1, “myKernel”);
cuLaunchKernel(function1, …);

这会增加应用程序中的代码复杂性，因为您必须检索和跟踪每个上下文和每个模块的类型。您还必须使用 cuModuleUnload API 显式卸载每个模块。

当库或框架主要使用 CUDA 驱动程序 API 来加载它们自己的模块时，问题会加剧。他们可能无法完全控制应用程序拥有的上下文的生命周期。

// Application code

libraryInitialize();
cuDevicePrimaryCtxRetain(&ctx0, device0);
libraryFunc();
cuDevicePrimaryCtxRetain(&ctx0, device1);
libraryFunc();
libraryDeinitialize();

// Library code

libraryInitialize() {
  map<CUcontext, CUmodule> moduleContextMap;
}

libraryFunc() {
  cuCtxGetCurrent(&ctx);
  if (!moduleContextMap.contains(ctx)){
    cuModuleLoad(&module, “myModule.cubin”);
    moduleContextMap[ctx] = module;
  }
  else {
    module = moduleContextMap[ctx];
  }

  cuModuleGetFuntion(&function, module, “myKernel”);
  cuLaunchKernel(function, …);
}

libraryDeinitialize() {
  moduleContextMap.clear();
}

在代码示例中，库必须检查新上下文并将模块显式加载到其中。它还必须维护状态以检查模块是否已加载到上下文中。

理想情况下，状态可以在上下文被销毁后释放。但是，如果库无法控制上下文的生命周期，这是不可能的。

这意味着资源的释放必须延迟到库取消初始化。这不仅增加了代码的复杂性，而且还导致库占用资源的时间超过了它必须占用的时间，可能会拒绝应用程序的另一部分使用该内存。

另一种选择是让库和框架对用户施加额外的限制，以确保他们对资源分配和清理有足够的控制权。

上下文独立的加载(Context-independent loading)

CUDA 12.0 通过添加 cuLibrary* 和 cuKernel* API 引入了上下文无关加载，从而解决了这些问题。通过上下文独立加载，在应用程序创建和销毁上下文时，CUDA 驱动程序会自动将模块加载和卸载到每个 CUDA 上下文中。

// Load library
cuLibraryLoadFromFile(&library,“myModule.cubin”, …);
cuLibraryGetKernel(&kernel, library, “myKernel”);

// Launch kernel on the primary context of device 0
cuDevicePrimaryCtxRetain(&ctx0, device0);
cuLaunchKernel((CUkernel)kernel, …);

// Launch kernel on the primary context of device 1
cuDevicePrimaryCtxRetain(&ctx1, device1);
cuLaunchKernel((CUkernel)kernel, …);

// Unload library
cuLibraryUnload(library);

如代码示例所示，cuLibraryLoadFromFile API 负责在创建或初始化上下文时加载模块。在示例中，这是在 cuDevicePrimaryCtxRetain 期间完成的。

此外，您现在可以使用与上下文无关的句柄 CUkernel 启动内核，而不必维护每个上下文的 CUfunction。 cuLibraryGetKernel 检索设备函数 myKernel 的上下文无关句柄。然后可以通过指定与上下文无关的句柄 CUkernel 来使用 cuLaunchKernel 启动设备功能。 CUDA 驱动程序负责根据此时处于活动状态的上下文在适当的上下文中启动设备功能。

库和框架现在可以分别在初始化和取消初始化期间简单地加载和卸载模块一次。

// Application code

libraryInitialize();
cuDevicePrimaryCtxRetain(&ctx0, device0);
libraryFunc();
cuDevicePrimaryCtxRetain(&ctx0, device1);
libraryFunc();
libraryDeinitialize();

// Library code

libraryInitialize() {
  cuLibraryLoadFromFile(&library,“myModule.cubin”, …);
  cuLibraryGetKernel(&kernel, library, “myKernel”);
}

libraryFunc() {
  cuLaunchKernel((CUkernel)kernel, …);
}

libraryDeinitialize() {
  cuLibraryUnload(library);
}