PyTorch 内部剖析（第一部分） – PyTorch

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 张量。

关键部分是定义新类型的“类型对象”。我们的 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));

这说明了我们需要进行类型特化的两种情况：

我们的输出代码将调用 THPTensor_New(...) 而不是 THPTensor_(New)
我们的输出代码将调用 THTensor_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.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 的代码都会将该 real 替换为 float。在相应的 THGenerateIntType.h 文件中，相同的宏会将 real 替换为 int。

这些输出文件从 split_types 返回并添加到源文件列表中，因此我们可以看到如何创建不同类型的 .cpp 代码。

这里有几点需要注意：首先，split_types 函数并非严格必要。我们可以将 Tensor.cpp 中的代码封装在一个文件中，为每种类型重复一遍。我们将代码拆分成单独的文件是为了加快编译速度。其次，当我们谈论类型替换（例如，将 real 替换为 float）时，我们的意思是 C 预处理器将在编译期间执行这些替换。仅仅用这些宏包围源代码在预处理之前没有副作用。

通用构建（第二部分）

现在我们已经有了所有 Tensor 类型的源文件，我们需要考虑如何创建相应的头文件声明，以及如何将 THTensor_(method) 和 THPTensor_(method) 转换为 THTensor_method 和 THPTensor_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 张量等效物上调用 TH 张量的所有方法。但是，为所有 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 后端接口，该工具生成一个 PyMethodDef 数组，该数组可以存储或附加到 THPTensor 的 tp_methods 字段。

在封装 Tensor 方法的特定情况下，构建过程首先从 TensorMethods.cwrap 生成输出源文件。此源文件包含在通用 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 对于理解如何为 Python 编写 C/C++ 扩展非常有价值

PyTorch 内部结构之旅（第一部分）