PyTorch 的基本单元是张量。本文将概述如何在 PyTorch 中实现张量，以便用户可以从 Python shell 与之交互。特别是，我们想要回答以下四个主要问题：
- PyTorch 是如何扩展 Python 解释器以定义一个可以从 Python 代码中操作的张量类型的？
- PyTorch 是如何包装定义张量属性和方法的 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 张量。
关键部分是定义一个新的“类型对象”。例如，我们 Python float 的类型对象的定义形式如下：
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 张量支持 Python 的映射协议。这使得我们可以执行如下操作：
x = torch.Tensor(10).fill_(1)
y = x[3] // y == 1
x[4] = 2
// etc.
**请注意，这种索引扩展到多维 Tensor
我们可以通过定义这里描述的三个映射方法来使用 [] -风格的表示法。
其中最重要的方法是 THPTensor_(getValue) 和 THPTensor_(setValue) ，它们描述了如何索引 Tensor，以返回新的 Tensor/标量或就地更新现有 Tensor 的值。阅读这些实现以更好地理解 PyTorch 如何支持基本的张量索引。
通用构建（第一部分）
我们可以花很多时间探索 THPTensor 的各个方面以及它与定义新 Python 对象的关系。但我们仍然需要看看 THPTensor_(init)() 函数是如何翻译成我们在模块初始化中使用的 THPIntTensor_init() 的。我们如何使用定义“通用”张量的 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 的代码都将被替换为 float。在相应的文件 THGenerateIntType.h 中，相同的宏将用 int 替换 real 。
这些输出文件由 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 是根据当时符号 Real 的定义 #define ‘d 得到的。因为我们的头文件和源代码被所有类型的宏定义所包围，如上所示，在预处理程序运行后，生成的代码就是我们预期的代码。 TH 库中的代码为 THTensor_(NAME) 定义了相同的宏，支持这些函数的翻译。这样，我们就得到了包含专用代码的头文件和源文件。
模块对象和类型方法
现在我们已经看到，我们是如何将 TH 的 Tensor 定义封装在 THP 中，并生成 THP 方法如 THPFloatTensor_init(...) 的。现在我们可以探索上述代码实际上在创建的模块中做了什么。 THPTensor_(init) 中的关键行是：
# THPTensorBaseStr, THPTensorType are also macros that are specific
# to each type
PyModule_AddObject(module, THPTensorBaseStr, (PyObject *)&THPTensorType);
此函数将我们的 Tensor 对象注册到扩展模块中，这样我们就可以在 Python 代码中使用 THPFloatTensor、THPIntTensor 等。
只能创建张量还不够有用 - 我们需要能够调用 TH 定义的 所有方法。一个简单的例子是调用张量上的 in-place 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 张量等价物上调用所有 TH 张量的方法。然而，为所有 TH 方法编写包装器将非常耗时且容易出错。我们需要一种更好的方法来完成这项工作。
PyTorch cwrap
PyTorch 实现了自己的 cwrap 工具，用于包装 TH Tensor 方法以在 Python 后端使用。我们定义一个包含一系列 C 方法声明的自定义 YAML 格式的 .cwrap 文件。cwrap 工具接收此文件，并输出包含包装方法的源文件，这些文件与我们的 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 工具的架构非常简单。它读取一个文件，然后通过一系列插件进行处理。有关插件如何更改代码的所有方式的文档，请参阅 tools/cwrap/plugins/__init__.py 。
源代码生成发生在一系列步骤中。首先，解析并处理 YAML“声明”。然后逐步生成源代码 - 添加诸如参数检查和提取、定义方法头以及调用底层库的实际调用（如 TH ）。最后，cwrap 工具允许一次处理整个文件。 addmv_ 的结果输出可以在此处探索。
为了与 CPython 后端接口，该工具生成一个 PyMethodDef 数组，可以存储或追加到 THPTensor 的 tp_methods 字段。
在包装 Tensor 方法的具体情况下，构建过程首先从 TensorMethods.cwrap 生成输出源文件。这个源文件在通用的 Tensor 源文件中 #include 。所有这些都在预处理器施展魔法之前发生。因此，所有生成的包装器方法都经过与上面 THPTensor 代码相同的流程。因此，为每种类型都专门化了一个通用的声明和定义。
整合一切
到目前为止，我们已经展示了如何扩展 Python 解释器以创建一个新的扩展模块，该模块如何定义我们的新 THPTensor 类型，以及我们如何为与 TH 接口的所有类型的 Tensors 生成源代码。简要来说，我们将简要介绍编译过程。
Setuptools 允许我们定义一个用于编译的扩展。整个 torch._C 扩展通过收集所有源文件、头文件、库等并创建 setuptools Extension 来编译。然后 setuptools 负责构建扩展本身。我将在后续的文章中更深入地探讨构建过程。
总结一下，让我们重新审视我们的四个问题：
- PyTorch 是如何扩展 Python 解释器以定义一个可以从 Python 代码中操作的 Tensor 类型的？
它使用 CPython 的框架来扩展 Python 解释器并定义新类型，同时特别注意为所有类型生成代码。
- PyTorch 是如何包装定义 Tensor 属性和方法的 C 库的？
通过定义一个新的类型 THPTensor ，该类型由 TH 张量支持来实现。函数调用通过 CPython 后端的约定转发到这个张量。
- PyTorch 的 cwrap 是如何生成 Tensor 方法代码的？
它通过一系列步骤使用多个插件处理我们的自定义 YAML 格式代码，并为每个方法生成源代码。
- PyTorch 的构建系统是如何将这些组件编译并生成一个可工作的应用程序的？
使用 Setuptools 构建扩展需要一大堆源文件、头文件、库和编译指令。
这只是 PyTorch 构建系统的一部分快照。其中包含更多细微之处和细节，但我希望这能作为对我们张量库众多组件的温和介绍。
资源：
- https://docs.python.org/3.7/extending/index.html 对于理解如何编写 C/C++扩展到 Python 非常有价值。
