PyTorch C10 CUDA 模块源码结构解读（参考版本：PyTorch 2.0.0 RC1）_cuda算子的源码

什么是 C10？

为实现精准的控制和对CUDA等硬件的使用，PyTorch底层采用C++编写；为给机器学习相关领域开发者提供更舒适的界面，PyTorch的高层模块使用Python编写，并提供Python接口。

其基本结构如下图所示：

使用C++完成对CUDA等底层硬件的对接，并十分高效地实现基础组件和部分算法；借助Python原生调用能力，将用C++实现的模块封装成接口，提供给Python代码调用，如此即可使用Python实现更多模块；在模块之上，封装更好用的接口，提供给机器学习领域的开发者使用。

PyTorch的跨语言环境接口主要有两大部分：C++与原生运行环境的对接、Python与C++的对接。

C++与原生运行环境的对接全部在ATen和C10内实现。如，C10的CUDAFunctions模块完成对NVIDIA CUDA Runtime API的二次封装，以支持上层更定制化的操作。

Python与C++的对接层为torch/_C模块。该部分接口在编译期自动生成，为Python代码提供支持。

PyTorch 整体模块依赖关系

经过对代码的阅读，以及与PyTorch C++ API维护者Brain Hirsh的邮件沟通，我们学习到，PyTorch的整体结构大致如下图所示：

ATen代表A Tensor Library，是PyTorch最原始的Tensor基础库。

C10是一个仍处于开发阶段的库，逐步代替ATen的地位。我们可以注意到，ATen内的很多实现已经改为直接包含C10内的同类实现。如，Tensor的实现已经完全迁移到C10内。因此，我们可以观察到，ATen对C10呈现出依赖关系。

ATen和C10封装了含CUDA在内的底层操作。基于ATen和C10的能力，PyTorch的开发者们在csrc（代表C Source）模块内使用C和C++实现autograd等更高级的功能。

如Java一样，Python具有原生调用接口的功能。PyTorch的原生接口层为_C模块，其中暴露的所有原生功能皆由csrc模块提供。

PyTorch的更高层功能使用Python编写，并在部分功能处通过_C调用原生代码，获得更好的性能和对CUDA等高级硬件的操控能力。

C10 CUDA 模块及其子模块设计解读

C10库的CUDA模块直接覆盖在CUDA运行环境上，为外层开发者提供基础资源管理服务。

该模块由多个子模块组成，含实现了GC和类Arena的GPU内存管理器和CUDA执行流管理器。它们通过封装思路较为奇特的CUDA Runtime API，为外层提供更友好的接口，提供更优的支持；通过对资源的复用，减少对底层API的频繁访问，提供更好的性能。

本节将对C10/CUDA模块内的多个子模块展开分析。

CUDAExceptions

基本结构

该模块为对CUDA的调用提供最基础的保障功能。

C10_CUDA_CHECK负责检查运行过程调用的CUDA函数是否支持。

C10_CUDA_KERNEL_LAUNCH_CHECK用于检测CUDA是否存在未处理的错误报告。

C10_CUDA_CHECK_WARN和IGNORE_ERROR皆用于包装某个运行过程。运行后，会对CUDA是否出错进行检查，并自动输出带文件名和代码行号的报错日志。

算法设计

c10_cuda_check_implementation通过读取CUDA Kernel Launch Registry（一个基于循环队列的日志记录器）单例，检查其是否登记错误信息，以判断是否出现过运行时错误。如果有出现，则将报错信息输出。

C10_CUDA_CHECK则是在执行被包裹语句后，调用c10_cuda_check_implementation，得以实现功能。

C10_CUDA_CHECK_WARN和IGNORE_ERROR皆是在执行被包裹语句后，立即读取CUDA错误信息，并对错误进行相应处理。

CUDAFunctions

基本结构

CUDAFunctions模块对CUDA的部分接口进行封装。这些接口包含设备信息、内存拷贝和流同步。CUDA Runtime函数的结果值返回通常采用传入引用的方式完成，那些函数本身的返回值是错误码。C10库将错误码统一处理，将外层调用者希望得到的结果通过函数返回的方式传递到外层。

算法设计

该模块基本只是对CUDA底层调用的一些封装。其中稍微有意思的是device count方法。

该方法最底层对接CUDA提供的cudaGetDeviceCount接口，并进行一些处理。

正常情况下，该函数应该返回一个不太大的数字（开发者认为，不可能有安装超过255个CUDA设备的主机）。但是，开发者们对这部分代码不自信。因此，它们将取设备个数的实现用try语句块包裹，检查得到的结果，并在结果超过255时，抛出一个异常，提醒外部开发者将这个Issue提交给PyTorch。

获取设备数量时，首先调用cudaGetDeviceCount。如果这样就成功了，直接返回结果就行；如果遇到错误，则需要做错误处理。如果错误由主机未安装CUDA设备产生，直接返回0，表示没有设备，即可；如果是CUDA驱动版本太低，则要提醒用户更新驱动程序；如果是驱动启动失败，则提醒用户这件事；如果是其他错误，似乎不好处理，但也要将这个unknown error通过报错信息告知用户。

CUDADeviceAssertionHost

基本结构

该模块提供多个与日志相关的结构。

DeviceAssertionData：该结构存储单条CUDA设备的检查（assert）报错信息。其使用的内存同时可被CPU和GPU访问。

