PyTorch源码浅析(二)

TensorApply宏

如同作者注释所言,tensor apply系列的宏的机制如下

  • 从最外部的角标开始,循环至第一个发生内存不连续的地址,然后将其记为张量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宏的操作。但是这里更专业一些。

可以参考这个

developer.download.nvidia.com

这一部分放在下一篇文章吧。完了讲完这个再说sparse部分和python胶水那部分。去吃饭了。

PyTorch源码浅析(目录)

来源:知乎 www.zhihu.com

作者:罗秀哲

【知乎日报】千万用户的选择,做朋友圈里的新鲜事分享大牛。
点击下载

原文链接>> http://zhuanlan.zhihu.com/p/34510108?utm_campaign=rss&utm_medium=rss&utm_source=rss&utm_content=title