[Feature] Add Ascend support for RoIAlign op (#2638)

* add roi_algin op

* add unit test

* update ops.md

* fix typo

---------

Co-authored-by: momo609 <963372609.qq.com>
Co-authored-by: Zaida Zhou <58739961+zhouzaida@users.noreply.github.com>

docs/en/understand_mmcv/ops.md

@@ -46,7 +46,7 @@ We implement common ops used in detection, segmentation, etc.
 | RoIPool                      |     | √    | √   |     | √      |
 | RoIAlignRotated              | √   | √    | √   |     |        |
 | RiRoIAlignRotated            |     | √    |     |     |        |
-| RoIAlign                     | √   | √    | √   |     |        |
+| RoIAlign                     | √   | √    | √   |     | √      |
 | RoIAwarePool3d               |     | √    | √   |     |        |
 | SAConv2d                     |     | √    |     |     |        |
 | SigmoidFocalLoss             |     | √    | √   |     | √      |

docs/zh_cn/understand_mmcv/ops.md

@@ -46,7 +46,7 @@ MMCV 提供了检测、分割等任务中常用的算子
 | RoIPool                      |     | √    | √   |     | √      |
 | RoIAlignRotated              | √   | √    | √   |     |        |
 | RiRoIAlignRotated            |     | √    |     |     |        |
-| RoIAlign                     | √   | √    | √   |     |        |
+| RoIAlign                     | √   | √    | √   |     | √      |
 | RoIAwarePool3d               |     | √    | √   |     |        |
 | SAConv2d                     |     | √    |     |     |        |
 | SigmoidFocalLoss             |     | √    | √   |     | √      |

mmcv/ops/csrc/pytorch/npu/roi_align_npu.cpp

@@ -0,0 +1,68 @@
+#include "pytorch_npu_helper.hpp"
+
+using namespace NPU_NAME_SPACE;
+using namespace std;
+
+void roi_align_forward_npu(Tensor input, Tensor rois, Tensor output,
+                           Tensor argmax_y, Tensor argmax_x, int aligned_height,
+                           int aligned_width, float spatial_scale,
+                           int sampling_ratio, int pool_mode, bool aligned) {
+  if (!aligned) {
+    LOG(WARNING) << "The [aligned] attr in roi_align op is false";
+  }
+  int64_t aligned_height_64 = aligned_height;
+  int64_t aligned_width_64 = aligned_width;
+  int64_t sampling_ratio_64 = sampling_ratio;
+  int64_t roi_end_mode = 0;
+  OpCommand cmd;
+  cmd.Name("ROIAlign")
+      .Input(input)
+      .Input(rois)
+      .Output(output)
+      .Attr("spatial_scale", spatial_scale)
+      .Attr("pooled_height", aligned_height_64)
+      .Attr("pooled_width", aligned_width_64)
+      .Attr("sample_num", sampling_ratio_64)
+      .Attr("roi_end_mode", roi_end_mode)
+      .Run();
+}
+
+void roi_align_backward_npu(Tensor grad_output, Tensor rois, Tensor argmax_y,
+                            Tensor argmax_x, Tensor grad_input,
+                            int aligned_height, int aligned_width,
+                            float spatial_scale, int sampling_ratio,
+                            int pool_mode, bool aligned) {
+  int64_t aligned_height_64 = aligned_height;
+  int64_t aligned_width_64 = aligned_width;
+  int64_t sampling_ratio_64 = sampling_ratio;
+  int64_t roi_end_mode = 0;
+  c10::SmallVector<int64_t, SIZE> xdiff_shape =
+      at_npu::native::array_to_small_vector(grad_input.sizes());
+  OpCommand cmd;
+  cmd.Name("ROIAlignGrad")
+      .Input(grad_output)
+      .Input(rois)
+      .Output(grad_input)
+      .Attr("xdiff_shape", xdiff_shape)
+      .Attr("pooled_width", aligned_width_64)
+      .Attr("pooled_height", aligned_height_64)
+      .Attr("spatial_scale", spatial_scale)
+      .Attr("sample_num", sampling_ratio_64)
+      .Attr("roi_end_mode", roi_end_mode)
+      .Run();
+}
+
+void roi_align_forward_impl(Tensor input, Tensor rois, Tensor output,
+                            Tensor argmax_y, Tensor argmax_x,
+                            int aligned_height, int aligned_width,
+                            float spatial_scale, int sampling_ratio,
+                            int pool_mode, bool aligned);
+
+void roi_align_backward_impl(Tensor grad_output, Tensor rois, Tensor argmax_y,
+                             Tensor argmax_x, Tensor grad_input,
+                             int aligned_height, int aligned_width,
+                             float spatial_scale, int sampling_ratio,
+                             int pool_mode, bool aligned);
+
+REGISTER_NPU_IMPL(roi_align_forward_impl, roi_align_forward_npu);
+REGISTER_NPU_IMPL(roi_align_backward_impl, roi_align_backward_npu);

tests/test_ops/test_roi_align.py

@@ -3,7 +3,7 @@ import numpy as np
 import pytest
 import torch
 
-from mmcv.utils import IS_CUDA_AVAILABLE, IS_MLU_AVAILABLE
+from mmcv.utils import IS_CUDA_AVAILABLE, IS_MLU_AVAILABLE, IS_NPU_AVAILABLE
 
 _USING_PARROTS = True
 try:
@@ -102,15 +102,19 @@ def _test_roialign_allclose(device, dtype):
     pytest.param(
         'mlu',
         marks=pytest.mark.skipif(
-            not IS_MLU_AVAILABLE, reason='requires MLU support'))
+            not IS_MLU_AVAILABLE, reason='requires MLU support')),
+    pytest.param(
+        'npu',
+        marks=pytest.mark.skipif(
+            not IS_NPU_AVAILABLE, reason='requires NPU support'))
 ])
 @pytest.mark.parametrize('dtype', [
     torch.float,
     pytest.param(
         torch.double,
         marks=pytest.mark.skipif(
-            IS_MLU_AVAILABLE,
-            reason='MLU does not support for 64-bit floating point')),
+            IS_MLU_AVAILABLE or IS_NPU_AVAILABLE,
+            reason='MLU and NPU do not support for 64-bit floating point')),
     torch.half
 ])
 def test_roialign(device, dtype):