[Feature] Add the support of arf op for ascend device (#2789)

open-mmlab · May 11, 2023 · e197eff · e197eff
1 parent 883d339
commit e197eff
Show file tree

Hide file tree

Showing 6 changed files with 46 additions and 5 deletions.
diff --git a/docs/en/compatibility.md b/docs/en/compatibility.md
@@ -9,7 +9,7 @@ At the same time, MMCV released [2.x](https://github.com/open-mmlab/mmcv/tree/2.
 - `mmcv.fileio` module, removed in PR [#2179](https://github.com/open-mmlab/mmcv/pull/2179). FileIO module from mmengine will be used wherever required.
 - `mmcv.runner`, `mmcv.parallel`, `mmcv. engine` and `mmcv.device`, removed in PR [#2216](https://github.com/open-mmlab/mmcv/pull/2216).
 - All classes in `mmcv.utils` (eg `Config` and `Registry`) and many functions, removed in PR [#2217](https://github.com/open-mmlab/mmcv/pull/2217). Only a few functions related to mmcv are reserved.
-- `mmcv.onnex`, `mmcv.tensorrt` modules and related functions, removed in PR [#2225](https://github.com/open-mmlab/mmcv/pull/2225).
+- `mmcv.onnx`, `mmcv.tensorrt` modules and related functions, removed in PR [#2225](https://github.com/open-mmlab/mmcv/pull/2225).
 
 (2) It added the [`mmcv.transforms`](https://github.com/open-mmlab/mmcv/tree/2.x/mmcv/transforms) data transformation module.
 

diff --git a/docs/en/understand_mmcv/ops.md b/docs/en/understand_mmcv/ops.md
@@ -4,7 +4,7 @@ We implement common ops used in detection, segmentation, etc.
 
 | Device                       | CPU | CUDA | MLU | MPS | Ascend |
 | ---------------------------- | --- | ---- | --- | --- | ------ |
-| ActiveRotatedFilter          | √   | √    |     |     |        |
+| ActiveRotatedFilter          | √   | √    |     |     | √      |
 | AssignScoreWithK             |     | √    |     |     |        |
 | BallQuery                    |     | √    | √   |     |        |
 | BBoxOverlaps                 |     | √    | √   | √   | √      |

diff --git a/docs/zh_cn/compatibility.md b/docs/zh_cn/compatibility.md
@@ -9,7 +9,7 @@ OpenMMLab 团队于 2022 年 9 月 1 日在世界人工智能大会发布了新
 - `mmcv.fileio` 模块，删除于 PR [#2179](https://github.com/open-mmlab/mmcv/pull/2179)。在需要使用 FileIO 的地方使用 mmengine 中的 FileIO 模块
 - `mmcv.runner`、`mmcv.parallel`、`mmcv.engine` 和 `mmcv.device`，删除于 PR [#2216](https://github.com/open-mmlab/mmcv/pull/2216)
 - `mmcv.utils` 的所有类（例如 `Config` 和 `Registry`）和大部分函数，删除于 PR [#2217](https://github.com/open-mmlab/mmcv/pull/2217)，只保留少数和 mmcv 相关的函数
-- `mmcv.onnex`、`mmcv.tensorrt` 模块以及相关的函数，删除于 PR [#2225](https://github.com/open-mmlab/mmcv/pull/2225)
+- `mmcv.onnx`、`mmcv.tensorrt` 模块以及相关的函数，删除于 PR [#2225](https://github.com/open-mmlab/mmcv/pull/2225)
 
 （2）新增了 [`mmcv.transforms`](https://github.com/open-mmlab/mmcv/tree/2.x/mmcv/transforms) 数据变换模块
 

diff --git a/docs/zh_cn/understand_mmcv/ops.md b/docs/zh_cn/understand_mmcv/ops.md
@@ -4,7 +4,7 @@ MMCV 提供了检测、分割等任务中常用的算子
 
 | Device                       | CPU | CUDA | MLU | MPS | Ascend |
 | ---------------------------- | --- | ---- | --- | --- | ------ |
-| ActiveRotatedFilter          | √   | √    |     |     |        |
+| ActiveRotatedFilter          | √   | √    |     |     | √      |
 | AssignScoreWithK             |     | √    |     |     |        |
 | BallQuery                    |     | √    | √   |     |        |
 | BBoxOverlaps                 |     | √    | √   | √   | √      |

diff --git a/mmcv/ops/csrc/pytorch/npu/active_rotated_filter_npu.cpp b/mmcv/ops/csrc/pytorch/npu/active_rotated_filter_npu.cpp
@@ -0,0 +1,36 @@
+#include "pytorch_npu_helper.hpp"
+
+using namespace NPU_NAME_SPACE;
+using namespace std;
+
+void active_rotated_filter_forward_impl(const Tensor input,
+                                        const Tensor indices, Tensor output);
+
+void active_rotated_filter_backward_impl(const Tensor grad_out,
+                                         const Tensor indices, Tensor grad_in);
+
+void active_rotated_filter_forward_npu(const Tensor input, const Tensor indices,
+                                       Tensor output) {
+  OpCommand cmd;
+  cmd.Name("ActiveRotatedFilter")
+      .Input(input)
+      .Input(indices)
+      .Output(output)
+      .Run();
+}
+
+void active_rotated_filter_backward_npu(const Tensor grad_out,
+                                        const Tensor indices, Tensor grad_in) {
+  OpCommand cmd;
+  cmd.Name("ActiveRotatedFilterGrad")
+      .Input(grad_out)
+      .Input(indices)
+      .Output(grad_in)
+      .Run();
+}
+
+REGISTER_NPU_IMPL(active_rotated_filter_forward_impl,
+                  active_rotated_filter_forward_npu);
+
+REGISTER_NPU_IMPL(active_rotated_filter_backward_impl,
+                  active_rotated_filter_backward_npu);
diff --git a/tests/test_ops/test_active_rotated_filter.py b/tests/test_ops/test_active_rotated_filter.py
@@ -4,6 +4,7 @@
 import torch
 
 from mmcv.ops import active_rotated_filter
+from mmcv.utils import IS_CUDA_AVAILABLE, IS_NPU_AVAILABLE
 
 np_feature = np.array([[[[[-1.4934e-01, 1.1341e+00, -1.6241e-01],
                           [-1.0986e+00, -1.1463e+00, -1.3176e+00],
@@ -245,7 +246,11 @@
     pytest.param(
         'cuda',
         marks=pytest.mark.skipif(
-            not torch.cuda.is_available(), reason='requires CUDA support')),
+            not IS_CUDA_AVAILABLE, reason='requires CUDA support')),
+    pytest.param(
+        'npu',
+        marks=pytest.mark.skipif(
+            not IS_NPU_AVAILABLE, reason='requires NPU support'))
 ])
 def test_active_rotated_filter(device):
     feature = torch.tensor(