系统软件层

TVM 通过使用 CUDA Driver API 与底层 GPU 进行交互的具体流程可以总结为以下几个步骤:

  • 在 TVM 中,首先通过 CUDA Driver API 初始化 GPU 上下文 (cuCtxCreate 或者 cuCtxSetCurrent)。这个上下文用于管理与特定 GPU 设备的交互,并确保后续的 GPU 操作在这个上下文中执行。

  • 在 GPU 上分配内存,TVM 通过调用 cuMemAlloc 来请求 GPU 显存,以存储计算所需的输入、输出及中间结果数据。可以使用 cuMemFree 来释放不再使用的显存。

  • TVM 需要将数据从主机(CPU)传输到设备(GPU)。通过调用 cuMemcpyHtoD(从主机到设备)和 cuMemcpyDtoH(从设备到主机)进行数据传输。

  • 使用 cuModuleGetFunction 获取 CUDA 函数句柄,通过 cuLaunchKernel 启动内核执行。启动时可以指定线程块和线程网格的配置,这样可以充分利用 GPU 的并行计算能力。

  • 为了确保所有的 GPU 操作都完成,可以使用 cuCtxSynchronize,它会等待所有在当前上下文中发起的操作执行完毕。

  • 当 GPU 计算任务完成之后,TVM 会通过调用相应的 CUDA Driver API 函数来释放内存和资源,包括 cuMemFreecuModuleUnload,以避免内存泄漏。