原版英文链接：Edward Z. Yang's PyTorch internals : Inside 245-5D
Mechanics
Code Flow
PyTorch源码仓库中包含众多的文件件，详细的介绍可以参考CONTRIBUTING。其中最重要的四个文件夹及源码模块为：
torch/: PyTorch模块库。包含PyTorch开发最常用的功能，通过import导入的预定义模块。PyTorch的Python前端(frontend)。
torch/csrc/: PyTorch前端模块的C++代码，实现C++代码与Python的绑定。此外包含自动求导引擎(autograd/)、JIT编译器(jit/)，以及PyTorch的C++前端(api/)。
aten: "A Tensor Library"的缩写，张量相关操作的实现，不包含自动求导功能。src/目录下包含两种实现：已经过时的c版本实现(TH/, THC/, THNN/, THCUNN/)，和基于C++实现(ATen/)。不同device上张量操作实现的方式分别位于不同文件夹(ATen/cpu, ATen/cuda, ATen/sparse, ...)
c10: Caffe2和ATen的混合缩写(caffe ten)，PyTorch的核心抽象和基础功能实现，包括张量的具体存储和实现方式，支持部署到服务器和移动端设别。PyTorch正在将ATen/core中的基础核心实现移植到c10/core。PyTorch的C++后端(backend)
核心模块支持上层的逻辑实现。以PyTorch的相加函数torch.add为例，模块间调用流程为：
- 将Python函数转换为C函数调用，解析Python参数为C++参数，由torch/csrc/中函数实现。例如，以PyTorch的
add函数为例(下列代码自动生成)：
// actual binding
static PyMethodDef torch_functions[] ={
...
{"add", (PyCFunction)THPVariable_add, METH_VARARGS | METH_VARKEYWORDS | METH_STATIC, NULL}
...
}
// auto-generated codes, needed to build PyTorch to generate it
static PyObject* torch._C.VariableFunctions.add.THPVariable_add(
PyObject* self_, PyObject* args, PyObject* kwargs){
static PythonArgParser parser(...);
ParsedArgs<4> parsed_args;
auto r = parser.parse(args, kwargs, parsed_args);
...
if(r.isNone(3)){
return wrap(dispatch_add(r.tensor(0), r.tensor(1), r.scalar(2)));
}else{
return wrap(dispatch_add(r.tensor(0), r.tensor(1), r.scalar(2), r.tensor(3)));
}
...
}
torch_functions定义了Python函数和C版本函数名称的对应关系，通过该映射表查询对应的C版本函数。PythonArgParser类实现了对PyTorch参数的C版本解析，然后通过dispatch_add调度底层C版本的add实现。计算结果通过wrap重新包装为PyObject对象，返回给Python层。
2.变量类型的调度
上一步中的dispathc_add函数调度实际上调用self.add(tensor, scalar)函数，即张量自身的实现版本，而该版本通过如下函数实现，调用不同变量类型的实现版本。
// inline functions defined on the 'type'
inline Tensor Tensor::add(const Tensor& other, Scalar alpha) const {
return type().add(*this, other, alpha);
}
函数type()确定变量的具体类型，并通过虚函数add调度实际类型的实现。这些处理在aten/src/ATen模块中完成。
3.设别类型和布局的调度
type()实际同时完成了变量和设备类型的调度，返回类似TypeDefault、GPUFloatType等包括数据类型和设备类型的描述。针对每种类型，PyTorch在build后会生成具体的类似如下的实现代码:
Tensor TypeDefault:add(const Tensor& self, const Tensor& other, Scalar alpha) const {
const OptionalDeviceGuard device_guard(device_of(self)) # device type checking
return at::native::add(self, other, alpha) # modern c++ impl.
}
由于add函数对于不同类型变量及设备类型的底层实现相同，通可以过TypeDefault统一封装。如果某种计算操作有不同实现，则需要扩展实现并调用对应版本，类似于GPUFloatType::add(...)。
4.核心代码的调用
第三步中的代码封装了更为底层的at::native::add(self, other, alpha)实现。这些实现依赖aten/src/ATen中的模块，通过C++版本(native/)或者过时的c版本实现(TH/, THC/, THNN/, THCUNN/)。
Kernels
PyTorch提供了一些工具和规范用于开发核心计算操作符，由aten/src/ATen模块支持。一段完整的自定义核心操作示例代码所示：
Tensor my_op(Tensor& result, const Tensor& self, const Tensor& other){
// error checking
TORCH_CHECK(result.is_cpu() && self.is_cpu() && other.is_cpu());
TORCH_CHECK(self.dim() == 1);
TORCH_CHECK(self.sizes() == other.sizes());
// output allocation
result.resize_(self.sizes());
// data type (dtype) dispatch
AT_DISPATCH_FORALL_TYPES(
self.scalar_type(), "my_op", [&]{
my_op_cpu<scalar_t>(result, self, other),
}
);
}
template<typename scalar_t>
void my_op_cpu(Tensor& result, const Tensor& self, const Tensor& other){
// data access
auto result_accessor = result.accessor<scalar_t, 1>();
auto self_accessor = self.accessor<scalar_t, 1>();
auto other_accessor = other.accessor<scalar_t, 1>();
// parallelization
parallel_for(0, self.size(0), 0, [&](int64_t start, int64_t end){
... self_accessor[i] ...
});
}
包含如下几个部分：
元数据注册：由PyTorch提供的元数据要求，用于自动化生成Python的绑定代码(如上节介绍的Python与C代码之间的转换和参数解析)。每个定义的核心操作都需要提供如下的元数据模式：
- func: func_name(ArgType arg0, ArgType arg1, ...) -> Return
variants: function, method
dispatch:
CPU: func_cpu
CUDA: func_cuda
其中：
func_name：所定义的核心计算操作函数的名称。
ArgType：参数类型，可以是Tensor, Tensor[], int, int[], float, Scalar等。
variants: 包含function和method两个类型，用于控制PyTorch自动生成Python版本函数的名称是张量方法(t.foo())还是命名空间的函数(at::foo())。当使用method变体时，需要包含self参数。在自动生成Python版本函数名称时，该self参数会从参数列表中去掉。例如对于where(BoolTensor cond, Tensor self, Tensor other)的函数声明。设置为method会自动生成self.where(cond, other)的函数名称；设置为function会自动生成at::where(cond, self, other)的函数名称。缺省情况下，ATen对native函数只生成function方式名称，对张量相关的核心操作符(e.g, add, sub等)可以使用method方式名称。
dispatch: 指定针对不同设别类型，该函数可以调度的实际函数名称。可以针对不同设别类型，指定生成不同的版本的函数名称。
更详细的规范要求可参考aten/src/ATen/native/README.md。任何自定义的核心计算函数的元数据需要按照如上要求编写，并添加到native_functions.yaml
文件中进行注册。PyTorch会对注册的函数按照元数据描述的要求，自动生成Python的绑定。
上述自定义的核心操作函数，一种可能的元数据描述如下：
-func: my_op(Tensor& result, const Tensor& self, const Tensor& other) -> Tensor
variants: function, method
dispath:
CPU: my_op_cpu
CUDA: my_op_cuda
对于需要支持反向梯度计算的核心计算操作，需要按照类似的方式提供求导操作函数的元数据。具体可参考derivatives.yaml
错误检测(Error Checking)
错误检查在编写核心代码时非常重要。PyTorch提供了两种错误检查的工具方便开发者：low level的方式是提供了TORCH_CHECK宏；High level方式通过将Tensor封装为TensorArg，并提供checkDim等检测函数。
输出存储分配(Output Allocation)
在输出结果前，需要预先分配内存用于存储。PyTorch支持预分配输出、原位输出、拷贝输出等方式输出结果。实现过程中，原位输出和拷贝输出只是预分配输出的简单封装。例如
// pre-allocate storage for 'result' outside
Tensor& abs_out(Tensor& result, const Tensor& self){
result.resize_(self.sizes());
// ... the real impl.
}
// a new allocated operation
Tensor& abs(const Tensor& self){
Tensor result = at::empty({0}, self.options());
abs_out(result, self);
return result
}
// in-place operation
Tensor& abs_(const Tensor& self){
return abs_out(self, self);
}
数据类型调度(Dtype Dispatch)
通过AT_DISPATCH_ALL_TYPES宏定义数据类型调度。该宏其实是一个模版函数，通过变量当前类型进行特化，调度实际匹配的类型实现。
数据访问(Data Access)
PyTorch支持三种不同的、针对张量的访问方式。封装的访问方式比访问原始数据指针方便，可以自动处理底层的stride或者布局。TensorAccesor支持访问张量某个特定位置的数据；TensorIterator支持规则方式轮询访问；针对CPU的序列化，提供了Vec256等序列化描述访问。
Notes on PyTorch Internals系列文章
Notes on PyTorch Internals I
Notes on PyTorch Internals II
Notes on PyTorch Internals III
