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 在为 PyFloatObject 调用 PyObject_New() 时就知道要分配多少内存。你可以设置的完整字段列表在 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 的**映射协议**。这使我们能够执行以下操作：
x = torch.Tensor(10).fill_(1)
y = x[3] // y == 1
x[4] = 2
// etc.
** 请注意，此索引扩展到具有多个维度的 Tensor
我们通过定义此处描述的三个映射方法，可以使用 [] 风格的表示法。
最重要的方法是 THPTensor_(getValue) 和 THPTensor_(setValue)，它们描述了如何索引 Tensor，以返回新的 Tensor/Scalar，或原地更新现有 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。此字段存储方法定义数组 (PyMethodDefs)，用于将方法（及其底层 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 工具的架构非常简单。它读取一个文件，然后通过一系列**插件**对其进行处理。有关插件如何更改代码的所有方式的文档，请参阅 tools/cwrap/plugins/__init__.py。
源代码生成通过一系列步骤进行。首先，解析并处理 YAML“声明”。然后，逐块生成源代码——添加参数检查和提取等内容，定义方法头，以及实际调用底层库（如 TH）。最后，cwrap 工具允许一次性处理整个文件。addmv_ 的最终输出可以在此处查看。
为了与 CPython 后端交互，该工具会生成一个 PyMethodDefs 数组，该数组可以存储或附加到 THPTensor 的 tp_methods 字段。
在包装 Tensor 方法的特定情况下，构建过程首先从 TensorMethods.cwrap 生成输出源文件。此源文件被 #include 到通用 Tensor 源文件中。所有这些都发生在预处理器发挥作用之前。因此，所有生成的方法包装器都经过与上述 THPTensor 代码相同的处理。因此，单个通用声明和定义也针对每种类型进行了专门化。
整合所有内容
到目前为止，我们已经展示了如何扩展 Python 解释器以创建新的扩展模块，该模块如何定义我们的新 THPTensor 类型，以及如何为所有类型的 Tensor 生成与 TH 交互的源代码。接下来，我们将简要介绍编译。
Setuptools 允许我们定义一个用于编译的扩展。整个 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 对于理解如何编写 C/C++ 扩展到 Python 来说是无价的。