PyTorch 的基本单元是 Tensor（张量）。本文将概述我们如何在 PyTorch 中实现 Tensor，以便用户能够通过 Python shell 与其进行交互。我们主要想要回答以下四个问题：
- PyTorch 如何扩展 Python 解释器以定义一个可以在 Python 代码中操作的 Tensor 类型？
- PyTorch 如何封装那些实际定义了 Tensor 属性和方法的 C 语言库？
- PyTorch 的 cwrap 工具是如何为 Tensor 方法生成代码的？
- PyTorch 的构建系统如何整合所有这些组件来编译并生成一个可用的应用程序？
扩展 Python 解释器
PyTorch 定义了一个名为 torch 的新包。在本文中，我们将讨论 ._C 模块。该模块被称为“扩展模块”——即用 C 语言编写的 Python 模块。此类模块允许我们定义新的内置对象类型（例如 Tensor）并调用 C/C++ 函数。
._C 模块定义在 torch/csrc/Module.cpp 中。init_C() / PyInit__C() 函数负责创建该模块并按需添加方法定义。该模块被传递给多个不同的 __init() 函数，这些函数负责向模块添加更多对象、注册新类型等。
其中一个 __init() 调用集合如下：
ASSERT_TRUE(THPDoubleTensor_init(module));
ASSERT_TRUE(THPFloatTensor_init(module));
ASSERT_TRUE(THPHalfTensor_init(module));
ASSERT_TRUE(THPLongTensor_init(module));
ASSERT_TRUE(THPIntTensor_init(module));
ASSERT_TRUE(THPShortTensor_init(module));
ASSERT_TRUE(THPCharTensor_init(module));
ASSERT_TRUE(THPByteTensor_init(module));
这些 __init() 函数将每种类型的 Tensor 对象添加到 ._C 模块中，以便它们可以在该模块中使用。让我们来了解这些方法是如何工作的。
THPTensor 类型
正如底层的 TH 和 THC 库一样，PyTorch 定义了一个“通用”Tensor，然后针对多种不同类型进行了专门化。在考虑这种专门化如何工作之前，我们先看看如何在 Python 中定义新类型，以及如何创建通用的 THPTensor 类型。
Python 运行时将所有 Python 对象视为 PyObject * 类型的变量，它充当所有 Python 对象的“基类”。每个 Python 类型都包含对象的引用计数，以及指向该对象类型对象的指针。类型对象决定了该类型的属性。例如，它可能包含与该类型关联的方法列表，以及实现这些方法所需的 C 函数。对象还包含表示其状态所需的任何字段。
定义新类型的公式如下：
- 创建一个结构体，定义新对象将包含的内容
- 定义该类型的类型对象
结构体本身可以非常简单。在 Python 中，所有的浮点类型实际上都是堆上的对象。Python 的浮点结构体定义如下：
typedef struct {
PyObject_HEAD
double ob_fval;
} PyFloatObject;
PyObject_HEAD 是一个宏，它引入了实现对象引用计数以及指向相应类型对象的指针的代码。因此，在这种情况下，要实现一个浮点数，唯一需要的其他“状态”就是浮点值本身。
现在，让我们看看 THPTensor 类型的结构体：
struct THPTensor {
PyObject_HEAD
THTensor *cdata;
};
很简单，对吧？我们只是通过存储一个指向底层 TH tensor 的指针来对其进行封装。
关键部分是定义新类型的“类型对象”。我们 Python 浮点数的类型对象定义示例如下：
static PyTypeObject py_FloatType = {
PyVarObject_HEAD_INIT(NULL, 0)
"py.FloatObject", /* tp_name */
sizeof(PyFloatObject), /* tp_basicsize */
0, /* tp_itemsize */
0, /* tp_dealloc */
0, /* tp_print */
0, /* tp_getattr */
0, /* tp_setattr */
0, /* tp_as_async */
0, /* tp_repr */
0, /* tp_as_number */
0, /* tp_as_sequence */
0, /* tp_as_mapping */
0, /* tp_hash */
0, /* tp_call */
0, /* tp_str */
0, /* tp_getattro */
0, /* tp_setattro */
0, /* tp_as_buffer */
Py_TPFLAGS_DEFAULT, /* tp_flags */
"A floating point number", /* tp_doc */
};
理解类型对象最简单的方法是将其视为定义对象属性的一组字段。例如，tp_basicsize 字段被设置为 sizeof(PyFloatObject)。这是为了让 Python 在调用 PyObject_New() 创建 PyFloatObject 时知道需要分配多少内存。你可以设置的完整字段列表定义在 CPython 后端的 object.h 中：https://github.com/python/cpython/blob/master/Include/object.h。
THPTensor 的类型对象是 THPTensorType，定义在 csrc/generic/Tensor.cpp 中。该对象定义了 THPTensor 的名称、大小、映射方法等。
举个例子，让我们看看我们在 PyTypeObject 中设置的 tp_new 函数：
PyTypeObject THPTensorType = {
PyVarObject_HEAD_INIT(NULL, 0)
...
THPTensor_(pynew), /* tp_new */
};
tp_new 函数支持对象创建。它负责创建（而不是初始化）该类型的对象，等同于 Python 层的 __new__() 方法。其 C 语言实现是一个静态方法，接收正在实例化的类型和任何参数，并返回一个新创建的对象。
static PyObject * THPTensor_(pynew)(PyTypeObject *type, PyObject *args, PyObject *kwargs)
{
HANDLE_TH_ERRORS
Py_ssize_t num_args = args ? PyTuple_Size(args) : 0;
THPTensorPtr self = (THPTensor *)type->tp_alloc(type, 0);
// more code below
我们这个新函数做的第一件事是分配 THPTensor。然后，它根据传递给函数的参数执行一系列初始化。例如，当从另一个 THPTensor y 创建一个 THPTensor x 时，我们设置新创建的 THPTensor 的 cdata 字段，使其结果为调用 THTensor_(newWithTensor) 并将 y 的底层 TH Tensor 作为参数。对于大小、存储、NumPy 数组和序列，也存在类似的构造函数。
** 请注意，我们仅使用 tp_new，而不结合使用 tp_new 和 tp_init（后者对应于 __init__() 函数）。
在 Tensor.cpp 中定义的另一个重要内容是索引的工作方式。PyTorch Tensor 支持 Python 的映射协议 (Mapping Protocol)。这允许我们执行如下操作：
x = torch.Tensor(10).fill_(1)
y = x[3] // y == 1
x[4] = 2
// etc.
** 请注意，此索引方式延伸到了多维 Tensor。
我们通过定义此处描述的三个映射方法，从而能够使用 [] 风格的符号。
最重要的方法是 THPTensor_(getValue) 和 THPTensor_(setValue)，它们描述了如何索引 Tensor，以便返回一个新的 Tensor/标量，或原地更新现有 Tensor 的值。请通读这些实现，以更好地理解 PyTorch 如何支持基本的 Tensor 索引。
通用构建（第一部分）
我们可以花大量时间探索 THPTensor 的各个方面及其与定义新 Python 对象的关系。但我们还需要了解 THPTensor_(init)() 函数是如何转换为我们在模块初始化中使用的 THPIntTensor_init() 的。我们如何获取定义了“通用”Tensor 的 Tensor.cpp 文件，并利用它为所有类型排列生成 Python 对象？换句话说，Tensor.cpp 中充斥着类似这样的代码行：
return THPTensor_(New)(THTensor_(new)(LIBRARY_STATE_NOARGS));
这展示了我们需要针对特定类型处理的两种情况：
- 我们的输出代码将调用
THP<Type>Tensor_New(...)来替代THPTensor_(New) - 我们的输出代码将调用
TH<Type>Tensor_new(...)来替代THTensor_(new)
换句话说，对于所有支持的 Tensor 类型，我们需要“生成”完成了上述替换的源代码。这是 PyTorch “构建”过程的一部分。PyTorch 依赖 Setuptools (https://setuptools.readthedocs.io/en/latest/) 来构建包，我们在顶层目录定义了一个 setup.py 文件来自定义构建过程。
使用 Setuptools 构建扩展模块的一个组件是列出编译涉及的源文件。然而，我们的 csrc/generic/Tensor.cpp 文件并没有列出！那么该文件中的代码是如何成为最终产品的一部分的呢？
回想一下，我们是从 generic 的上一级目录调用 THPTensor* 函数（例如 init）的。如果我们查看该目录，会发现定义了另一个文件 Tensor.cpp。该文件的最后一行很重要：
//generic_include TH torch/csrc/generic/Tensor.cpp
请注意，这个 Tensor.cpp 文件被包含在 setup.py 中，但它被包裹在一个名为 split_types 的 Python 辅助函数调用中。该函数接收一个文件作为输入，并查找文件内容中的 “//generic_include” 字符串。如果找到，它将为每种 Tensor 类型生成一个新的输出文件，并进行以下更改：
- 输出文件被重命名为
Tensor<Type>.cpp - 输出文件进行了如下轻微修改：
# Before:
//generic_include TH torch/csrc/generic/Tensor.cpp
# After:
#define TH_GENERIC_FILE "torch/src/generic/Tensor.cpp"
#include "TH/THGenerate<Type>Type.h"
在第二行包含头文件会产生包含 Tensor.cpp 中源代码的副作用，并带有额外的上下文定义。让我们看看其中一个头文件：
#ifndef TH_GENERIC_FILE
#error "You must define TH_GENERIC_FILE before including THGenerateFloatType.h"
#endif
#define real float
#define accreal double
#define TH_CONVERT_REAL_TO_ACCREAL(_val) (accreal)(_val)
#define TH_CONVERT_ACCREAL_TO_REAL(_val) (real)(_val)
#define Real Float
#define THInf FLT_MAX
#define TH_REAL_IS_FLOAT
#line 1 TH_GENERIC_FILE
#include TH_GENERIC_FILE
#undef accreal
#undef real
#undef Real
#undef THInf
#undef TH_REAL_IS_FLOAT
#undef TH_CONVERT_REAL_TO_ACCREAL
#undef TH_CONVERT_ACCREAL_TO_REAL
#ifndef THGenerateManyTypes
#undef TH_GENERIC_FILE
#endif
这样做是为了引入通用 Tensor.cpp 文件中的代码，并用以下宏定义将其包围。例如，我们将 real 定义为 float，因此通用 Tensor 实现中任何引用 real 的代码都会将 real 替换为 float。在相应的 THGenerateIntType.h 文件中，相同的宏会将 real 替换为 int。
这些输出文件由 split_types 返回并添加到源文件列表中，因此我们可以看到不同类型的 .cpp 代码是如何创建的。
这里有几点需要注意：首先，split_types 函数并非严格必要。我们可以将 Tensor.cpp 中的代码封装在一个文件中，并为每种类型重复一遍。将代码拆分为单独文件的原因是加快编译速度。其次，当我们谈论类型替换（例如将 real 替换为 float）时，是指 C 预处理器将在编译期间执行这些替换。仅仅用这些宏包围源代码在预处理之前不会产生任何副作用。
通用构建（第二部分）
现在我们有了所有 Tensor 类型的源文件，我们需要考虑相应的头文件声明是如何创建的，以及从 THTensor_(method) 和 THPTensor_(method) 到 TH<Type>Tensor_method 和 THP<Type>Tensor_method 的转换是如何工作的。例如，csrc/generic/Tensor.h 中有如下声明：
THP_API PyObject * THPTensor_(New)(THTensor *ptr);
对于头文件的代码生成，我们使用了相同的策略。在 csrc/Tensor.h 中，我们执行以下操作：
#include "generic/Tensor.h"
#include <TH/THGenerateAllTypes.h>
#include "generic/Tensor.h"
#include <TH/THGenerateHalfType.h>
这具有相同的效果，我们为每种类型从通用头文件中提取代码，并用相同的宏定义进行包装。唯一的区别是，生成的代码全部包含在同一个头文件中，而不是拆分为多个源文件。
最后，我们需要考虑如何“转换”或“替换”函数类型。如果我们查看同一个头文件，会看到一堆 #define 语句，包括：
#define THPTensor_(NAME) TH_CONCAT_4(THP,Real,Tensor_,NAME)
这个宏表示源代码中任何匹配 THPTensor_(NAME) 格式的字符串都应替换为 THPRealTensor_NAME，其中 Real 取决于当时 #define 的符号 Real 是什么。因为我们的头文件代码和源文件代码被如上所述的所有类型的宏定义所包围，所以在预处理器运行后，生成的代码正是我们所期望的。TH 库中的代码也为 THTensor_(NAME) 定义了相同的宏，同样支持这些函数的转换。通过这种方式，我们最终得到了包含专门代码的头文件和源文件。
模块对象和类型方法
现在我们已经了解了如何将 TH 的 Tensor 定义封装在 THP 中，并生成了诸如 THPFloatTensor_init(...) 之类的 THP 方法。现在我们可以探讨上述代码在我们正在创建的模块中实际做了什么。THPTensor_(init) 中的关键行是：
# THPTensorBaseStr, THPTensorType are also macros that are specific
# to each type
PyModule_AddObject(module, THPTensorBaseStr, (PyObject *)&THPTensorType);
此函数将我们的 Tensor 对象注册到扩展模块，以便我们可以在 Python 代码中使用 THPFloatTensor、THPIntTensor 等。
仅仅能够创建 Tensor 是不够的——我们需要能够调用 TH 定义的所有方法。一个简单的示例展示了在 Tensor 上调用原地 zero_ 方法。
x = torch.FloatTensor(10)
x.zero_()
让我们从如何向新定义的类型添加方法开始。类型对象中的字段之一是 tp_methods。此字段持有一组方法定义（PyMethodDef），用于将方法（及其底层的 C/C++ 实现）与类型关联起来。假设我们想在 PyFloatObject 上定义一个新方法来替换该值。我们可以按如下方式实现：
static PyObject * replace(PyFloatObject *self, PyObject *args) {
double val;
if (!PyArg_ParseTuple(args, "d", &val))
return NULL;
self->ob_fval = val;
Py_RETURN_NONE
}
这等同于 Python 方法：
def replace(self, val):
self.ob_fval = val
阅读更多关于 CPython 中如何定义方法的内容很有启发性。通常，方法将对象的实例作为第一个参数，并可选择位置参数和关键字参数。此静态函数被注册为我们浮点数的一个方法：
static PyMethodDef float_methods[] = {
{"replace", (PyCFunction)replace, METH_VARARGS,
"replace the value in the float"
},
{NULL} /* Sentinel */
}
这注册了一个名为 replace 的方法，由同名的 C 函数实现。METH_VARARGS 标志表示该方法接收一个包含函数所有参数的元组。该数组被设置为类型对象的 tp_methods 字段，然后我们就可以在该类型的对象上使用 replace 方法了。
我们希望能够在我们的 THP Tensor 等价物上调用 TH Tensor 的所有方法。然而，为所有 TH 方法编写包装器既耗时又容易出错。我们需要一种更好的方法来做到这一点。
PyTorch cwrap
PyTorch 实现了自己的 cwrap 工具，用于封装 TH Tensor 方法，以便在 Python 后端中使用。我们在自定义的 YAML 格式中定义了一个包含一系列 C 方法声明的 .cwrap 文件。cwrap 工具获取该文件并输出 .cpp 源文件，其中包含以与我们的 THPTensor Python 对象及 Python C 扩展方法调用格式兼容的格式封装好的方法。该工具不仅用于生成封装 TH 的代码，也用于封装 CuDNN。它被设计为可扩展的。
原地 addmv_ 函数的一个 YAML “声明”示例如下：
[[
name: addmv_
cname: addmv
return: self
arguments:
- THTensor* self
- arg: real beta
default: AS_REAL(1)
- THTensor* self
- arg: real alpha
default: AS_REAL(1)
- THTensor* mat
- THTensor* vec
]]
cwrap 工具的架构非常简单。它读取一个文件，然后通过一系列插件 (plugins) 处理它。有关插件如何改变代码的所有方法的文档，请参阅 tools/cwrap/plugins/__init__.py。
源代码生成通过一系列过程发生。首先，解析并处理 YAML “声明”。然后，源代码被逐段生成——添加参数检查和提取、定义方法头，以及实际调用底层库（如 TH）。最后，cwrap 工具允许一次性处理整个文件。addmv_ 的最终输出可以在此处查看。
为了与 CPython 后端进行接口对接，该工具生成了一个 PyMethodDef 数组，可以存储或附加到 THPTensor 的 tp_methods 字段中。
在封装 Tensor 方法的特定情况下，构建过程首先从 TensorMethods.cwrap 生成输出源文件。此源文件被 #include 到通用 Tensor 源文件中。这一切都发生在预处理器施展魔法之前。因此，所有生成的包装方法都会经历与上述 THPTensor 代码相同的过程。因此，单个通用声明和定义也针对每种类型进行了专门化。
总结
到目前为止，我们已经展示了如何扩展 Python 解释器以创建新的扩展模块，该模块如何定义我们的新 THPTensor 类型，以及我们如何为所有与 TH 交互的类型生成 Tensor 源代码。简要地，我们将触及编译。
Setuptools 允许我们定义一个用于编译的扩展 (Extension)。整个 torch._C 扩展是通过收集所有源文件、头文件、库等并创建一个 setuptools Extension 来编译的。然后，setuptools 处理扩展本身的构建。我将在后续文章中进一步探讨构建过程。
总而言之，让我们回顾一下提出的四个问题：
- PyTorch 如何扩展 Python 解释器以定义一个可以在 Python 代码中操作的 Tensor 类型？
它使用了 CPython 扩展 Python 解释器和定义新类型的框架，同时特别注意为所有类型生成代码。
- PyTorch 如何封装那些实际定义了 Tensor 属性和方法的 C 语言库？
它通过定义一个由 TH Tensor 支持的新类型 THPTensor 来实现。函数调用通过 CPython 后端的约定转发给此 Tensor。
- PyTorch 的 cwrap 工具是如何为 Tensor 方法生成代码的？
它采用我们自定义的 YAML 格式代码，通过使用多个插件的一系列步骤处理它，为每个方法生成源代码。
- PyTorch 的构建系统如何整合所有这些组件来编译并生成一个可用的应用程序？
它获取一堆源文件/头文件、库和编译指令，使用 Setuptools 构建一个扩展。
这只是 PyTorch 构建系统部分内容的缩影。还有更多的细微差别和细节，但我希望这能作为一个温和的介绍，让你了解我们 Tensor 库的许多组件。
资源：
- https://docs.pythonlang.cn/3.7/extending/index.html 对于了解如何编写 Python 的 C/C++ 扩展非常有价值。