modify ROIAlign (#1540) (#1589)

Author: jiayisunx

csrc/cpu/aten/ROIAlign.h

@@ -80,10 +80,10 @@ namespace {
 
 template <typename T>
 struct PreCalc {
-  int pos1;
-  int pos2;
-  int pos3;
-  int pos4;
+  int64_t pos1;
+  int64_t pos2;
+  int64_t pos3;
+  int64_t pos4;
   T w1;
   T w2;
   T w3;
@@ -94,27 +94,27 @@ template <typename T, typename ACC_T>
 inline void roi_align_single_framework_forward(
     const T* input,
     const ACC_T count,
-    int channels,
-    int height,
-    int width,
-    int pooled_height,
-    int pooled_width,
-    int roi_bin_grid_h,
-    int roi_bin_grid_w,
+    int64_t channels,
+    int64_t height,
+    int64_t width,
+    int64_t pooled_height,
+    int64_t pooled_width,
+    int64_t roi_bin_grid_h,
+    int64_t roi_bin_grid_w,
     const std::vector<PreCalc<ACC_T>>& pre_calc,
     T* output);
 
 template <typename T, typename ACC_T>
 inline void roi_align_single_framework_channels_last_forward(
     const T* input,
     const ACC_T count,
-    int channels,
-    int height,
-    int width,
-    int pooled_height,
-    int pooled_width,
-    int roi_bin_grid_h,
-    int roi_bin_grid_w,
+    int64_t channels,
+    int64_t height,
+    int64_t width,
+    int64_t pooled_height,
+    int64_t pooled_width,
+    int64_t roi_bin_grid_h,
+    int64_t roi_bin_grid_w,
     const std::vector<PreCalc<ACC_T>>& pre_calc,
     T* output);
 
@@ -124,13 +124,13 @@ inline void roi_align_single_framework_channels_last_forward<
     float>(
     const at::BFloat16* input,
     const float count,
-    int channels,
-    int height,
-    int width,
-    int pooled_height,
-    int pooled_width,
-    int roi_bin_grid_h,
-    int roi_bin_grid_w,
+    int64_t channels,
+    int64_t height,
+    int64_t width,
+    int64_t pooled_height,
+    int64_t pooled_width,
+    int64_t roi_bin_grid_h,
+    int64_t roi_bin_grid_w,
     const std::vector<PreCalc<float>>& pre_calc,
     at::BFloat16* output);
 
@@ -141,57 +141,57 @@ template <typename T, typename ACC_T>
 inline void roi_align_single_framework_backward(
     const T* grad_output,
     const ACC_T count,
-    int channels,
-    int height,
-    int width,
-    int pooled_height,
-    int pooled_width,
-    int roi_bin_grid_h,
-    int roi_bin_grid_w,
+    int64_t channels,
+    int64_t height,
+    int64_t width,
+    int64_t pooled_height,
+    int64_t pooled_width,
+    int64_t roi_bin_grid_h,
+    int64_t roi_bin_grid_w,
     const std::vector<PreCalc<ACC_T>>& pre_calc,
     T* grad_input);
 
 template <typename T, typename ACC_T>
 inline void roi_align_single_framework_channels_last_backward(
     const T* grad_output,
     const ACC_T count,
-    int channels,
-    int height,
-    int width,
-    int pooled_height,
-    int pooled_width,
-    int roi_bin_grid_h,
-    int roi_bin_grid_w,
+    int64_t channels,
+    int64_t height,
+    int64_t width,
+    int64_t pooled_height,
+    int64_t pooled_width,
+    int64_t roi_bin_grid_h,
+    int64_t roi_bin_grid_w,
     const std::vector<PreCalc<ACC_T>>& pre_calc,
     T* grad_input);
 
 template <typename T, typename ACC_T>
 void roi_align_forward_kernel_body(
-    int n_rois,
+    int64_t n_rois,
     const T* input,
     const ACC_T& spatial_scale,
-    int channels,
-    int height,
-    int width,
-    int pooled_height,
-    int pooled_width,
-    int sampling_ratio,
+    int64_t channels,
+    int64_t height,
+    int64_t width,
+    int64_t pooled_height,
+    int64_t pooled_width,
+    int64_t sampling_ratio,
     bool aligned,
     const ACC_T* rois,
     T* output,
     bool is_channels_last);
 
 template <typename T, typename ACC_T>
 void roi_align_backward_kernel_body(
-    int n_rois,
+    int64_t n_rois,
     const T* grad_output,
     const ACC_T& spatial_scale,
-    int channels,
-    int height,
-    int width,
-    int pooled_height,
-    int pooled_width,
-    int sampling_ratio,
+    int64_t channels,
+    int64_t height,
+    int64_t width,
+    int64_t pooled_height,
+    int64_t pooled_width,
+    int64_t sampling_ratio,
     bool aligned,
     T* grad_input,
     const ACC_T* rois,