TensorApply宏
如同作者注释所言，tensor apply系列的宏的机制如下
- 从最外部的角标开始，循环至第一个发生内存不连续的地址，然后将其记为张量A，A所在的内存都是连续的，把剩下的记为B。
- 然后接下来有限对B从最外面的角标进行遍历，而对于Ａ由于内存本身就是连续的，我们直接这一整块内存进行遍历
- 然后为了减少循环嵌套，将A中在内存上连续（具体来说就是和stride乘积相等）的维度组合到一起。
注释原文如下
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 | /* * The basic strategy for apply is as follows: * * 1. Starting with the outermost index, loop until we reach a dimension where the * data is no longer contiguous, i.e. the stride at that dimension is not equal to * the size of the tensor defined by the outer dimensions. Let's call this outer * (contiguous) tensor A. Note that if the Tensor is contiguous, then A is equal * to the entire Tensor. Let's call the inner tensor B. * * 2. We loop through the indices in B, starting at its outermost dimension. For * example, if B is a 2x2 matrix, then we do: * * B[0][0] * B[0][1] * B[1][0] * B[1][1] * * We set the offset into the underlying storage as (storageOffset + stride_B * index_B), * i.e. basically we compute the offset into the storage as we would normally for a * Tensor. But because we are guaranteed the subsequent data is contiguous in memory, we * can simply loop for sizeof(A) iterations and perform the operation, without having to * follow the order described by the strides of A. * * 3. As an optimization, we merge dimensions of A that are contiguous in memory. For * example, if A is a 3x3x3x3 tensor narrowed from a 3x3x4x3 tensor, then the first two * dimensions can be merged for the purposes of APPLY, reducing the number of nested * loops. */ |
具体实现暂且不表，我们讲讲这写宏要如何使用，我从 THTensorMath.c 里选了一个cadd函数来举例子
1 2 3 4 5 6 7 8 9 10 11 12 13 | void THTensor_(cadd)(THTensor *r_, THTensor *t, real value, THTensor *src) { THTensor_(resizeAs)(r_, t); if (THTensor_(isContiguous)(r_) && THTensor_(isContiguous)(t) && THTensor_(isContiguous)(src) && THTensor_(nElement)(r_) == THTensor_(nElement)(src)) { if(r_ == t) { THBlas_(axpy)(THTensor_(nElement)(t), value, THTensor_(data)(src), 1, THTensor_(data)(r_), 1); } else { TH_TENSOR_APPLY3_CONTIG(real, r_, real, t, real, src, THVector_(cadd)(r__data, t_data, src_data, value, r__len);); } } else { TH_TENSOR_APPLY3(real, r_, real, t, real, src, *r__data = *t_data + value * *src_data;); } } |
这里cadd的作用是遍历张量t和src中的元素，将src中的元素乘以value之后加上t中的元素赋值给r_。
1 | *r__data = *t_data + value * *src_data; |
这个函数首先会确认t和r_的大小，如果r_没有声明是一个空指针，THTensor_(resizeAs)函数会按照t的大小分配一块新的内存给r_这个指针。if的第一段暂且不说，这是为了增加向量化操作而写的代码，我们先看通用的TH_TENSOR_APPLY3这个宏。这个宏的声明如下
1 | #define TH_TENSOR_APPLY3(TYPE1, TENSOR1, TYPE2, TENSOR2, TYPE3, TENSOR3, CODE) |
后面的数字３是说这个宏会对三个张量进行遍历。TYPE分别是各个TENSOR对应的类型名称，CODE是你想要进行的操作。例如在这里三个张量分别为r_，t，src，那么他们在循环中对应的元素指针为其名称后加_data后缀，分别为r__data，t_data, src_data。所以上面cadd函数中的这段代码的意思就是遍历相同大小的r_, t, src然后应用代码
1 | *r__data = *t_data + value * *src_data; |
这类似于一些多维数组库里的map函数，一般来说一个map函数大约长这样，由于CUDA部分是有C++的，后面就会发现在CUDA部分THC库里面大约是按照map函数的思路来封装的，而不再使用宏。
1 | map(f, array) |
CPU上的向量化操作
刚刚在cadd函数里还有一段代码，是有关于向量化操作的。很多CPU都提供了向量化指令（SIMD），这包括AVX, AVX2, SSE等等。通过支持向量化操作可以使得你的计算速度获得很大的提升（具体提升视数据类型，所占位数而定，因为寄存器的大小是固定的）。不同的CPU型号所支持的向量化指令集可能有所不同。PyTorch在支持不同CPU上使用了多重派发的方法，在运行时会自动根据当前所能使用的指令对向量化函数进行分配。在无法获得SIMD指令支持的时候会自动退回到普通的实现上。
我在支持复数的过程中简单地实现了一些对复数的SIMD指令操作，详见我的Github: CSIMD
具体还是举例说明
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 | static void (*THVector_(fill_DISPATCHPTR))(real *, const real, const ptrdiff_t) = &THVector_(fill_DEFAULT); static FunctionDescription THVector_(fill_DISPATCHTABLE)[] = { #if defined(__NEON__) #if defined(TH_REAL_IS_FLOAT) FUNCTION_IMPL(THVector_(fill_NEON), SIMDExtension_NEON), #endif #endif #if defined(__PPC64__) #if defined(TH_REAL_IS_DOUBLE) || defined(TH_REAL_IS_FLOAT) FUNCTION_IMPL(THVector_(fill_VSX), SIMDExtension_VSX), #endif #endif #if defined(USE_AVX) #if defined(TH_REAL_IS_DOUBLE) || defined(TH_REAL_IS_FLOAT) FUNCTION_IMPL(THVector_(fill_AVX), SIMDExtension_AVX), #endif #endif #if defined(USE_SSE2) || defined(USE_SSE3) || defined(USE_SSSE3) \ || defined(USE_SSE4_1) || defined(USE_SSE4_2) #if defined(TH_REAL_IS_DOUBLE) || defined(TH_REAL_IS_FLOAT) FUNCTION_IMPL(THVector_(fill_SSE), SIMDExtension_SSE), #endif #endif FUNCTION_IMPL(THVector_(fill_DEFAULT), SIMDExtension_DEFAULT) }; void THVector_(fill)(real *x, const real c, const ptrdiff_t n) { THVector_(fill_DISPATCHPTR)(x, c, n); } |
以上对于fill这个操作，实现了NEON,PPC64,AVX,SSE2,SSE3,SSSE3,SSE4指令的支持，其具体实现分别在THVector_(fill_ARCH)里，这里ARCH代表具体的SIMD指令型号。在编译时会编译所有支持的指令，但是具体使用时会按照以上的声明顺序进行调用，ARCH为DEFAULT的函数是默认实现，没有向量化支持，优先级最低。
具体如何使用SIMD指令由于指令集不同，并且读了指令集文档之后使用起来并不困难，不做介绍。
到了具体在表达式中使用时，PyTorch实现了另外一个宏，它会将内部的操作用向量化指令加速，然后再使用openmp的轻量级线程进一步加速。
1 | TH_TENSOR_APPLY_CONTIG(TYPE, TENSOR, CODE) |
这个宏内已经完成了openmp的相关操作，所以在使用的时候非常方便，非常顺滑。
CUDA张量后端THC
THC除了使用之前提到的通过C的宏命令产生泛型的方法以外，还使用了cmake命令进行简单的代码生成。一般来说一个THC的部分会有四个部分组成：C头文件 xxx.h，C源文件 xxx.c，CUDA C++头文件和源文件 xxx.cuh, xxx.cu.
THC中重新对存储在GPU上的张量进行了定义，分别为THCStorage和THCTensor。其结构类似于TH中的结构，但是注意在Copy的实现上，THCStorage的copy是依赖于THCTensor的，而非TH中THTensor依赖于THStorage。
类似于TH中，为了实现元素遍历，在THC中实现了几个reduce函数用来完成类似于TH_TENSOR_APPLY宏的操作。但是这里更专业一些。
这一部分放在下一篇文章吧。完了讲完这个再说sparse部分和python胶水那部分。去吃饭了。
源码浅析系列目录
文章来源：罗秀哲知乎专栏
本站微信群、QQ群(三群号 726282629)：
