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]
- end
- end
而 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 + 2
- MIN(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 + B
- void 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 Float
- struct 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 Float
- struct 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 Float
- struct Vector
- {
- num *data;
- int n;
- };
- void Vector_(add)(Vector *C, Vector *A, Vector *B)
- {
- //codes
- }
- #undef num
- #undef Num
- #define num double
- #define Num Double
- struct 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 *.a
- add.c.o:
- 000000000000007e T DoubleVector_add
- 0000000000000000 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\n", C->data[i]);
- }
- free(A);
- free(B);
- free(C);
- return 0;
- }
- 0.000000
- 3.000000
- 6.000000
- 9.000000
- 12.000000
- 15.000000
- 18.000000
- 21.000000
- 24.000000
- 27.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"
- #else
- typedef 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"
- #else
- void 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
来源: http://www.tuicool.com/articles/go/y226NjA