MX底层im2col代码

基本的文件配置参考官方教程
自己写的一个卷积，大家可以看看

关于GPU的kernel函数的撰写是gpu运算的核心，其中涉及到一个宏CUDA_KERNEL_LOOP，它定义在src/operator/mxnet_op.h:L57，具体定义如下：

#define CUDA_KERNEL_LOOP(i, n) \
  for (int i = blockIdx.x * blockDim.x + threadIdx.x; \ // blockIdx.x: 一个线程格grid在x维度线程块block的索引，
                                                        // blockDim.x: 一个线程块block在x维度上的线程thread数量, 
                                                                      //在mx中往往等于kBaseThreadNum
                                                        // threadIdx.x: 一个线程块block在x维度上线程thread的索引
      i < (n); \                                        // n代表核函数所要处理的元素总个数
      i += blockDim.x * gridDim.x)                      //gridDim.x: 一个线程格grid在x维度上的block数量, 
                                                                     //由核函数<<<>>>中的第一个参数决定
                                                                     //往往等于cuda_get_num_blocks(N)

说一下，根据这个指出，blockDim.x* gridDim.x表示的是该线程格所有线程的数量，n表示核函数总共要处理的元素个数有时候，n会大于blockDim.x* gridDim.x，因此并不能一个线程处理一个元素。由此通过上面的方法，让一个线程串行（for循环）处理几个元素。这其实是常用的伎俩，得借鉴学习一下。在incubator-mxnet/3rdparty/mshadow/mshadow/cuda/tensor_gpu_inl.cuh:L25-31中可以看到，每个block的线程数最多可以达到1024，在L41中可以看到grid的最大个数为65535。

这里涉及到cuda中的线程概念，可以参考这篇博文，还有这篇，或者谷歌搜索线程格，也有很多解释。另外发现一个写的很棒的博主，对cuda编程有较深的理解和通俗的解释。

下面看im2col_gpu_kernel这个函数：

template <typename DType>
__global__ void im2col_gpu_kernel(const int n, const DType* data_im, // __global__ 表示共享内存，并行调用的关键字，gpu核函数必须以此声明
                                                                     // n 代表conv kernel的个数，n=Cin*H*W
                                                                     // DType*  data_im 代表一个张量，为输入图像(Cin,H,W)，以_im为后缀的都与输入图像相关
    const int height, const int width, const int kernel_h, const int kernel_w, //feature map的高宽和kernel的高宽
    const int pad_h, const int pad_w,
    const int stride_h, const int stride_w,
    const int dilation_h, const int dilation_w,
    const int height_col, const int width_col, // 输出column的高宽
    DType* data_col) { // 输出column张量，(Kh*Kw*Cin,H,W)
  CUDA_KERNEL_LOOP(index, n) { // 在0-gridDim.x*blockDim.x之间并行，到n结束，index代表第i个conv kernel的索引
    // index index of output matrix
    // 说明一下，若%某个维度的大小，表示要索引这个维度上的位置，而/这个维度的大小则要索引下个维度的位置
    // 这里说的下个维度，是从右向左。如这里的data_im，shape为(Cin,H,W)，那么顺序为W，H，Cin
    // / 是总的索引，%是在某个维度上某个范围内的索引
    const int h_index = index / width_col;
    const int h_col = h_index % height_col; // // 在某个c_in的维度下高的索引
    const int w_col = index % width_col;    // 在某个高的维度下宽的索引
    const int c_im = h_index / height_col;  // 输入通道索引
    const int c_col = c_im * kernel_h * kernel_w; // 输出通道索引
    const int h_offset = h_col * stride_h - pad_h; // 输出h的偏移
    const int w_offset = w_col * stride_w - pad_w; // 输出w的偏移
    DType* data_col_ptr = data_col;  //获得输出张量的指针拷贝
    // 指针向前移动，由于index是0-Cin*H*W，c_col,h_col和w_col有Cin、H和W种取值，正好对应index
    data_col_ptr += (c_col * height_col + h_col) * width_col + w_col; 
    const DType* data_im_ptr = data_im; //获取输入张量的指针拷贝
    data_im_ptr += (c_im * height + h_offset) * width + w_offset; //指针向前移动
    for (int i = 0; i < kernel_h; ++i) {
      for (int j = 0; j < kernel_w; ++j) { // 对单个kernel进行循环
        int h_im = h_offset + i * dilation_h;
        int w_im = w_offset + j * dilation_w;
        *data_col_ptr = // *+指针是只取指针所指位置的数值，这里赋值给对应位置
            (h_im >= 0 && w_im >= 0 && h_im < height && w_im < width) ? // 若索引不越界
            data_im_ptr[i * dilation_h * width + j * dilation_w] : static_cast<DType>(0);
        data_col_ptr += height_col * width_col;
      }
    }
  }
}

然后是col2im

template <typename DType>
__global__ void col2im_gpu_kernel(const int n, const DType* data_col,
    const int channels, const int height, const int width,
    const int kernel_h, const int kernel_w,
    const int pad_h, const int pad_w,
    const int stride_h, const int stride_w,
    const int dilation_h, const int dilation_w,
    const int height_col, const int width_col,
    DType* data_im, OpReqType req) {
  CUDA_KERNEL_LOOP(index, n) {
    DType val = 0;
    const int w_im = index % width + pad_w;
    const int h_im = (index / width) % height + pad_h;
    const int c_im = index / (width * height);
    int kernel_extent_w = (kernel_w - 1) * dilation_w + 1;
    int kernel_extent_h = (kernel_h - 1) * dilation_h + 1;
    // compute the start and end of the output
    const int w_col_start =
        (w_im < kernel_extent_w) ? 0 : (w_im - kernel_extent_w) / stride_w + 1;
    const int w_col_end = min(w_im / stride_w + 1, width_col);
    const int h_col_start =
        (h_im < kernel_extent_h) ? 0 : (h_im - kernel_extent_h) / stride_h + 1;
    const int h_col_end = min(h_im / stride_h + 1, height_col);
    // TODO(caffe): use LCM of stride and dilation to avoid unnecessary loops
    for (int h_col = h_col_start; h_col < h_col_end; h_col += 1) {
      for (int w_col = w_col_start; w_col < w_col_end; w_col += 1) {
        int h_k = (h_im - h_col * stride_h);
        int w_k = (w_im - w_col * stride_w);
        if (h_k % dilation_h == 0 && w_k % dilation_w == 0) {
          h_k /= dilation_h;
          w_k /= dilation_w;
          int data_col_index = (((c_im * kernel_h + h_k) * kernel_w + w_k) *
                                height_col + h_col) * width_col + w_col;
          val += data_col[data_col_index];
        }
      }
    }
    KERNEL_ASSIGN(data_im[index], req, val);
  }
}