Optimize deformable kernel

Signed-off.by: Ignacio Sañudo Olmedo<ignacio.sanudoolmedo@unimore.it>
Signed-off-by: Davide Sapienza <sapienza.dav@gmail.com>

CMakeLists.txt

@@ -17,6 +17,7 @@ endif()
 find_package(CUDA 9.0 REQUIRED)
 SET(CUDA_SEPARABLE_COMPILATION ON)
 #set(CUDA_NVCC_FLAGS "${CUDA_NVCC_FLAGS} -arch=sm_30 --compiler-options '-fPIC'")
+#set(CUDA_NVCC_FLAGS ${CUDA_NVCC_FLAGS} --maxrregcount=32)
 
 find_package(CUDNN REQUIRED)
 

include/tkDNN/kernels.h

@@ -30,14 +30,20 @@ void upsampleForward(dnnType* srcData, dnnType* dstData,
 void float2half(float* srcData, __half* dstData, int size, const cudaStream_t stream = cudaStream_t(0));
 
 
+// void modulated_deformable_im2col_cuda(cudaStream_t stream,
+//                                 const float *data_im, const float *data_offset, const float *data_mask,
+//                                 const int batch_size, const int channels, const int height_im, const int width_im,
+//                                 const int height_col, const int width_col, const int kernel_h, const int kenerl_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 deformable_group, float *data_col);
 void modulated_deformable_im2col_cuda(cudaStream_t stream,
                                 const float *data_im, const float *data_offset, const float *data_mask,
                                 const int batch_size, const int channels, const int height_im, const int width_im,
-                                const int height_col, const int width_col, const int kernel_h, const int kenerl_w,
+                                const int height_col, const int width_col, 
-                                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 deformable_group, float *data_col);
 
+
 void dcn_v2_cuda_forward(cublasStatus_t stat, cublasHandle_t handle, 
                          float *input, float *weight,
                          float *bias, float *ones,

src/kernels/deformable_conv.cu

@@ -9,7 +9,7 @@
       i < (n);                                          \
       i += blockDim.x * gridDim.x)
 
-const int CUDA_NUM_THREADS = 1024;
+const int CUDA_NUM_THREADS = 512;
 inline int GET_BLOCKS(const int N)
 {
   return (N + CUDA_NUM_THREADS - 1) / CUDA_NUM_THREADS;
@@ -17,84 +17,89 @@ inline int GET_BLOCKS(const int N)
 
 
 __device__ float dmcn_im2col_bilinear(const float *bottom_data, const int data_width,
-                                      const int height, const int width, float h, float w)
+  const int height, const int width, float h, float w)
 {
-  int h_low = floor(h);
+int h_low = floor(h);
-  int w_low = floor(w);
+int w_low = floor(w);
-  int h_high = h_low + 1;
+int h_high = h_low + 1;
-  int w_high = w_low + 1;
+int w_high = w_low + 1;
 
-  float lh = h - h_low;
+float lh = h - h_low;
-  float lw = w - w_low;
+float lw = w - w_low;
-  float hh = 1 - lh, hw = 1 - lw;
+float hh = 1 - lh, hw = 1 - lw;
 
-  float v1 = 0;
+float v1 = ( (h_low >= 0 && w_low >= 0) ? bottom_data[h_low * data_width + w_low]:0);
-  if (h_low >= 0 && w_low >= 0)
+float v2 = ( (h_low >= 0 && w_high <= width - 1) ? bottom_data[h_low * data_width + w_high]:0);
-    v1 = bottom_data[h_low * data_width + w_low];
+float v3 = ( (h_high <= height - 1 && w_low >= 0) ? bottom_data[h_high * data_width + w_low]:0);
-  float v2 = 0;
+float v4 = ( (h_high <= height - 1 && w_high <= width - 1) ? bottom_data[h_high * data_width + w_high]:0);
-  if (h_low >= 0 && w_high <= width - 1)
-    v2 = bottom_data[h_low * data_width + w_high];
-  float v3 = 0;
-  if (h_high <= height - 1 && w_low >= 0)
-    v3 = bottom_data[h_high * data_width + w_low];
-  float v4 = 0;
-  if (h_high <= height - 1 && w_high <= width - 1)
-    v4 = bottom_data[h_high * data_width + w_high];
 
-  float w1 = hh * hw, w2 = hh * lw, w3 = lh * hw, w4 = lh * lw;
+float w1 = hh * hw, w2 = hh * lw, w3 = lh * hw, w4 = lh * lw;
 
-  float val = (w1 * v1 + w2 * v2 + w3 * v3 + w4 * v4);
+float val = (w1 * v1 + w2 * v2 + w3 * v3 + w4 * v4);
-  return val;
+return val;
 }
 
 __global__ void modulated_deformable_im2col_gpu_kernel(const int n,
-                                                       const float *data_im, const float *data_offset, const float *data_mask,
+  const float *data_im, const float *data_offset, const float *data_mask,
-                                                       const int height, const int width, const int kernel_h, const int kernel_w,
+  const int height, const int width,
-                                                       const int pad_h, const int pad_w,
+  const int batch_size, const int num_channels, const int deformable_group,
-                                                       const int stride_h, const int stride_w,
+  const int height_col, const int width_col,
-                                                       const int dilation_h, const int dilation_w,
+  float *data_col)
-                                                       const int channel_per_deformable_group,
-                                                       const int batch_size, const int num_channels, const int deformable_group,
-                                                       const int height_col, const int width_col,
-                                                       float *data_col)
 {
   CUDA_KERNEL_LOOP(index, n)
   {
+    //If n is a power of 2, ( i / n ) is equivalent to ( i ≫ log2 n ) and ( i % n ) is equivalent to ( i & n - 1 ).
+    const int ind_on_w = index / width_col;
+    const int ind_on_w_on_h = ind_on_w / height_col; 
+    const int kk = 3 * 3;
     // index index of output matrix
     const int w_col = index % width_col;
-    const int h_col = (index / width_col) % height_col;
+    const int h_col = (ind_on_w) % height_col;
-    const int b_col = (index / width_col / height_col) % batch_size;
+    const int b_col = (ind_on_w_on_h) % batch_size;
-    const int c_im = (index / width_col / height_col) / batch_size;
+    const int c_im = (ind_on_w_on_h) / batch_size;
-    const int c_col = c_im * kernel_h * kernel_w;
+    const int c_col = c_im * kk;
 
     // compute deformable group index
-    const int deformable_group_index = c_im / channel_per_deformable_group;
+    const int deformable_group_index = c_im / (int)(num_channels / deformable_group);
 
-    const int h_in = h_col * stride_h - pad_h;
+    const int h_in = h_col - 1;
-    const int w_in = w_col * stride_w - pad_w;
+    const int w_in = w_col - 1;
+    const int s_col = height_col * width_col;
+    const int s_col2 = 2 * s_col;
 
     float *data_col_ptr = data_col + ((c_col * batch_size + b_col) * height_col + h_col) * width_col + w_col;
     //const float* data_im_ptr = data_im + ((b_col * num_channels + c_im) * height + h_in) * width + w_in;
     const float *data_im_ptr = data_im + (b_col * num_channels + c_im) * height * width;
-    const float *data_offset_ptr = data_offset + (b_col * deformable_group + deformable_group_index) * 2 * kernel_h * kernel_w * height_col * width_col;
+    const int add_ptr = (b_col * deformable_group + deformable_group_index) * kk * s_col;
+    const float *data_offset_ptr = data_offset + add_ptr + add_ptr;
 
-    const float *data_mask_ptr = data_mask + (b_col * deformable_group + deformable_group_index) * kernel_h * kernel_w * height_col * width_col;
+    const float *data_mask_ptr = data_mask + add_ptr;
 
-    for (int i = 0; i < kernel_h; ++i)
+    const int first_member = w_col + width_col * h_col;
+    float val = static_cast<float>(0);
+    #pragma unroll
+    for (int i = 0; i < 3; ++i)
     {
-      for (int j = 0; j < kernel_w; ++j)
+      #pragma unroll
+      for (int j = 0; j < 3; ++j)
       {
-        const int data_offset_h_ptr = ((2 * (i * kernel_w + j)) * height_col + h_col) * width_col + w_col;
+        const int iter_member = (i * 3 + j);
-        const int data_offset_w_ptr = ((2 * (i * kernel_w + j) + 1) * height_col + h_col) * width_col + w_col;
+        // const int data_offset_h_ptr = ((2 * (i * kernel_w + j)) * height_col + h_col) * width_col + w_col;
-        const int data_mask_hw_ptr = ((i * kernel_w + j) * height_col + h_col) * width_col + w_col;
+        const int data_offset_h_ptr = first_member + s_col2 * iter_member;
+        
+        // const int data_offset_w_ptr = ((2 * (i * kernel_w + j) + 1) * height_col + h_col) * width_col + w_col;
+        const int data_offset_w_ptr = s_col + first_member + s_col2 * iter_member;
+        
+        // const int data_mask_hw_ptr = ((i * kernel_w + j) * height_col + h_col) * width_col + w_col;
+        const int data_mask_hw_ptr = first_member + s_col * iter_member;
+
         const float offset_h = data_offset_ptr[data_offset_h_ptr];
         const float offset_w = data_offset_ptr[data_offset_w_ptr];
         const float mask = data_mask_ptr[data_mask_hw_ptr];
-        float val = static_cast<float>(0);
+        const float h_im = offset_h + h_in + i;
-        const float h_im = h_in + i * dilation_h + offset_h;
+        const float w_im = offset_w + w_in + j;
-        const float w_im = w_in + j * dilation_w + offset_w;
         //if (h_im >= 0 && w_im >= 0 && h_im < height && w_im < width) {
-        if (h_im > -1 && w_im > -1 && h_im < height && w_im < width)
+        if (h_im < height && w_im < width && h_im > -1 && w_im > -1)
         {
           //const float map_h = i * dilation_h + offset_h;
           //const float map_w = j * dilation_w + offset_w;
@@ -104,7 +109,89 @@ __global__ void modulated_deformable_im2col_gpu_kernel(const int n,
           val = dmcn_im2col_bilinear(data_im_ptr, width, height, width, h_im, w_im);
         }
         *data_col_ptr = val * mask;
-        data_col_ptr += batch_size * height_col * width_col;
+        data_col_ptr += batch_size * s_col;
+        //data_col_ptr += height_col * width_col;
+      }
+    }
+  }
+}
+
+__global__ void modulated_deformable_im2col_gpu_kernel2(const int n,
+  const float *data_im, const float *data_offset, const float *data_mask,
+  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 channel_per_deformable_group,
+  const int batch_size, const int num_channels, const int deformable_group,
+  const int height_col, const int width_col,
+  float *data_col)
+{
+  CUDA_KERNEL_LOOP(index, n)
+  {
+    //If n is a power of 2, ( i / n ) is equivalent to ( i ≫ log2 n ) and ( i % n ) is equivalent to ( i & n - 1 ).
+    // printf("--- %d %d %d %d %d %d %d %d\n",kernel_h, kernel_w, pad_h, pad_w, stride_h, stride_w, dilation_h, dilation_w);
+    const int ind_on_w = index / width_col;
+    const int ind_on_w_on_h = ind_on_w / height_col; 
+    const int kk = kernel_h * kernel_w;
+    // index index of output matrix
+    const int w_col = index % width_col;
+    const int h_col = (ind_on_w) % height_col;
+    const int b_col = (ind_on_w_on_h) % batch_size;
+    const int c_im = (ind_on_w_on_h) / batch_size;
+    const int c_col = c_im * kk;
+
+    // compute deformable group index
+    const int deformable_group_index = c_im / channel_per_deformable_group;
+
+    const int h_in = h_col * stride_h - pad_h;
+    const int w_in = w_col * stride_w - pad_w;
+    const int s_col = height_col * width_col;
+    const int s_col2 = 2 * s_col;
+
+    float *data_col_ptr = data_col + ((c_col * batch_size + b_col) * height_col + h_col) * width_col + w_col;
+    //const float* data_im_ptr = data_im + ((b_col * num_channels + c_im) * height + h_in) * width + w_in;
+    const float *data_im_ptr = data_im + (b_col * num_channels + c_im) * height * width;
+    const int add_ptr = (b_col * deformable_group + deformable_group_index) * kk * s_col;
+    const float *data_offset_ptr = data_offset + add_ptr + add_ptr;
+
+    const float *data_mask_ptr = data_mask + add_ptr;
+
+    const int first_member = w_col + width_col * h_col;
+    float val = static_cast<float>(0);
+    #pragma unroll
+    for (int i = 0; i < kernel_h; ++i)
+    {
+      #pragma unroll
+      for (int j = 0; j < kernel_w; ++j)
+      {
+        const int iter_member = (i * kernel_w + j);
+        // const int data_offset_h_ptr = ((2 * (i * kernel_w + j)) * height_col + h_col) * width_col + w_col;
+        const int data_offset_h_ptr = first_member + s_col2 * iter_member;
+        
+        // const int data_offset_w_ptr = ((2 * (i * kernel_w + j) + 1) * height_col + h_col) * width_col + w_col;
+        const int data_offset_w_ptr = s_col + first_member + s_col2 * iter_member;
+        
+        // const int data_mask_hw_ptr = ((i * kernel_w + j) * height_col + h_col) * width_col + w_col;
+        const int data_mask_hw_ptr = first_member + s_col * iter_member;
+
+        const float offset_h = data_offset_ptr[data_offset_h_ptr];
+        const float offset_w = data_offset_ptr[data_offset_w_ptr];
+        const float mask = data_mask_ptr[data_mask_hw_ptr];
+        const float h_im = offset_h + h_in + i * dilation_h;
+        const float w_im = offset_w + w_in + j * dilation_w;
+        //if (h_im >= 0 && w_im >= 0 && h_im < height && w_im < width) {
+        if (h_im < height && w_im < width && h_im > -1 && w_im > -1)
+        {
+          //const float map_h = i * dilation_h + offset_h;
+          //const float map_w = j * dilation_w + offset_w;
+          //const int cur_height = height - h_in;
+          //const int cur_width = width - w_in;
+          //val = dmcn_im2col_bilinear(data_im_ptr, width, cur_height, cur_width, map_h, map_w);
+          val = dmcn_im2col_bilinear(data_im_ptr, width, height, width, h_im, w_im);
+        }
+        *data_col_ptr = val * mask;
+        data_col_ptr += batch_size * s_col;
         //data_col_ptr += height_col * width_col;
       }
     }
@@ -113,6 +200,28 @@ __global__ void modulated_deformable_im2col_gpu_kernel(const int n,
 
 
 void modulated_deformable_im2col_cuda(cudaStream_t stream,
+  const float* data_im, const float* data_offset, const float* data_mask,
+  const int batch_size, const int channels, const int height_im, const int width_im, 
+  const int height_col, const int width_col,
+  const int deformable_group, float* data_col) {
+  // num_axes should be smaller than block size
+  // const int channel_per_deformable_group = channels / deformable_group;
+  const int num_kernels = channels * batch_size * height_col * width_col;
+  modulated_deformable_im2col_gpu_kernel
+      <<<GET_BLOCKS(num_kernels), CUDA_NUM_THREADS,
+          0, stream>>>(
+      num_kernels, data_im, data_offset, data_mask, height_im, width_im, 
+      batch_size, channels, deformable_group, height_col, width_col, data_col);
+  
+  cudaError_t err = cudaGetLastError();
+  if (err != cudaSuccess)
+  {
+    printf("error in modulated_deformable_im2col_cuda: %s\n", cudaGetErrorString(err));
+  }
+
+}
+
+void modulated_deformable_im2col_cuda2(cudaStream_t stream,
   const float* data_im, const float* data_offset, const float* data_mask,
   const int batch_size, const int channels, const int height_im, const int width_im, 
   const int height_col, const int width_col, const int kernel_h, const int kenerl_w,
@@ -122,7 +231,7 @@ void modulated_deformable_im2col_cuda(cudaStream_t stream,
   // num_axes should be smaller than block size
   const int channel_per_deformable_group = channels / deformable_group;
   const int num_kernels = channels * batch_size * height_col * width_col;
-  modulated_deformable_im2col_gpu_kernel
+  modulated_deformable_im2col_gpu_kernel2
       <<<GET_BLOCKS(num_kernels), CUDA_NUM_THREADS,
           0, stream>>>(
       num_kernels, data_im, data_offset, data_mask, height_im, width_im, kernel_h, kenerl_w,
@@ -137,7 +246,6 @@ void modulated_deformable_im2col_cuda(cudaStream_t stream,
 
 }
 
-
 void dcn_v2_cuda_forward(cublasStatus_t stat, cublasHandle_t handle, 
                          float *input, float *weight,
                          float *bias, float *ones,
@@ -182,9 +290,14 @@ void dcn_v2_cuda_forward(cublasStatus_t stat, cublasHandle_t handle,
                                     input, offset,
                                     mask,
                                     1, channels, height, width,
-                                    height_out, width_out, kernel_h, kernel_w,
+                                    height_out, width_out, deformable_group, columns);
-                                    pad_h, pad_w, stride_h, stride_w, dilation_h, dilation_w,
+  // modulated_deformable_im2col_cuda2(stream,
-                                    deformable_group, columns);
+  //                                   input, offset,
+  //                                   mask,
+  //                                   1, channels, height, width,
+  //                                   height_out, width_out, kernel_h, kernel_w,
+  //                                   pad_h, pad_w, stride_h, stride_w, dilation_h, dilation_w,
+  //                                   deformable_group, columns);
   
   //(k * m)  x  (m * n)
   // Y = WC