Skip to content

Commit 8c0cdf7

Browse files
committed
use common cuda helpers in ROIAlign
1 parent 06c4018 commit 8c0cdf7

File tree

1 file changed

+4
-7
lines changed

1 file changed

+4
-7
lines changed

torchvision/csrc/cuda/ROIAlign_cuda.cu

Lines changed: 4 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -5,10 +5,7 @@
55
#include <THC/THCAtomics.cuh>
66
#include <THC/THCDeviceUtils.cuh>
77

8-
// TODO make it in a common file
9-
#define CUDA_1D_KERNEL_LOOP(i, n) \
10-
for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < n; \
11-
i += blockDim.x * gridDim.x)
8+
#include "cuda_helpers.h"
129

1310

1411
template <typename T>
@@ -48,7 +45,7 @@ __device__ T bilinear_interpolate(const T* input,
4845
T ly = y - y_low;
4946
T lx = x - x_low;
5047
T hy = 1. - ly, hx = 1. - lx;
51-
48+
5249
// do bilinear interpolation
5350
T v1 = input[y_low * width + x_low];
5451
T v2 = input[y_low * width + x_high];
@@ -171,7 +168,7 @@ __device__ void bilinear_interpolate_gradient(
171168
}
172169

173170
template <typename T>
174-
__global__ void RoIAlignBackwardFeature(const int nthreads, const T* grad_output,
171+
__global__ void RoIAlignBackward(const int nthreads, const T* grad_output,
175172
const int num_rois, const T spatial_scale,
176173
const int channels, const int height, const int width,
177174
const int pooled_height, const int pooled_width,
@@ -327,7 +324,7 @@ at::Tensor ROIAlign_backward_cuda(const at::Tensor& grad,
327324
int w_stride = grad.stride(3);
328325

329326
AT_DISPATCH_FLOATING_TYPES(grad.type(), "ROIAlign_backward", [&] {
330-
RoIAlignBackwardFeature<scalar_t><<<grid, block, 0, stream>>>(
327+
RoIAlignBackward<scalar_t><<<grid, block, 0, stream>>>(
331328
grad.numel(),
332329
grad.data<scalar_t>(),
333330
num_rois,

0 commit comments

Comments
 (0)