最近在準備學習PyTorch原始碼，在看到網上的一些博文和分析後，發現他們發的PyTorch的Tensor原始碼剖析基本上是0.4.0版本以前的。比如說：在0.4.0版本中，你是無法找到a = torch.FloatTensor()中FloatTensor的usage的，只能找到a = torch.FloatStorage()。這是因為在PyTorch中，將基本的底層THTensor.h THStorage.h都放在名為Aten的後端中了（TH是torch7下面的一個重要的庫），並將之前放在torch/csrc/generic中的Tensor.h刪除。即相比之前做了模組解耦的工作。

0.前言（楔子）

我們知道，PyTorch中的Tensor的底層資料結構是Storage。那麼Storage是什麼？其實很簡單，Storage是一個連續（對應記憶體中的一段連續地址）的一維陣列，且裡面的元素型別是一樣的（比如都為Int

前面提到，基於PyTorch 0.4.0版本及目前最新的開原始碼中，我發現：使用者是無法找到a = torch.FloatTensor()中FloatTensor的usage的，只能找到a = torch.FloatStorage()。PyTorch開發者為了避免冗雜程式碼，所以在torch/csrc/generic中，將Tensor.h和Tensor.cpp都刪掉了。只保留了Storage.h和Storage.cpp，注意csrc目錄的作用：
將ATen中的基於torch 7的原生THTensor轉換為Torch Python的THPTensor

什麼是THTensor，什麼是THPTensor，包括後面還會見到的如THDPTensor、THCSPTensor等，都會在後面介紹。

下面，我將從原始碼中找到Storage，並逐步分析，究竟它是如何被封裝成我們日常使用的torch.FloatTensor等型別的。

class DoubleStorage(_C.DoubleStorageBase, _StorageBase):
    pass


class FloatStorage(_C.FloatStorageBase, _StorageBase):
    pass

...

class IntStorage(_C.IntStorageBase,
 
 _StorageBase):
    pass

不過，為了更好的學習程式碼，我們需要一些預備知識：

• 1）Python如何拓展C/C++庫
• 2）Python的實現機制

這些內容將放在本筆記最後，我將使用常見的API，用C語言寫module，然後被Python呼叫的例子進行展示。

1. 在Python擴充套件C

由class IntStorage(_C.IntStorageBase, _StorageBase): 可以看出，IntStorage
關於這塊的詳細介紹將在最後介紹，Pytorch中的拓展模組定義程式碼主要在torch/csrc/Module.cpp中，直接在Module.cpp找到我們關注的地方來進行說明：

 #include "torch/csrc/python_headers.h"
 #include <ATen/ATen.h>
 #include "THP.h"

 #ifdef USE_CUDNN
 #include "cudnn.h"
 #endif

 #ifdef USE_C10D
 #include "torch/csrc/distributed/c10d/c10d.h"
 #endif
...

#define ASSERT_TRUE(cmd) if (!(cmd)) return NULL
...

static PyObject* initModule() {
...
#if PY_MAJOR_VERSION == 2
  ASSERT_TRUE(module = Py_InitModule("torch._C", methods.data()));
// python3不支援Py_InitModule. 
// 現在, 使用者可以建立一個PyModuleDef structure，並將其引用傳遞給 PyModule_Create.
#else
  static struct PyModuleDef torchmodule = {
     PyModuleDef_HEAD_INIT,
     "torch._C",
     NULL,
     -1,
     methods.data()
  };
...
}
...
// 各種Torch Python型別的Storage初始化
ASSERT_TRUE(THPDoubleStorage_init(module));
ASSERT_TRUE(THPFloatStorage_init(module));
ASSERT_TRUE(THPHalfStorage_init(module));
ASSERT_TRUE(THPLongStorage_init(module));
...
#if PY_MAJOR_VERSION == 2
PyMODINIT_FUNC init_C()
#else
PyMODINIT_FUNC PyInit__C()
#endif
{
#if PY_MAJOR_VERSION == 2
  initModule();
#else
  return initModule();
#endif
}
// 到達結尾

那幾個標頭檔案很重要，#include <ATen/ATen.h>是因為PyTorch的很多模組，即這裡要分析的Storage就是基於ATen中的TH，TH表示Torch，因為PyTorch是從Torch 7移植過來的。相應地，THP表示Torch Python。TH和 THP的轉換定義在torch/csrc下的標頭檔案#include "THP.h"。

