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 的 float 结构体定义如下：
typedef struct {
PyObject_HEAD
double ob_fval;
} PyFloatObject;
PyObject_HEAD 是一个宏，它引入了实现对象引用计数以及指向相应类型对象的指针的代码。因此，在这种情况下，要实现一个 float，唯一需要的其他“状态”就是浮点值本身。
现在，让我们看看我们的 THPTensor 类型的结构体：
struct THPTensor {
PyObject_HEAD
THTensor *cdata;
};
很简单，对吧？我们只是通过存储一个指向底层 TH tensor 的指针来对其进行封装。
关键部分是定义新类型的“类型对象”。我们的 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 在为 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 Tensors 支持 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/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 等。
光是能够创建 Tensors 并不十分有用——我们需要能够调用 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 中定义方法工作原理的更多信息会很有启发性。通常，方法将对象实例作为第一个参数，并可选地接收位置参数和关键字参数的参数。此静态函数被注册为我们的 float 上的一个方法：
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 后端中使用。我们定义一个 .cwrap 文件，其中包含一系列采用自定义 YAML 格式的 C 方法声明。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 类型，以及如何为与 TH 交互的所有类型 Tensor 生成源代码。接下来，我们将简要介绍编译。
Setuptools 允许我们定义一个用于编译的 Extension。整个 torch._C 扩展是通过收集所有源文件、头文件、库等并创建一个 setuptools Extension 来编译的。然后 setuptools 处理构建扩展本身。我将在后续文章中更深入地探讨构建过程。
总结一下，让我们回顾一下之前的四个问题：
- PyTorch 如何扩展 Python 解释器以定义可以从 Python 代码中操作的 Tensor 类型？
它利用 CPython 的框架来扩展 Python 解释器和定义新类型，同时特别注意为所有类型生成代码。
- PyTorch 如何封装实际定义 Tensor 属性和方法的 C 库？
它通过定义一个新的类型 THPTensor 来实现，该类型由 TH Tensor 支持。函数调用通过 CPython 后端的约定转发给此 Tensor。
- PyTorch 的 cwrap 工具如何工作来生成 Tensor 方法的代码？
它接收我们自定义的 YAML 格式代码，并通过使用多个插件的一系列步骤处理，为每个方法生成源代码。
- PyTorch 的构建系统如何整合所有这些组件来编译和生成可运行的应用程序？
它汇集了一系列源文件/头文件、库和编译指令，使用 Setuptools 构建扩展。
这只是 PyTorch 构建系统部分内容的概览。其中有更多的细微之处和细节，但我希望这能作为我们 Tensor 库许多组件的温和介绍。
资源
- https://docs.pythonlang.cn/3.7/extending/index.html 对于理解如何编写 Python 的 C/C++ 扩展非常有价值
