PyTorch源碼淺析（一）
CPU上的張量（多維數組）庫
TH庫的實現使用了用C語言的宏產生的泛型，並且通過命名規則來產生類似面向對象的效果。這部分我們在這一章後面介紹。
TH負責實現CPU上的張量（Tensor）運算，存儲，內存分配。張量的內存管理和運算類型通過THTensor和THStorage兩個C泛型來進行建模。張量這個數學對象被TH分解為THTensor和THStorage，THTensor提供一種查看THStorage的方法，THStorage負責管理張量的存儲方式。
數據存儲
存儲的數據結構聲明如下
typedef struct THStorage{ real *data; ptrdiff_t size; int refcount; char flag; THAllocator *allocator; void *allocatorContext; struct THStorage *view;} THStorage;
所有在CPU上的張量實際上都是內存中的一個一維C數組（C指針）data來存儲，並且使用引用計數（reference count）來管理內存。
構造函數
所有構造新THStorage的函數都以new開頭，後面跟具有相關含義的後綴名。
TH_API THStorage* THStorage_(new)(void);TH_API THStorage* THStorage_(newWithSize)(ptrdiff_t size);TH_API THStorage* THStorage_(newWithSize1)(num);TH_API THStorage* THStorage_(newWithSize2)(num, num);TH_API THStorage* THStorage_(newWithSize3)(num, num, num);TH_API THStorage* THStorage_(newWithSize4)(num, num, num, num);TH_API THStorage* THStorage_(newWithMapping)(const char *filename, ptrdiff_t size, int flags);/* takes ownership of data */TH_API THStorage* THStorage_(newWithData)(num *data, ptrdiff_t size);TH_API THStorage* THStorage_(newWithAllocator)(ptrdiff_t size, THAllocator* allocator, void *allocatorContext);TH_API THStorage* THStorage_(newWithDataAndAllocator)(num* data, ptrdiff_t size, THAllocator* allocator, void *allocatorContext);
析構函數都以`free`開頭（實際上只有一個名為`free`的函數）
張量
張量在TH中是一種查看存儲（Storage）的方法。它包括以下內容：
- long *size：一個用來存儲張量各個維度大小的一維數組
- long *stride：一個用來存儲張量各個角標偏移量的數組
- int nDimension：維度
- THStorage *storage：存儲的指針 （作者在這裡也註明瞭，張量大小實際上是小於等於存儲大小的）
- ptrdiff_t: 存儲偏移
- refcount：引用計數
- char flag：（暫時還沒完全看懂`flag`有啥用）
它的具體聲明如下
typedef struct THTensor{ long *size; long *stride; int nDimension; // 注意: storage->size 可能比張量的大小要大 THStorage *storage; ptrdiff_t storageOffset; int refcount; char flag;} THTensor;
我們接下來具體解釋幾個可能不太容易理解的地方。首先是`stride`，說道`stride`我們要先簡單介紹諸如**NumPy**，**Eigen**等提供了BLAS（基本線性代數運算）是如何存儲一個矩陣的。首先矩陣在內存中實際上都作為一個內存塊進行存儲，在C語言看來它是一個一維數組或者說是由`malloc`或者`calloc`分配的某個給定大小的內存塊，例如下表是一個有20個浮點類型（雙精度）的內存塊，它可能存儲了一個4x5矩陣的值，也有可能存儲了一個2x5x2的三階張量的值。
向系統申請這個內存塊，在不再使用之後刪除所分配的內存，將內存塊固定到硬碟等存儲中，以及訪問指定地址等任務實際上就可以單獨交給THStorage來完成，因為我們並不需要知道其對應張量的大小。甚至有可能幾個元素數目不同但是總數相同的張量（比如4x4,2x2x2x2,1x16的不同大小張量）可以通過用不同的THTensor共享一塊內存（共用一個THStorage，但此THStorage的引用計數將會大於等於3）。但當我們需要完成張量的一些運算，例如對於矩陣，他們的乘積（matrix product），點積（dot product）等運算會需要使用維度的信息（各個維度的大小）並且這個時候我們將按照維度來訪問不同位置的元素，這使得我們首先需要存儲各個維度的大小long *size，但是這還不夠，我們實際上在訪問一塊連續內存的時候實際上使用的是不同維度上的間隔，例如第一個維度上的間隔一般是0，第二個維度上的間隔是第一個維度的大小size[0]，依次類推，但也有可能由於是由某個較大的張量分割來的，並不滿足上述間隔分配方式，所以我們有必要再用一個數組存儲各個維度的間隔大小long *stride，同時再加上內存的偏移量storageOffset。這樣在訪問某個角標ijk對應的內存地址時就可以用
storageOffset + i * stride[0] + j * stride[1] + k * stride[2]
來獲得其真實內存地址了。而類似於存儲，一個張量也有可能被不用的變數所使用，這也需要一個引用計數refcount來管理內存。
張量構造
張量的構造相比存儲對象的構造就麻煩多了，但很多時候這些操作的共性就是對每一個或者部分張量元素使用某一個函數，在一些語言或者框架中，這被稱為map函數。在TH中，使用了宏函數來做到一個高性能的map函數，我們首先介紹一下TH是如何使用宏函數做到高性能的map的。
TensorApply宏
`TensorApply`系列的宏函數是TH實現各種張量元素操作最重要的操作，它們負責把一個針對某些標量的操作應用到多個張量元素上去。在GPU部分是相當於一個map的操作。大致方法是優先去操作內存連續部分，然後再操作不連續的部分，以增加CPU cache命中率。詳細內容留到下一篇文章講。
使用C語言實現面向對象以及泛型
在PyTorch/Torch中，後端的庫都使用了宏來進行泛型等功能的實現。下面我們用一個例子介紹這一部分。面向對象這一點可以通過命名規範來完成，例如我們的向量結構體如果是`Vector`，那麼屬於這個結構體的方法就是`Vector_xxx`。下面主要介紹泛型。
需求
現在我們需要在C語言中實現對兩個向量的加法`add`。並且向量的類型可以是：`float`, `double`。
實現一
很容易想到的一個方法就是針對不同類型編寫按照規則命名的`add`函數以及向量結構體`Vector`，例如我們可以分別實現如下的`Vector`類型：`Float_Vector`, `Double_Vector`。同時實現其各自對應的加法函數（假設函數輸入類型必須一致）：`Float_Vector_add`, `Double_Vector_add`。
實現二
上述的實現方法實際上重複寫了很多代碼，我們知道兩個向量的加法就是各個元素對應相加。以上所有類型所需的演算法是完全相同的。假如在實現了泛型語言中去做這件事情是非常簡單的，比如在C++中我們可以使用模板函數
// 這裡的Vector是某個自己實現的類型template<typename T>void add(Vector<T> &c, Vector<T> &a, Vector<T> &b){ for(int i=0; i<a.size(); i++) {c.data[i] = a.data[i] + b.data[i] }}
或者對於一些有自動類型匹配的語言，比如Julia，直接將變數指定為這些類型的抽象類型即可
function add!{T<:Number}(c::Vector{T}, a::Vector{T}, b::Vector{T}) for i=1:size(c) c[i] = a[i] + b[i] endend
而C並不具備這樣的功能。但不難從實現一中發現，不同類型的命名方式是固定的，這使得我們可以通過藉助文本替換的方式來完成自動命名，也就間接實現了泛型。而文本替換可以藉助外部程序來完成例如一些模板語言（template language）,也可以自己來寫。好在我們現在的後端是用C語言而不是Fortran95，C自身提供了宏來實現類似的功能。而對於Fortran95，就只能使用像Jinja這樣的模板語言來完成泛型的支持了。
PyTorch選擇了兩種方案，在後端代碼中利用宏來完成泛型的支持，而在中間的膠水代碼部分，使用了一個用Python實現的，通過一種YAML標記語言的變體生成泛型膠水代碼的生成器。不過這一部分我們著重關注第一種實現方案。下面我們繼續。
回顧一下C語言的宏
關於C語言的宏，可以說是C語言中最有趣的一部分。下面關於C語言宏預處理器的介紹來自於GNU的宏命令在線文檔我們只是簡單的回顧，如果有疑問請詳細閱讀這份在線文檔。
指令（Directive）
#define MACRO_NAME VALUE
定義一個宏（Macro），其名稱為MACRO_NAME，值（將被展開的形式）VALUE
#line digit "finename"
改變編譯器存儲的當前行號`digit`和文件名`finename`為指定的行號和文件名。
#include "filepath"#include <filepath>
預讀取指定文件`filepath`，對於雙引號中的文件，將在本地目錄查找。對尖括弧中的內容將在環境目錄中查找。
宏變數
宏變數是最簡單的宏，例如
#define BUFFER_SIZE 1024
在預處理器工作的時候，當後面的代碼出現了`BUFFER_SIZE`，就會將其替換為`1024`，例如下面的代碼
BUFFER_SIZE + 2
就會被替換為
1024 + 2
宏的定義支持續行符``，當一個宏命令過長時，我們可以通過使用續行符來整理你的代碼。這個我們會在後面遇到。
所以也正式因為它只是簡單的**文本替換**使用它也是很危險的，如果定義不當，程序可能會出現作者沒有預料的行為。所以一定要小心。
有時候，我恰好和需要再次使用相同宏變數的名字，這個時候需要取消定義
#undef BUFFER_SIZE
這樣在此之後預處理器就不會將`BUFFER_SIZE`替換為宏後面的內容了
宏函數
宏也可以具有參數，其行為類似於函數，但實際上很不一樣。例如
#define MIN(X, Y) X < Y? X : Y
這個宏函數實現了比較其輸入變數大小的功能，例如執行
// 獲得最小的數字MIN(2, 3); //注意一定要有分號
將會得到`2`，這是因為預處理器將宏`MIN`替換成了
2 < 3? 2 : 3;
這個表達式將返回`2`。可見實際上宏函數也是某種文本替換，但是不當的聲明是很危險的，例如上面的這個宏，若我們
#define G 1 + 2MIN(G, 2);
預處理器將替換為
1 + 2 < 2? 2 : 2;
這是不符合我們原本的意圖的。所以我們要修改原來的定義來防止不必要的意外發生。
#define MIN(X, Y) ((X) < (Y)? (X): (Y))
還有就是一定不要在宏的最後使用分號，這是為了保證代碼樣式的統一。例如
#define UglyMIN(X, Y) ((X) < (Y) ? (X): (Y));
會使得在使用時沒有分號，看起來和正常的語句不同。
將宏名稱轉換為字元串
如果我們使用宏來產生泛型，那麼在拋出錯誤等場景可能會需要輸出是哪個類型錯了在宏內部可以使用`#`來產生字元串，例如
#define WARN(EXP) printf(#EXP)
會將輸入的變數變為字元串再替換
WARN(test);
被替換為
printf("test");
組合名字
當我們使用不同的宏產生名字時，我們最終需要將它們組合起來。
#define CONCAT(A, B, C) A ## B ## C
例如這個宏可以用來產生`Double_Matrix_add`這個變數名
Double_Matrix CONCAT(Double, Matrix, add)(Double_Matrix *A, Double_Matrix *B);
一些預定義的宏
C語言的預處理器有一些預定義的宏
- `__FILE__` 當前輸入文件名稱，是一個C字元串常量，這個變數會展開為完整路徑，而不是像是在`#include`中使用的簡寫。
- `__LINE__` 當前輸入行號，是一個整數常量，這個宏的值會隨著預處理器讀入的行的改變而改變，所以其行為與其它預定義宏略有不同。
構建你的C泛型
首先假設我們已經有了泛型`num`，接下來我們試著按照實現一中的命名規則寫出利用這個泛型構造的向量類型和`add`函數
struct NumVector{ num *data; int n;}// C = A + Bvoid NumVector_add(NumVector *C, NumVector *A, NumVector *B){// check size if(!((C->n == A->n) && (C->n == B->n))) { exit(1); // 稍後再說產生異常的問題，先這麼退出 } int i,j, n; n = C->n; for(i=0; i<n; i++) { C->data[i] = A->data[i] + B->data[i]; }}
現在考慮如何將類似於`Num_add`的形式特例化為`FloatVector_add`等類型名稱。這個可以用宏函數實現
#define Vector_(NAME) Num ## Vector_ ## NAME#define Vector Num ## Vector#define num float#define Num Floatstruct Vector{ num *data; int n;};void Vector_(add)(Vector *C, Vector *A, Vector *B){//codes}
我們期望這些宏將把以上函數和結構體替換為
struct FloatVector{ float *data; int n;};void FloatVector_add(FloatVector *C, FloatVector *A, FloatVector *B){//codes}
但是實際上以上代碼只能產生`NumVector`的名字，這是因為C的宏定義在出現`#`和`##`時不會展開宏名，我們需要使用一個中間宏來讓編譯器先展開宏名，然後再組合它們。修改後如下
#define CONCAT_2_EXPAND(A, B) A ## B#define CONCAT_2(A, B) CONCAT_2_EXPAND(A, B)#define CONCAT_3_EXPAND(A, B, C) A ## B ## C#define CONCAT_3(A, B, C) CONCAT_3_EXPAND(A, B, C)#define Vector_(NAME) CONCAT_3(Num, Vector_, NAME)#define Vector CONCAT_2(Num, Vector)#define num float#define Num Floatstruct Vector{ num *data; int n;};void Vector_(add)(Vector *C, Vector *A, Vector *B){//codes}
但是這隻能產生一種類型對應的函數，如果要產生多種類型的函數就需要有如下的結構
// add.c#define CONCAT_2_EXPAND(A, B) A ## B#define CONCAT_2(A, B) CONCAT_2_EXPAND(A, B)#define CONCAT_3_EXPAND(A, B, C) A ## B ## C#define CONCAT_3(A, B, C) CONCAT_3_EXPAND(A, B, C)#define Vector_(NAME) CONCAT_3(Num, Vector_, NAME)#define Vector CONCAT_2(Num, Vector)#define num float#define Num Floatstruct Vector{ num *data;int n;};void Vector_(add)(Vector *C, Vector *A, Vector *B){//codes}#undef num#undef Num#define num double#define Num Doublestruct Vector{ num *data; int n;};void Vector_(add)(Vector *C, Vector *A, Vector *B){//codes}#undef num#undef Num// etc.
這樣不斷複製粘貼之前的帶宏命令的代碼肯定是不現實的。但如果這部分泛型代碼在另外一個文件裏的話，那麼豈不是每次從這個文件開始讀取不就好了？我們現在將這部分代碼分離出去，放在`generic/`文件夾下(這樣就可以取相同的名字，方便記憶)，現在工程目錄如下
.├── add.c # 用來展開generic/add.c├── add.h # 用來展開generic/add.h├── general.h # 用來包含其它頭文件└── generic ├── add.c # 泛型add函數定義 └── add.h # 泛型Vector類型的定義
現在`add.h`和`add.c`裏變成了這樣
// add.h#include "general.h"#define Vector_(NAME) CONCAT_3(Num, Vector_, NAME)#define Vector CONCAT_2(Num, Vector)#define num float#define Num Float#include "generic/add.h"#undef num#undef Num#define num double#define Num Double#include "generic/add.h"#undef num#undef Num
// add.c#include "add.h"#define num float#define Num Float#include "generic/add.c"#undef num#undef Num#define num double#define Num Double#include "generic/add.c"#undef num#undef Num
用`nm`命令查看一下鏈接庫裏的函數名
>>> nm *.aadd.c.o:000000000000007e T DoubleVector_add0000000000000000 T FloatVector_add
成功了，現在寫一個測試文件來看看是否正確
#include "general.h"#include "add.h"int main(int argc, char const *argv[]){int i, n;FloatVector *A, *B, *C;A = (FloatVector *)malloc(sizeof(FloatVector));B = (FloatVector *)malloc(sizeof(FloatVector));C = (FloatVector *)malloc(sizeof(FloatVector));n = 10;A->data = (float *)calloc(n, sizeof(float));B->data = (float *)calloc(n, sizeof(float));C->data = (float *)calloc(n, sizeof(float));A->n = n;B->n = n;C->n = n;for(i=0;i<n;i++){ A->data[i] = i; B->data[i] = 2 * i; C->data[i] = 0;}FloatVector_add(C, A, B);for(i=0;i<n;i++){printf("%f
", C->data[i]);}free(A);free(B);free(C);return 0;}
0.0000003.0000006.0000009.00000012.00000015.00000018.00000021.00000024.00000027.000000
正確無誤！
唔，但是我們總不能每次都寫一遍`num`這泛型的宏定義，我們現在把它打包到一個頭文件`GenerateFloat.h`裏去，然後用一個宏`GENERIC_FILE`來存儲要進行特例化的文件名。首先判斷是否定義了這個宏
#ifndef GENERIC_FILE#error "You must define GENERIC_FILE before including GenerateFloat.h"#endif
然後把剛才的特例化宏代碼挪進來，加入`#line`使得編譯器每次載入`GENERIC_FILE`的時候`__LINE__`都是從1開始，就好像是重新讀入一樣。
// GenerateFloat.h#ifndef GENERIC_FILE#error "You must define GENERIC_FILE before including GenerateFloat.h"#endif#define num float#define Num Float#line 1 GENERIC_FILE#include GENERIC_FILE#undef num#undef Num#define num double#define Num Double#line 1 GENERIC_FILE#include GENERIC_FILE#undef num#undef Num
現在再修改`generic/add.h`和`generic/add.c`定義`GENERIC_FILE`這個宏
// generic/add.h#ifndef GENERIC_FILE#define GENERIC_FILE "generic/add.h"#elsetypedef struct Vector{num *data;int n;} Vector;extern void Vector_(add)(Vector *C, Vector *A, Vector *B);#endif
// generic/add.c#ifndef GENERIC_FILE#define GENERIC_FILE "generic/add.c"#elsevoid Vector_(add)(Vector *C, Vector *A, Vector *B){int i, n;n = C->n;for(i=0;i<n;i++){C->data[i] = A->data[i] + B->data[i];}}#endif
完成！
下一篇具體講TensorApply，CPU的THNN庫部分。然後講THC也就是CUDA部分。
羅秀哲：PyTorch源碼淺析（目錄）
推薦閱讀：
