pytorch源码：C拓展
Posted 你的奋斗
tags:
篇首语：本文由小常识网(cha138.com)小编为大家整理，主要介绍了pytorch源码：C拓展相关的知识，希望对你有一定的参考价值。
读了pytorch的Python部分源码，不断追溯代码，很多类都会继承“_C”模块里的内容，如：
class IntTensor(_C.IntTensorBase, _TensorBase)
def is_signed(self):
return True
@classmethod
def storage_type(cls):
return IntStorage
其中TensorBase为Python类，定义了多种Tensor类的共同操作。本文重点看C.IntTensorBase。主要尝试回答自己的两个疑问：
C扩展中的各种Tensor是如何定义和实现的。
上层的Python是如何调用C中定义的类（或结构体）。
事实证明，看似简单的两个问题并不那么简单，直接读Pytorch中C的源码完全不知所云，查了一些资料才明白自己严重需要预备知识：
Python如何拓展C库
Python的实现机制
这里跳过预备知识直入正题，相关链接整理在文章末尾。为了解答自己的疑惑，同样结合一个简单的例子：
t = torch.IntTensor()
下面尝试说明这条语句的背后底层做了哪些事情。再正式开始之前，需要先搞明白pytorch C拓展部分的代码生成套路。
Python C拓展
和普通的CPython方式拓展类似，下面主要看pytorch中的拓展模块“_C"的定义和相应其他模块的添加方式。pytorch中的拓展模块定义代码主要在torch/csrc/Module.cpp中，在预备知识中熟悉Python如何拓展C后，直接在Module.cpp找到感兴趣的初始化部分：
#if PY_MAJOR_VERSION == 2
PyMODINIT_FUNC init_C()
#else
PyMODINIT_FUNC PyInit__C()
#endif
{
...
#if PY_MAJOR_VERSION == 2
ASSERT_TRUE(module = Py_InitModule("torch._C", methods.data()));
#else
static struct PyModuleDef torchmodule = {
PyModuleDef_HEAD_INIT,
"torch._C",
NULL,
-1,
methods.data()
};
ASSERT_TRUE(module = PyModule_Create(&torchmodule));
#endif
...
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));
...
}
在编译过程中PyMODINITFUNC方法被调用，完成了"torch.C"的定义，接着就是各种类型Tensor的初始化函数调用，该部分在后面详细来看。 和普通C拓展套路一致，最终在编译阶段的setup.py文件中,声明Extension 执行setup加入拓展和用到的lib:
C = Extension("torch._C",
libraries=main_libraries,
sources=main_sources,
language='c++',
extra_compile_args=main_compile_args + extra_compile_args,
include_dirs=include_dirs,
library_dirs=library_dirs,
extra_link_args=extra_link_args + main_link_args + [make_relative_rpath('lib')],
)
extensions.append(C)
...
...
setup(name="torch", version=version,
description="Tensors and Dynamic neural networks in Python with strong GPU acceleration",
ext_modules=extensions,
cmdclass=cmdclass,
packages=packages,
package_data={'torch': [
'lib/*.so*', 'lib/*.dylib*',
'lib/torch_shm_manager',
'lib/*.h',
'lib/include/TH/*.h', 'lib/include/TH/generic/*.h',
'lib/include/THC/*.h', 'lib/include/THC/generic/*.h',
'lib/include/ATen/*.h',
]},
install_requires=['pyyaml', 'numpy'],
)
代码生成
一开始阅读代码，并没有发现IntTensor部分的实现，如上面PyMODINITFUNC函数中的THPIntTensorinit(module)，和最开始提到的IntTensorBase。另一个奇怪的是很多文件同时在根目录（pytorch/torch/csrc/）和generic目录下出现。以根目录下Tensor.cpp为例：
#include <Python.h>
#include <structmember.h>
#define THP_HOST_HALF
#include <stdbool.h>
#include <vector>
#include <stack>
#include <tuple>
#include <TH/THMath.h>
#include "torch/csrc/THP.h"
#include "torch/csrc/copy_utils.h"
#include "torch/csrc/DynamicTypes.h"
//generic_include TH torch/csrc/generic/Tensor.cpp
这个cpp文件和一般的cpp文件不同，注意最后一句，在pytorch的setup.py中准备source列表时会调用：
main_sources += split_types("torch/csrc/Tensor.cpp")
其中splittypes在pytorch/tools/setuphelpers/目录下实现，主要实现了两个功能：
重命名Tensor.cpp为Tensor[Type].cpp，其中Type为Float、Int等
各个Tensor[Type].cpp内容的最后一行内容改变为：
#define TH_GENERIC_FILE "torch/src/generic/Tensor.cpp"
#include "TH/THGenerate[Type]Type.h"
在lib/TH下则可以看到对应的THGenerate[Type]Type.h，如TH/THGenerateIntType.h：
#ifndef TH_GENERIC_FILE
#error "You must define TH_GENERIC_FILE before including THGenerateIntType.h"
#endif
#define real int32_t
#define ureal uint32_t
#define accreal int64_t
#define TH_CONVERT_REAL_TO_ACCREAL(_val) (accreal)(_val)
#define TH_CONVERT_ACCREAL_TO_REAL(_val) (real)(_val)
#define Real Int
#define THInf INT_MAX
#define TH_REAL_IS_INT
#line 1 TH_GENERIC_FILE
#include TH_GENERIC_FILE
#undef real
#undef ureal
#undef accreal
#undef Real
#undef THInf
#undef TH_REAL_IS_INT
#undef TH_CONVERT_REAL_TO_ACCREAL
#undef TH_CONVERT_ACCREAL_TO_REAL
#ifndef THGenerateManyTypes
#undef TH_GENERIC_FILE
#endif
头文件中有两个重要的宏定义：
#define real int32_t
#define Real Int
而在对应的Tensor.h中则有下面的宏定义:
#ifndef THP_TENSOR_INC
#define THP_TENSOR_INC
#define THPTensor TH_CONCAT_3(THP,Real,Tensor)
#define THPTensorStr TH_CONCAT_STRING_3(torch.,Real,Tensor)
#define THPTensorClass TH_CONCAT_3(THP,Real,TensorClass)
#define THPTensor_(NAME) TH_CONCAT_4(THP,Real,Tensor_,NAME)
......
其中THCONCAT*同样是宏定义，实现拼接字符串的功能，如对于IntTensor类型会有：
THPTensor_(init) will be convert to THPIntTensor_init
而类似THPTensor_(NAME)格式的字串在generic目录下的代码实现中大量出现。
这里的命名值得一提，源码中会经常遇到THP和TH前缀的变量，前者是pytorch中的变量前缀，区别于后者，源自Torch库中的对应变量。pytorch底层实现中调用了大量的Torch库。
THPTensor实现
在弄清楚上面两部分之后，再看generic目录下的代码就清晰很多了，还是以THPIntTensor为例来看。这里的THPIntTensor实际上是pytorch拓展的一个新Python类型。如果接触过Python源码的话会很清楚，定义一个新类型需要：
定义该对象包括哪些东西
为该对象定义类型
下面通过对比Python和pytorch的对象机制来简单说明：
Python对象机制
以C实现的Python为例，对于int类型，需要为其定义该类型：
typedef struct tagPyIntObject
{
PyObject_HEAD;
int value;
} PyIntObject;
对应类型有：
PyTypeObject PyInt_Type =
{
PyObject_HEAD_INIT(&PyType_Type),
"int",
...
};
其中PyObjectHEAD为宏定义，定义了所有对象所共有的部分，包括对象的引用计数和对象类型等共有信息，这也是Python中多态的来源。PyObjectHEAD_INIT是类型初始化的宏定义，简单来看如下：
#define PyObject_HEAD \
int refCount;\
struct tagPyTypeObject *type
#define PyObject_HEAD_INIT(typePtr)\
0, typePtr
如果对Python源码很感兴趣，强烈推荐陈儒（Robert Chen）的《Python源码剖析》，很是精炼。按照书中的介绍，我这里也实现了一个傻瓜Python以供参考：https://github.com/zqhZY/smallpy
pytorch对象机制
pytorch拓展的Tensor类型与Python的一般类型的定义类似，generic目录下的Tensor.h中有如下定义：
struct THPTensor {
PyObject_HEAD
// Invariant: After __new__ (not __init__), this field is always non-NULL.
THTensor *cdata;
};
同样的简洁明了，一个PyObject_HEAD头，一个THTensor类型指针指向具体内容。这里的THTensor类型就涉及到了TH库，也是源码的最底层的C语言实现。TH库的内容到最后再看，这里简单理解为一个类型。同样值得一提，根据上面的内容，这里的THPTensor和THTensor最终会转换成不同的类型：如Int对应THPIntTensor和THIntTensor。 同样，对于THPTensor对象也对应一个类型对象的定义，generic/Tensor.cpp中有：
PyTypeObject THPTensorType = {
PyVarObject_HEAD_INIT(NULL, 0)
"torch._C." THPTensorBaseStr, /* tp_name */
sizeof(THPTensor), /* tp_basicsize */
0, /* tp_itemsize */
(destructor)THPTensor_(dealloc), /* tp_dealloc */
0, /* tp_print */
...
...
0, /* tp_alloc */
THPTensor_(pynew), /* tp_new */
};
结构体中包括了很多指针，这里看最后的THPTensor(pynew)，该方法在该类型对象创建时调用，对应Python层面的new_函数。找到该函数：
static PyObject * THPTensor_(pynew)(PyTypeObject *type, PyObject *args, PyObject *kwargs)
{
...
// torch.Tensor()
if (num_args == 0) {
self->cdata = THPTensor_(_new)();
return (PyObject*)self.release();
}
...
}
THPTensor(pynew)先是为THPTensor申请内存，之后根据参数不同创建并初始化。 这里简单看代码中不传参数情况下创建Tensor调用THPTensor(_new):
static THTensor* THPTensor_(_new)()
{
THTensorPtr tensor(THTensor_(new)(LIBRARY_STATE_NOARGS));
if (!tensor->storage) {
tensor->storage = THStorage_(new)(LIBRARY_STATE_NOARGS);
}
return tensor.release();
}
该函数返回THTensor类型，方法中调用的THTensor(new)和THStorage(new)均为TH库中的内容，实现底层封装和内存申请。其中Storage保存了Tensor的值，之后具体看。
Python C拓展Tensor类型
看完 Tensor类型的初始化后，接下来是看如何把各种Tensor类型加入到”C“模块下供上层Python调用。这里回到一开始的torch/csrc/Module.cpp中PyMODINITFUNC方法的一系列初始化：
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));
该部分初始化对应到generic/Tensor.cpp中的THPTensor_(init):
bool THPTensor_(init)(PyObject *module)
{
...
THPTensorType.tp_methods = THPTensor_(methods);
...
PyModule_AddObject(module, THPTensorBaseStr, (PyObject *)&THPTensorType);
THPTensor_(initCopyMethods)();
return true;
}
该段代码有两点需要解释，一是Tensor模块的添加，二是Tensor对象的方法集的指定。
Tensor模块添加
PyModule_AddObject函数为Python C拓展API：
int PyModule_AddObject(PyObject *module, const char *name, PyObject *value)
Add an object to module as name.
而THPTensorBaseStr则是另外一个宏定义：
#define THPTensorBaseStr TH_CONCAT_STRING_2(Real,TensorBase)
同样是字符串拼接宏,对不同类型THPTensorBaseStr最终转换成[Type]TensorBase，再回到文章最开始的：
class IntTensor(_C.IntTensorBase, _TensorBase)
由此得到了Python层可以继承的_C.IntTensorBase。
Tensor类型添加方法
参考Pyhton实现，在定义一个对象后，对应的类型结构体中包含一个指针，指向该类型可以调用的方法集，例如list类型的用法：
a = []
a.append(1)
在pytorch的Tensor类型中该指针即为tpmethods，该指针的赋值，如上面THPTensor(init)中:
THPTensorType.tp_methods = THPTensor_(methods);
奇怪的是THPTensor_(methods)并不能在源码中找到，而在generic/methods可以看到一些.cwrap文件，如对ones函数有：
[[
name: ones
variants:
- function
auto_gpu: False
return: argument 0
arguments:
- arg: THTensor* result
output: True
- arg: THSize* size
long_args: True
]]
pytorch中自己实现了插件，根据这种YAML格式文本对TH库代码进行封装生成对应代码，插件代码在tools/cwrap/plugins。如addmv_函数对应生成的内容： https://gist.github.com/killeent/c00de46c2a896335a52552604cc4d74b 。
小结
到此，主要说明了pytorch中Tensor类型的定义及其模块拓展机制，可以使上层的Python调用C拓展的类型和相应方法。可以看到，pytorch中使用了代码生成方式，只定义一个模板，不同类型的Tensor对象通过该模板生成，避免了大量重复代码，虽然一开始一头雾水，但确实比较巧妙。
篇幅原因，这里并没有深入去看TH库部分的代码，pytorch对torch库做了CPython类封装，重用了大量代码，TH中主要的一个部分是THTensor的实现，后面再继续整理TH部分的代码。
参考文献
Python拓展CPython源码剖析精简版A Tour of PyTorch Internals (Part I)A quick tour of Torch internals
原创文章， 转载注明出处
以上是关于pytorch源码：C拓展的主要内容，如果未能解决你的问题，请参考以下文章
Pytorch拓展进阶：Pytorch结合C++以及Cuda拓展
源码学习成功运行RDAA代码及踩坑记录 - Pytorch版