[Feature] Add NPU operator RotatedFeatureAlign (#2994)

open-mmlab · Nov 29, 2023 · 2e44eae · 2e44eae
1 parent 5040148
commit 2e44eae
Show file tree

Hide file tree

Showing 4 changed files with 59 additions and 3 deletions.
diff --git a/docs/en/understand_mmcv/ops.md b/docs/en/understand_mmcv/ops.md
@@ -41,7 +41,7 @@ We implement common ops used in detection, segmentation, etc.
 | PointsInBoxes                | √   | √    |     |     |        |
 | PointsInPolygons             |     | √    |     |     | √      |
 | PSAMask                      | √   | √    | √   |     | √      |
-| RotatedFeatureAlign          | √   | √    | √   |     |        |
+| RotatedFeatureAlign          | √   | √    | √   |     | √      |
 | RoIPointPool3d               |     | √    | √   |     |        |
 | RoIPool                      |     | √    | √   |     | √      |
 | RoIAlignRotated              | √   | √    | √   |     |        |

diff --git a/docs/zh_cn/understand_mmcv/ops.md b/docs/zh_cn/understand_mmcv/ops.md
@@ -41,7 +41,7 @@ MMCV 提供了检测、分割等任务中常用的算子
 | PointsInBoxes                | √   | √    |     |     |        |
 | PointsInPolygons             |     | √    |     |     |        |
 | PSAMask                      | √   | √    | √   |     | √      |
-| RotatedFeatureAlign          | √   | √    | √   |     |        |
+| RotatedFeatureAlign          | √   | √    | √   |     | √      |
 | RoIPointPool3d               |     | √    | √   |     |        |
 | RoIPool                      |     | √    | √   |     | √      |
 | RoIAlignRotated              | √   | √    | √   |     |        |

diff --git a/mmcv/ops/csrc/pytorch/npu/rotated_feature_align_npu.cpp b/mmcv/ops/csrc/pytorch/npu/rotated_feature_align_npu.cpp
@@ -0,0 +1,52 @@
+#include "pytorch_npu_helper.hpp"
+
+using namespace NPU_NAME_SPACE;
+using namespace std;
+
+void rotated_feature_align_forward_impl(const Tensor features,
+                                        const Tensor best_bboxes,
+                                        const float spatial_scale,
+                                        const int points, Tensor output);
+
+void rotated_feature_align_backward_impl(const Tensor top_grad,
+                                         const Tensor best_bboxes,
+                                         const float spatial_scale,
+                                         const int points, Tensor bottom_grad);
+
+void rotated_feature_align_forward_npu(const Tensor features,
+                                       const Tensor best_bboxes,
+                                       const float spatial_scale,
+                                       const int points, Tensor output) {
+  int64_t points_ = (int64_t)points;
+  at::Tensor best_bboxes_ = best_bboxes.transpose(2, 3).transpose(1, 2);
+  OpCommand cmd;
+  cmd.Name("RotatedFeatureAlign")
+      .Input(features)
+      .Input(best_bboxes_)
+      .Output(output)
+      .Attr("spatial_scale", spatial_scale)
+      .Attr("points", points_)
+      .Run();
+}
+
+void rotated_feature_align_backward_npu(const Tensor top_grad,
+                                        const Tensor best_bboxes,
+                                        const float spatial_scale,
+                                        const int points, Tensor bottom_grad) {
+  int64_t points_ = (int64_t)points;
+  at::Tensor best_bboxes_ = best_bboxes.transpose(2, 3).transpose(1, 2);
+  OpCommand cmd;
+  cmd.Name("RotatedFeatureAlignGrad")
+      .Input(top_grad)
+      .Input(best_bboxes_)
+      .Output(bottom_grad)
+      .Attr("spatial_scale", spatial_scale)
+      .Attr("points", points_)
+      .Run();
+}
+
+REGISTER_NPU_IMPL(rotated_feature_align_forward_impl,
+                  rotated_feature_align_forward_npu);
+
+REGISTER_NPU_IMPL(rotated_feature_align_backward_impl,
+                  rotated_feature_align_backward_npu);
diff --git a/tests/test_ops/test_rotated_feature_align.py b/tests/test_ops/test_rotated_feature_align.py
@@ -3,7 +3,7 @@
 import torch
 
 from mmcv.ops import rotated_feature_align
-from mmcv.utils import IS_CUDA_AVAILABLE, IS_MLU_AVAILABLE
+from mmcv.utils import IS_CUDA_AVAILABLE, IS_MLU_AVAILABLE, IS_NPU_AVAILABLE
 
 
 @pytest.mark.skipif(
@@ -17,6 +17,10 @@
         'mlu',
         marks=pytest.mark.skipif(
             not IS_MLU_AVAILABLE, reason='requires MLU support')),
+    pytest.param(
+        'npu',
+        marks=pytest.mark.skipif(
+            not IS_NPU_AVAILABLE, reason='requires NPU support')),
     pytest.param(
         'cpu',
         marks=pytest.mark.skipif(