pytorch源码：C拓展

2017-10-18 本文已影响3802人 zqh_zy

head

读了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));
...
}

在编译过程中PyMODINIT_FUNC方法被调用，完成了"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部分的实现，如上面PyMODINIT_FUNC函数中的THPIntTensor_init(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")

其中split_types在pytorch/tools/setup_helpers/目录下实现，主要实现了两个功能：

重命名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)
......

其中TH_CONCAT_*同样是宏定义，实现拼接字符串的功能，如对于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",
     ...
};

其中PyObject_HEAD为宏定义，定义了所有对象所共有的部分，包括对象的引用计数和对象类型等共有信息，这也是Python中多态的来源。PyObject_HEAD_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中PyMODINIT_FUNC方法的一系列初始化：

  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类型中该指针即为tp_methods，该指针的赋值，如上面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拓展C
Python源码剖析精简版
 A Tour of PyTorch Internals (Part I)
A quick tour of Torch internals

原创文章，转载注明出处
更多关注公众号：

wechat