後面的USE_CUDNN和USE_C10D分別對應是否使用CUDNN和分散式。這裡的分析以最基礎的CPU上的Storage為例進行說明，不關注CUDNN和分散式。

在編譯過程中使用者可以建立一個PyModuleDef structure，並將其引用傳遞給 PyModule_Create，完成了torch._C的定義，接下來就是各種Torch Python型別的Storage初始化。

下面就是寫setup.py了，在setup.py中，主要就是寫Extension和setup：

• torch._C的Extension編寫
  
• setup編寫
  

寫好了setup.py就可以直接用python setup.py install安裝，安裝成功的話提示類似如下：

這樣就可以直接在.py檔案引用torch這個包了。

2. THPDoubleStorage_init(module)的來由

現在讓我們迴歸重點，那就是THPDoubleStorage_init(module)是從哪裡來的？直接在原始碼中查詢是找不到的。通過剛才的鋪墊，應該瞭解到THP是由TH轉換而成的。

2.1 Python C 物件對映

以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

同樣地，Pytorch拓展的Tensor型別與Python的一般型別的定義類似，torch/csrc/generic目錄下的Storage.h中有類似定義：

struct THPStorage {
  PyObject_HEAD
  THWStorage *cdata;
};

現在的重點就變成了THWStorage *cdata，還記得在Module.cpp中的#include 'THP.h'嗎？THP.h的第27行開始，將THWStorage定義為THStorage。現在是不是感覺有點懂了？對的，我們通過Storage.h和THP.h將THPStorage結構體裡面的資料型別變成了原來Torch 7框架中的基本資料型別THStorage了！

所以，雖然我們看起來是在用THPStorage，但是實際上，Pytorch對映為由ATen中TH庫的THStorage和THTensor。

#define THWStorage THStorage
#define THWStorage_(NAME) THStorage_(NAME)
#define THWTensor THTensor
#define THWTensor_(NAME) THTensor_(NAME)

2.2 ATen的TH庫

好了，由上面的分析，我們將一個THPStorage的底層定位到了ATen/src/TH中。下面，我們從THStorage.h，一步一步開始分析：

• ① THStorage.h
  由程式碼可以看出，其實THStorage.h儲存的目的就是為了相容性，重點在於THStorageFunctions.h。

#pragma once
#include "THStorageFunctions.h"

// Compatability header. Use THStorageFunctions.h instead if you need this.

• ② THStorageFunctions.h
  這個標頭檔案我們重點關注下面幾行

#define THStorage_(NAME) TH_CONCAT_4(TH,Real,Storage_,NAME)

#include "generic/THStorage.h"
#include "THGenerateAllTypes.h"

#include "generic/THStorage.h"
#include "THGenerateHalfType.h"

#include "generic/THStorageCopy.h"
#include "THGenerateAllTypes.h"

#include "generic/THStorageCopy.h"
#include "THGenerateHalfType.h"

其中#define THStorage_(NAME) TH_CONCAT_4(TH,Real,Storage_,NAME)是定義了一個字串拼接巨集。

它的作用很直觀，比如NAME = init, Real = Float的時候，那麼我們通過這個巨集，就會得到：

THStorage_init -------> THFloatStorage_init
而THFloatStorage_init就是在Module.cpp初始化中的內容：

現在，我們好奇的是在巨集命令中的Real是在哪裡定義的？容易發現，Real是由aten/src/TH/目錄下包含的一系列THGenerateDoubleType.h、THGenerateFloatType.h等THGenerate[Tensor型別]Type.h中。

• ③ THGenerateDoubleType.h
  以Double為例，看一下它的標頭檔案內容。

這裡需要注意的重點是第5行和第9行，那麼我們就知道Real是如何定義的了。

#define real double
#define Real Double

Real定義找到使用場景了，那麼real呢？

• ④ THStorageClass.hpp

現在，從THStorageClass.h定位到THStorageClass.hpp，其從40行開始定義了THStorage的結構體。這裡重點關注這些成員裡重點關注at::ScalarType scalar_type、at::DataPtr data_ptr、 ptrdiff_t size就可以了。

scalar_type 是變數型別：int，float等等；
data_ptr 是一維陣列的地址
比如 int a[3] = {1,2,3}，data_ptr是陣列a的地址，對應的size是3，不是sizeof(a)，scalar_type是int。

