TensorApply 宏
如同作者注释所言, tensorapply 系列的宏的机制如下
从最外部的角标开始, 循环至第一个发生内存不连续的地址, 然后将其记为张量 A,A 所在的内存都是连续的, 把剩下的记为 B
然后接下来有限对 B 从最外面的角标进行遍历, 而对于A由于内存本身就是连续的, 我们直接这一整块内存进行遍历
然后为了减少循环嵌套, 将 A 中在内存上连续 (具体来说就是和 stride 乘积相等) 的维度组合到一起
注释原文如下
/*
* 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 函数来举例子
- 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_
*r__data = *t_data + value * *src_data;
这个函数首先会确认 t 和 r_的大小, 如果 r_没有声明是一个空指针, THTensor_(resizeAs)函数会按照 t 的大小分配一块新的内存给 r_这个指针 if 的第一段暂且不说, 这是为了增加向量化操作而写的代码, 我们先看通用的 TH_TENSOR_APPLY3 这个宏这个宏的声明如下
#define TH_TENSOR_APPLY3(TYPE1, TENSOR1, TYPE2, TENSOR2, TYPE3, TENSOR3, CODE)
后面的数字3是说这个宏会对三个张量进行遍历 TYPE 分别是各个 TENSOR 对应的类型名称, CODE 是你想要进行的操作例如在这里三个张量分别为 r_,t,src, 那么他们在循环中对应的元素指针为其名称后加_data 后缀, 分别为 r__data,t_data, src_data 所以上面 cadd 函数中的这段代码的意思就是遍历相同大小的 r_, t, src 然后应用代码
*r__data = *t_data + value * *src_data;
这类似于一些多维数组库里的 map 函数, 一般来说一个 map 函数大约长这样, 由于 CUDA 部分是有 C++ 的, 后面就会发现在 CUDA 部分 THC 库里面大约是按照 map 函数的思路来封装的, 而不再使用宏
map(f, array)
CPU 上的向量化操作
刚刚在 cadd 函数里还有一段代码, 是有关于向量化操作的很多 CPU 都提供了向量化指令 (SIMD), 这包括 AVX, AVX2, SSE 等等通过支持向量化操作可以使得你的计算速度获得很大的提升(具体提升视数据类型, 所占位数而定, 因为寄存器的大小是固定的) 不同的 CPU 型号所支持的向量化指令集可能有所不同 PyTorch 在支持不同 CPU 上使用了多重派发的方法, 在运行时会自动根据当前所能使用的指令对向量化函数进行分配在无法获得 SIMD 指令支持的时候会自动退回到普通的实现上
我在支持复数的过程中简单地实现了一些对复数的 SIMD 指令操作, 详见我的 Github: CSIMD
具体还是举例说明
- 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 的轻量级线程进一步加速
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 宏的操作但是这里更专业一些
可以参考这个
- http://developer.download.nvidia.com/compute/cuda/1.1-Beta/x86_website/projects/reduction/doc/reduction.pdf
- developer.download.nvidia.com
这一部分放在下一篇文章吧完了讲完这个再说 sparse 部分和 python 胶水那部分去吃饭了
PyTorch 源码浅析(目录)
来源: http://www.tuicool.com/articles/go/iyYZrqa