pytorch的c++扩展官方文档阅读笔记(一)

首先给出官方文档的链接

Custom C++ and CUDA Extensions — PyTorch Tutorials 2.3.0+cu121 documentation

一般来说用cpp扩展pytorch主要是为了添加一些算子或者其他运算支持。首先推荐的方法是使用python的方式实现需求，继承Function类和Module类实现相应的运算。官方文档给了两种情况下需要使用c++扩展的例子，一个是运算速度要求非常高的场景，一个是开发过程中需要用到一些cpp依赖库或者环境。

pytorch扩展的设计是尽可能与后端分离，这个目的主要是方便开发者快速应用到项目，但是如果想要将自己的扩展写到原生的pytorch后端中，也可以向项目贡献代码并重新组织代码结构。

c++扩展其实有两种方式，一个是通过setuptools去提前编译好，另一种可以用torch.utils.cpp_extension.load()去即时编译(just in time)。

文档中提到了一个cuda-11.5 nvcc的编译错误，需要用<ATen/Aten.h>代替<torch/extension.h>，后者是一个一站式头文件，包含了pytorch后端绑定python的所有内容。

文档提到：c++的扩展并不会自动生成反向传播函数，所以需要为算子自定义backward函数，最终用torch.autograd.Function去封装。

文档说到：在cpp前项和方向函数使用pybind11绑定python的部分，有一个变量明明为TORCH_EXTENSION_NAME，这个其实是python中对应的模块名，在setup.py中定义，那么这里为了避免开发者两处定义冲突，采用了宏定义，运行setup的时候会自动展开为setup.py中定义的模块名。

文档提到：pytorch编译器和cpp扩展编译器必须ABI兼容，否则可能需要相同的编译器编译pytorch和扩展。编译器的ABI（应用程序二进制接口，Application Binary Interface）问题主要涉及到不同编译器生成的二进制代码之间的兼容性。ABI是一种规范，它定义了编译后的程序如何与操作系统、库以及其他二进制组件交互的规则，包括但不限于调用约定、数据类型大小和对齐、名字修饰（Name Mangling）、运行时支持等。具体到编译器的ABI问题，主要体现在调用约定不一致、数据类型表示差异、名称修饰、运行时行为、字节序和对齐等方面。解决这些问题通常需要确保编译和链接过程中使用的工具链（包括编译器、链接器、库）与目标运行环境的ABI相匹配，或者在必要时使用兼容性层来桥接不同编译器产生的代码。在某些情况下，开发者可能需要手动调整代码或编译选项以确保跨编译器的兼容性。

文档提到：使用cpp重写前向推理过程，可以使得前向有小幅度的速度提升，但是对backwards速度提升是有限的，这主要是因为自动求导机制对计算图有并行优化。

文档中提到了用jit方式编译cpp文件的方法，用以下代码替换setup.py，后续的python调用可以保持不变。

from torch.utils.cpp_extension import load

lltm_cpp = load(name="lltm_cpp", sources=["lltm.cpp"])

这种方式第一次运行python会慢一些，但是后面不会增加额外的时间开销。