...
struct TH_CPP_API THStorage
{
  THStorage() = delete;
  THStorage(at::ScalarType, ptrdiff_t, at::DataPtr, at::Allocator*, char);
  THStorage(at::ScalarType, ptrdiff_t, at::Allocator*, char);
  // 關注下面3個成員變數
  at::ScalarType scalar_type;
  at::DataPtr data_ptr;
  ptrdiff_t size;
  // -----
  std::atomic<int> refcount;
  std::atomic<int> weakcount;
  char flag;
  at::Allocator* allocator;
  std::unique_ptr<THFinalizer> finalizer;
  struct THStorage* view;
  THStorage(THStorage&) = delete;
  THStorage(const THStorage&) = delete;
  THStorage(THStorage&&) = delete;
  THStorage(const THStorage&&) = delete;

  template <typename T>
  inline T* data() const {
    auto scalar_type_T = at::CTypeToScalarType<th::from_type<T>>::to();
    if (scalar_type != scalar_type_T) {
      AT_ERROR(
          "Attempt to access Storage having data type ",
          at::toString(scalar_type),
          " as data type ",
          at::toString(scalar_type_T));
    }
    return unsafe_data<T>();
  }

  template <typename T>
  inline T* unsafe_data() const {
    return static_cast<T*>(this->data_ptr.get());
  }
};

現在我們知道了THStorage的結構體，那麼接下來，就去THStorageClass.cpp檢視其建構函式：

#include "THStorageClass.hpp"

THStorage::THStorage(
    at::ScalarType scalar_type,
    ptrdiff_t size,
    at::DataPtr data_ptr,
    at::Allocator* allocator,
    char flag)
    : scalar_type(scalar_type),
      data_ptr(std::move(data_ptr)),
      size(size),
      refcount(1),
      weakcount(1), // from the strong reference
      flag(flag),
      allocator(allocator),
      finalizer(nullptr) {}

THStorage::THStorage(
    at::ScalarType scalar_type,
    ptrdiff_t size,
    at::Allocator* allocator,
    char flag)
    : THStorage(
		  // 標量型別
          scalar_type,
          size,
          allocator->allocate(at::elementSize(scalar_type) * size),
          allocator,
flag) {}

現在，可能細心的讀者會發現，之前預定義的real還沒用到啊？這東西到底在哪裡用呢？

• ⑤ generic/THStorage.cpp

答案就是TH庫的generic/THStorage.cpp 裡用！下面的程式碼就是使用的例子。通過將 THStorageClass.hpp 、THStorageClass.cpp THStorage.cpp聯合分析，終於找到了在THGenerate[Tensor型別]Type.h定義real的使用地點。

THStorage* THStorage_(newWithSize)(ptrdiff_t size)
{
  THStorage* storage = new THStorage(
      at::CTypeToScalarType<th::from_type<real>>::to(),
      size,
      getTHDefaultAllocator(),
      TH_STORAGE_REFCOUNTED | TH_STORAGE_RESIZABLE);
  return storage;
}

2.3 轉向Tensor

通過2.1和2.2的分析，我們能夠明白一個Storage的組成方式：

THPStorage(Torch Python層的結構體定義，位於csrc/generic/Storage.h)
——>
THWStorage——（THWSorage型別的具體內容，位於csrc/generic/Storage.h)
——>
THStorage(巨集定義轉換，位於csrc/THP.h)
——>
THStorage的結構體 (位於ATen/src/TH/THStorageClass.hpp)
——>
THStorage的兩種構造方法(位於ATen/src/TH/THStorageClass.cpp)

跟Storage類似，Tensor的結構體定義在aten/src/TH/THTensor.hpp中，可以看出，它完全是基於Storage來構建的，對應的是THStorageClass.cpp的第一種建構函式。

...
struct THTensor
{
    THTensor(THStorage* storage)
      : refcount_(1)
      , storage_(storage)
      , storage_offset_(0)
      , sizes_{0}
      , strides_{1}
      , is_zero_dim_(false)
      {}

    ~THTensor() {
      if (storage_) {
        THStorage_free(storage_);
      }
	}
...
}
...

3. THPStorage的實現

目前，前面的內容已經梳理明白了。那麼就讓我們把目光轉回到對映關係：C/C++物件————>Python型別

接觸過Python原始碼的人會比較清楚，定義一個新型別需要：

• ① 定義該物件包括哪些內容

• ② 為物件定義型別

3.1 定義物件包含內容

現在，我們找到pytorch/torch/csrc/generic目錄下的Storage.cpp。
這裡面就定義了型別中包含的內容：

