Factorize CUDA_KERNEL_LOOP used in CUDA kernels

Signed-off-by: Serge Panev <[email protected]>
apache · Sep 18, 2019 · b40615e · b40615e
1 parent 3dacabe
commit b40615e
Show file tree

Hide file tree

Showing 5 changed files with 6 additions and 23 deletions.
diff --git a/src/operator/contrib/count_sketch.cu b/src/operator/contrib/count_sketch.cu
@@ -33,10 +33,6 @@
 #define WARPS_PER_BLOCK 1
 #define THREADS_PER_BLOCK 512
 
-#define CUDA_KERNEL_LOOP(i, n) \
-  for (int i = blockIdx.x * blockDim.x + threadIdx.x; \
-       i < (n); \
-       i += blockDim.x * gridDim.x)
 namespace mshadow {
 namespace cuda {
 // wrappers to deal with atomic add

diff --git a/src/operator/contrib/deformable_psroi_pooling.cu b/src/operator/contrib/deformable_psroi_pooling.cu
@@ -38,10 +38,6 @@
     cudaError_t error = condition; \
     CHECK_EQ(error, cudaSuccess) << " " << cudaGetErrorString(error); \
   } while (0)
-#define CUDA_KERNEL_LOOP(i, n) \
-for (int i = blockIdx.x * blockDim.x + threadIdx.x; \
-      i < (n); \
-      i += blockDim.x * gridDim.x)
 
 namespace mshadow {
 namespace cuda {

diff --git a/src/operator/contrib/psroi_pooling.cu b/src/operator/contrib/psroi_pooling.cu
@@ -39,10 +39,6 @@
     cudaError_t error = condition; \
     CHECK_EQ(error, cudaSuccess) << " " << cudaGetErrorString(error); \
   } while (0)
-#define CUDA_KERNEL_LOOP(i, n) \
-for (int i = blockIdx.x * blockDim.x + threadIdx.x; \
-      i < (n); \
-      i += blockDim.x * gridDim.x)
 
 namespace mshadow {
 namespace cuda {

diff --git a/src/operator/contrib/roi_align.cu b/src/operator/contrib/roi_align.cu
@@ -24,15 +24,12 @@
  * Adapted from Caffe2
 */
 #include "./roi_align-inl.h"
+#include "../mxnet_op.h"
 
 
 namespace mxnet {
 namespace op {
 
-#define CUDA_1D_KERNEL_LOOP(i, n)                                 \
-  for (size_t i = blockIdx.x * blockDim.x + threadIdx.x; i < (n); \
-       i += blockDim.x * gridDim.x)
-
 using namespace mshadow::cuda;
 
 // The maximum number of blocks to use in the default kernel call.
@@ -120,7 +117,7 @@ __global__ void RoIAlignForwardKernel(
     const int sampling_ratio,
     const T* bottom_rois,
     T* top_data) {
-  CUDA_1D_KERNEL_LOOP(index, nthreads) {
+  CUDA_KERNEL_LOOP(index, nthreads) {
     // (n, c, ph, pw) is an element in the pooled output
     int pw = index % pooled_width;
     int ph = (index / pooled_width) % pooled_height;
@@ -259,7 +256,7 @@ __global__ void RoIAlignBackwardKernel(
     const int sampling_ratio,
     T* bottom_diff,
     const T* bottom_rois) {
-  CUDA_1D_KERNEL_LOOP(index, nthreads) {
+  CUDA_KERNEL_LOOP(index, nthreads) {
     // (n, c, ph, pw) is an element in the pooled output
     int pw = index % pooled_width;
     int ph = (index / pooled_width) % pooled_height;
@@ -353,7 +350,7 @@ __global__ void RoIAlignBackwardKernel(
         }  // if
       }  // ix
     }  // iy
-  }  // CUDA_1D_KERNEL_LOOP
+  }  // CUDA_KERNEL_LOOP
 }  // RoIAlignBackward
 
 template<typename xpu>

diff --git a/src/operator/correlation.cu b/src/operator/correlation.cu
@@ -28,6 +28,7 @@
 #include <mshadow/cuda/reduce.cuh>
 #include <algorithm>
 #include <vector>
+#include "./mxnet_op.h"
 
 #define ROUND_OFF 50000
 #define WARPS_PER_BLOCK 1
@@ -38,10 +39,7 @@
     cudaError_t error = condition; \
     CHECK_EQ(error, cudaSuccess) << " " << cudaGetErrorString(error); \
   } while (0)
-#define CUDA_KERNEL_LOOP(i, n) \
-for (int i = blockIdx.x * blockDim.x + threadIdx.x; \
-      i < (n); \
-      i += blockDim.x * gridDim.x)
+
 namespace mshadow {
 namespace cuda {
 // == Correlation Kernel