DeviceAssertionsData：用于存储CUDA设备产生的所有检查（assert）日志。其内部是一个DeviceAssertionData列表和一个计数器。

CUDAKernelLaunchInfo：该结构存储单条CUDA内核启动信息。

CUDAKernelLaunchRegistry：该结构基于循环队列，存储CUDA内核启动信息，以及CUDA设备产生的检查日志。该结构对外提供错误信息检查等接口。该结构以单例模式运行。

算法设计

CUDAKernelLaunchRegistry为每条待插入数据分配一个id。这个id每次自增。通过对id取模，即可找到该记录应该插入到循环队列的哪个位置。

CUDAKernelLaunchRegistry基于懒汉式单例模式运行。借助C++的函数内static机制，在第一次需要使用时构建对象，并在后续使用时返回同一对象，以此实现单例。然而，开发者们并没有将该单例类的构造函数对外屏蔽，令人感到疑惑。

CUDACachingAllocator

基本结构

该模块主要完成内存管理的基础部分。模块内实现一个CUDA内存分配器的基础类，名为CUDAAllocator。该类基于C10/Allocator，提供内存分配、内存释放、初始化、日志等方法。

模块内存储一个供全局使用的内存分配器对象。该对象可以是Native Caching Allocator，也可以是CUDA Malloc Async Allocator。具体选择通过环境变量的设置来决定。

Native Caching Allocator直接在本模块内实现；CUDA Malloc Async Allocator在另一模块内实现。

此外，我们很惊喜地发现，模块内包含一个分配器配置文件解析器。它相当于一个简易编译器，带有很简单的词法分析和语义分析模块，用于解析文本格式的配置文件。

算法设计

模块内的最大亮点是完整实现的Device Caching Allocator和Native Caching Allocator。

Device Caching Allocator

该工具用于管理单个显卡设备的内存池。

外层开发时，可能产生大量针对小内存的申请与释放操作。如果每次都直接向CUDA Runtime申请，可能会非常耗时。Device Caching Allocator作为CUDA Runtime的内存管理器和外层程序之间的桥梁，每次申请较大的内存，分割后供外层程序使用。外层程序释放内存时，该部分内存不直接还给CUDA Runtime，而是暂时留在分配器内，下一个申请请求到来时可以快速响应。

分配器将小块内存和大块内存放到不同的位置管理。如果外层希望申请的内存非常小，小于某个阈值，分配器会直接返回一块更大的内存，以减少内存碎片问题；如果用户希望申请的内存特别大，则会对该大内存申请请求进行特殊处理。这个思路与Arena相似，但其空闲内存块管理算法与Arena完全不同。

当然，分配器并不是永远持有得到的内存。开发者们在该工具内实现一个简单的GC，它会对空闲内存块进行代数计数，将较老代的内存交还给CUDA Runtime。

当外部前来申请内存，分配器前往自己管理的内存池里寻找是否有空闲内存。如果有，登记后直接返回即可。若无，则会尝试启动GC，并向CUDA Runtime API申请新的GPU内存块。当然，向CUDA API申请内存不一定成功。如果GPU内存耗尽，分配器会向外报告一个OOM。

内存分配过程中，分配器可能会对内存块进行分割。释放时，分配器会尝试将内存块进行合并。内存块之间通过链表形式连接，切割时产生一个新的链表节点，合并时尝试将相邻节点合并。同时，内存块在内存池内使用C++ Set容器存放。我们认为，这个设计很巧妙。相比基于线性表、位图或树的内存管理形式，基于链表的管理形式仅使用前后索引即可完成定位，而不是使用地址和长度，算法被简化。由于底层采用的是基于红黑树的Set，内存块查改速度仅为对数级。

Native Caching Allocator

该工具基于Device Caching Allocator提供的功能，为外部程序提供GPU内存。它使用前者的接口获取并接管显存块，将申请得到的显存通过设备ID（本质为显存地址）形式交给外部调用者。外部可以凭借设备ID得到显存块描述结构。

CUDAMallocAsyncAllocator

基本结构

该模块在逻辑上与Native Caching Allocator属于同一层，但是它的底层不是Device Caching Allocator，而是直接调用CUDA Runtime的异步内存申请API。

借助异步机制，申请内存过程中，程序可以并行执行其他操作。等到内存申请成功，通过回调方式继续完成需要使用那些内存的操作。
几个分配器的关系如下图所示：

算法设计

该模块的实现没有太多特色。它直接操作CUDA Runtime API，并加以处理，以对外提供功能。

CUDAStream

基本结构

该模块是对CUDA执行流的一层封装，借助“池”技术避免流的反复创建销毁带来的开销。

为每个设备创建三个CUDA流池，分别负责存储默认优先级流、低优先级流和高优先级流。这些“流池”都采用懒汉法进行初始化。

算法设计

init device stream state

该方法用于为每个设备初始化“流池”。按照PyTorch 2.0的设定，在初始化时会为每台设备创建32个高优先级执行流和32个低优先级执行流。这些流在未来会按照循环的方式供外层使用。

get stream from pool

该方法用于从流池获取一个暂时未被使用的流。C10库的开发者希望调用该方法的开发者认为自己创建了一个新的流。然而在底层，这个流是预先统一创建好的，被用来循环使用。

方法允许调用者声明希望得到的流是高优先级的还是低优先级的。对于每种优先级，为每次创建的请求按照自增方式创建一个ID。通过对ID取模，即可实现流的循环使用。