PyTypeObject THPStorageType = {
  PyVarObject_HEAD_INIT(NULL, 0)
  "torch._C." THPStorageBaseStr,         /* tp_name */
  sizeof(THPStorage),                    /* tp_basicsize */
  0,                                     /* tp_itemsize */
  (destructor)THPStorage_(dealloc),      /* tp_dealloc */
  0,                                     /* tp_print */
  0,                                     /* tp_getattr */
  0,                                     /* tp_setattr */
  0,                                     /* tp_reserved */
  0,                                     /* tp_repr */
  0,                                     /* tp_as_number */
  0,                                     /* tp_as_sequence */
  &THPStorage_(mappingmethods),          /* 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 | Py_TPFLAGS_BASETYPE, /* tp_flags */
  NULL,                                  /* tp_doc */
  0,                                     /* tp_traverse */
  0,                                     /* tp_clear */
  0,                                     /* tp_richcompare */
  0,                                     /* tp_weaklistoffset */
  0,                                     /* tp_iter */
  0,                                     /* tp_iternext */
  0,   /* will be assigned in init */    /* tp_methods */
  0,   /* will be assigned in init */    /* tp_members */
  0,                                     /* tp_getset */
  0,                                     /* tp_base */
  0,                                     /* tp_dict */
  0,                                     /* tp_descr_get */
  0,                                     /* tp_descr_set */
  0,                                     /* tp_dictoffset */
  0,                                     /* tp_init */
  0,                                     /* tp_alloc */
  THPStorage_(pynew),                    /* tp_new */
};

顯然，結構體中包括了很多指標，如最後的THPStorage_(pynew)，該方法在該型別物件建立時呼叫，對應Python類層面中的__new__函式。

THPStorage_(pynew)定義在當前的Storage.cpp中，主要工作就是申請記憶體和分配（並檢查引數，將資料轉移到gpu視訊記憶體等等————36行開始），有興趣的同學自行看吧…

3.2 為物件定義型別

現在，要看的是把各種Storage型別加入到”_C“模組下供上層Python呼叫。

回到torch/csrc/Module.cpp中的一系列初始化：

// 各種Torch Python型別的Storage初始化
ASSERT_TRUE(THPDoubleStorage_init(module));
ASSERT_TRUE(THPFloatStorage_init(module));
ASSERT_TRUE(THPHalfStorage_init(module));
ASSERT_TRUE(THPLongStorage_init(module));

該部分初始化對應到torch/csrc/generic/Storage.cpp中的THPStorage_(init)(PyObject *module):

該段程式碼中需要解釋的主要就是：
1）Storage模組的新增

上面的第329行，PyModule_AddObject的作用就是像module裡面新增模組，其定義如下：

//將名為name的PyObject指標value加入到模組module中去
int PyModule_AddObject(PyObject *module, const char *name, PyObject *value){
                ...
}

用法如下：一般是判斷是否將模組匯入成功

而其中的第2個引數THPStorageBaseStr則是一個在Storage.h中定義的**拼接巨集**引數：

作為一個字串拼接巨集,對不同型別，THPStorageBaseStr最終轉換成[Type]StorageBase：
以Real為Int為例：
經過此THPStorageBaseStr這個字串拼接巨集，我們得到了IntStorageBase。

即通過 ① THPStorageBaseStr字串拼接巨集 ② 函式PyModule_AddObject就將IntStorageBase、FloatStorageBase等內容新增到_C下面。

由此，我們得到了Python層可以繼承的_C.FloatStorageBase,_C.DoubleStorageBase等等。

2）Storage物件的方法集的指定

在Python中，在定義一個物件後，對應的型別結構體中，會包含一個指標，指向該型別可以呼叫的方法集，例如Python內建型別set的用法：

a = set()
a.add(10)

在PyTorch的Storage型別中,這個可以指向可以呼叫的方法集的指標即為tp_methods，該指標的賦值如下，等於methods.data()。

其中methods是由上面（319,321）的THPUtils_addPyMethodDefs(methods, THPStorage_(xxx))來將xxx匯入到methods中的。

319行-321行含義：新增自定義的方法集，如果THD_GENERIC_FILE的巨集沒有定義，那麼就將通用方法集新增到Tensor中去。

這些方法包括max()、min()等等，詳細內容請檢視官方文件。

4. 預備知識

4.1 如何寫Python/C 擴充套件

