[Feature] torch_npu support aclnn and add op (#2997)

open-mmlab · Feb 1, 2024 · 5494299 · 5494299
1 parent 4c01b02
commit 5494299
Show file tree

Hide file tree

Showing 16 changed files with 904 additions and 68 deletions.
diff --git a/docs/en/understand_mmcv/ops.md b/docs/en/understand_mmcv/ops.md
@@ -11,7 +11,7 @@ We implement common ops used in detection, segmentation, etc.
 | BorderAlign                  |     | √    |     |     |        |
 | BoxIouRotated                | √   | √    | √   |     | √      |
 | BoxIouQuadri                 | √   | √    |     |     |        |
-| CARAFE                       |     | √    | √   |     |        |
+| CARAFE                       |     | √    | √   |     | √      |
 | ChamferDistance              |     | √    |     |     |        |
 | CrissCrossAttention          |     | √    |     |     |        |
 | ContourExpand                | √   |      |     |     |        |
@@ -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
@@ -12,7 +12,7 @@ MMCV 提供了检测、分割等任务中常用的算子
 | BoxIouRotated                | √   | √    | √   |     | √      |
 | BoxIouQuadri                 | √   | √    |     |     |        |
 | CARAFE                       |     | √    | √   |     |        |
-| ChamferDistance              |     | √    |     |     |        |
+| ChamferDistance              |     | √    |     |     | √      |
 | CrissCrossAttention          |     | √    |     |     |        |
 | ContourExpand                | √   |      |     |     |        |
 | ConvexIoU                    |     | √    |     |     |        |
@@ -41,7 +41,7 @@ MMCV 提供了检测、分割等任务中常用的算子
 | PointsInBoxes                | √   | √    |     |     |        |
 | PointsInPolygons             |     | √    |     |     |        |
 | PSAMask                      | √   | √    | √   |     | √      |
-| RotatedFeatureAlign          | √   | √    | √   |     |        |
+| RotatedFeatureAlign          | √   | √    | √   |     | √      |
 | RoIPointPool3d               |     | √    | √   |     |        |
 | RoIPool                      |     | √    | √   |     | √      |
 | RoIAlignRotated              | √   | √    | √   |     |        |

diff --git a/mmcv/ops/chamfer_distance.py b/mmcv/ops/chamfer_distance.py
@@ -44,8 +44,8 @@ def forward(ctx, xyz1: Tensor, xyz2: Tensor) -> Sequence[Tensor]:
         xyz1 = xyz1.contiguous()
         xyz2 = xyz2.contiguous()
 
-        dist1 = torch.zeros(batch_size, n).to(device)
-        dist2 = torch.zeros(batch_size, m).to(device)
+        dist1 = torch.zeros(batch_size, n).type(xyz1.dtype).to(device)
+        dist2 = torch.zeros(batch_size, m).type(xyz2.dtype).to(device)
         idx1 = torch.zeros(batch_size, n).type(torch.IntTensor).to(device)
         idx2 = torch.zeros(batch_size, m).type(torch.IntTensor).to(device)
 
@@ -81,8 +81,8 @@ def backward(ctx,
         device = grad_dist1.device
         grad_dist1 = grad_dist1.contiguous()
         grad_dist2 = grad_dist2.contiguous()
-        grad_xyz1 = torch.zeros(xyz1.size()).to(device)
-        grad_xyz2 = torch.zeros(xyz2.size()).to(device)
+        grad_xyz1 = torch.zeros(xyz1.size()).type(xyz1.dtype).to(device)
+        grad_xyz2 = torch.zeros(xyz2.size()).type(xyz2.dtype).to(device)
 
         ext_module.chamfer_distance_backward(xyz1, xyz2, idx1, idx2,
                                              grad_dist1, grad_dist2, grad_xyz1,

diff --git a/mmcv/ops/csrc/common/pytorch_npu_helper.hpp b/mmcv/ops/csrc/common/pytorch_npu_helper.hpp
@@ -23,6 +23,7 @@
 
 #include "pytorch_cpp_helper.hpp"
 #include "pytorch_device_registry.hpp"
+#include "pytorch_npu_util.hpp"
 
 #define NPU_NAME_SPACE at_npu::native