提到寫擴充套件，首先要問問為什麼我們需要寫擴充套件呢？ 答案很如下：

1） You want speed and you know C is about 50x faster than Python.

2） Certain legacy C libraries work just as well as you want them to, so you don’t want to rewrite them in python.

3） Certain low level resource access - from memory to file interfaces.

4） Just because you want to.

主要有3種方法：1）Ctypes 2）SWIG 3）Python/C API（最廣泛使用）

我們以第3種為例進行說明

4.1.1 簡介

所有的Python物件（objects）都以PyObject結構體的形式存在，Python.h的標頭檔案中包含很多函式來操作它。

舉個例子，一個PyObject物件是一個PyListType（即Python中的list）,我們就可以對結構體使用PyList_Size()函式來獲得這個列表的長度（相當於len(list)）。

假設我們要寫一個很簡單的函式，官網的例子是對list求和（list裡面都是int）。

程式碼看起來長這樣，看起來很正常。但是唯一不同之處在於：Package addList是用C寫的

#Though it looks like an ordinary python import, the addList module is implemented in C
import addList

l = [1,2,3,4,5]
print "Sum of List - " + str(l) + " = " +  str(addList.add(l))

4.1.2 寫adder.c

1. include <Python.h>隱含了一些標準的標頭檔案： stdio.h, string.h, errno.h, limits.h, assert.h and stdlib.h (if available)

2.addList_add(...)接收PyObject型別的結構體。傳過來的引數 通過 PyArg_ParseTuple()將tuple拆分成一個個單獨的element。
其中，
第一個引數是要解析的引數變數，

第二個引數是解析方法，也就是下面的"O", "siO"等，剩下的引數就是指解析出的內容的對應物件地址。

  int n;
  char *s;
  PyObject* list;
  PyArg_ParseTuple(args, "siO", &s, &n, &list);

另外，我們不需要PyArg_ParseTuple()的返回值。下面是adder.c的程式碼
（需要注意，這裡面最後跟一些教程不一樣，是我自己改的，因為那些教程是基於Python2的寫法，對於Python3是不能用的）：

```C
//Python.h這個標頭檔案擁有所有我們需要的資料型別（用以表徵Python物件型別）和函式定義（用以操作Python物件）
#include <Python.h>

 //這就是在Python程式碼裡面需要呼叫的函式————通常的命名規則是
 //{module-name}_{function-name}
static PyObject* addList_add(PyObject* self, PyObject* args){

  PyObject * listObj;
  
  //解析輸入引數args(型別為PyObject指標) 引數傳過來的預設形式是tuple（元組），我們將它解析
  // 這裡只有一個list，下面會介紹當有多個輸入時，應該如何解析。
  // 在，PyArg_ParseTuple裡面，第2個引數中：‘i’ 表示 integer, ‘s’ 表示 string ‘O’ 表示一個 Python object
  // 如果解析多個引數：
  // int n;
  // char *s;
  // PyObject* list;
  // PyArg_ParseTuple(args, "siO", &s, &n, &list);
  
  if (! PyArg_ParseTuple( args, "O", &listObj))
    return NULL;
  
  // 現在已經將引數args 解析到 listObj物件中了
  long length = PyList_Size(listObj);
   
  // 求和
  long i = 0;
  // 
  long sum = 0; // short sum = 0;
  for(i = 0; i < length; i++){
    // 從ListObj中逐個取元素，每個元素同樣地，也是一個python物件
    
    PyObject* temp = PyList_GetItem(listObj, i);
    
    // 因為這個temp實際上也是一個python物件，所以將它轉換為C中原生型別中的Long  （我試試Short）
    long elem = PyInt_AsLong(temp);
    // short elem = PyInt_AsShort(temp); 
    
    sum += elem;
  }

  //value returned back to python code - another python object
  //build value here converts the C long to a python integer
  
  // 將值返回給Python程式碼，即還需要將C long/short 轉換成Python Integer
  return Py_BuildValue("i", sum);
}

// 文件說明：
static char addList_docs[] =
    "add( ): add all elements of the list\n";
     
/* This table contains the relavent info mapping -
  <Python模組中的函式名稱>, <對應C/C++中的函式體>,
  <函式期望的引數格式>, <函式的文件說明>
*/

static PyMethodDef addList_funcs[] = {
    {"add", (PyCFunction)addList_add, METH_VARARGS, addList_docs},
    {NULL, NULL, 0, NULL}
};

/*
注意：Python3不