Skip to content

Commit

Permalink
add npu roipoint_pool3d_forward (#3148)
Browse files Browse the repository at this point in the history
* add npu roipoint_pool3d_forward

* nms_rotated npu add new attr is_angle
  • Loading branch information
huhongsun committed Jul 18, 2024
1 parent a65aa0f commit 174b647
Show file tree
Hide file tree
Showing 4 changed files with 47 additions and 7 deletions.
1 change: 1 addition & 0 deletions mmcv/ops/csrc/pytorch/npu/nms_rotated_npu.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,7 @@ Tensor nms_rotated_npu(const Tensor dets, const Tensor scores,
.Output(selectedBox)
.Output(selectedIndex)
.Attr("iou_threshold", (float)iou_threshold)
.Attr("is_angle", false)
.Run();
selectedIndex = selectedIndex.to(at::kLong);
return selectedIndex;
Expand Down
38 changes: 38 additions & 0 deletions mmcv/ops/csrc/pytorch/npu/roipoint_pool3d_forward.cpp
Original file line number Diff line number Diff line change
@@ -0,0 +1,38 @@
#include "pytorch_npu_helper.hpp"

using namespace NPU_NAME_SPACE;
using namespace std;

void roipoint_pool3d_forward_impl_npu(int batch_size, int pts_num,
int boxes_num, int feature_in_len,
int sampled_pts_num, const Tensor xyz,
const Tensor boxes3d,
const Tensor pts_feature,
Tensor pooled_features,
Tensor pooled_empty_flag) {
auto points_trans = xyz.transpose(1, 2).contiguous();
auto point_features_trans = pts_feature.transpose(1, 2).contiguous();
c10::SmallVector<int64_t, SIZE> features_trans_size = {
xyz.size(0), boxes3d.size(1), xyz.size(2) + pts_feature.size(2),
sampled_pts_num};
at::Tensor pooled_features_trans =
at::empty(features_trans_size, xyz.options());
c10::SmallVector<int64_t, SIZE> empty_flag_size = {boxes3d.size(0),
boxes3d.size(1)};
EXEC_NPU_CMD(aclnnRoipointPool3dForward, points_trans, point_features_trans,
boxes3d, sampled_pts_num, pooled_features_trans,
pooled_empty_flag);
auto pooled_features_cache =
pooled_features_trans.transpose(2, 3).contiguous();
pooled_features.copy_(pooled_features_cache);
}

void roipoint_pool3d_forward_impl(int batch_size, int pts_num, int boxes_num,
int feature_in_len, int sampled_pts_num,
const Tensor xyz, const Tensor boxes3d,
const Tensor pts_feature,
Tensor pooled_features,
Tensor pooled_empty_flag);

REGISTER_NPU_IMPL(roipoint_pool3d_forward_impl,
roipoint_pool3d_forward_impl_npu);
4 changes: 0 additions & 4 deletions mmcv/ops/nms.py
Original file line number Diff line number Diff line change
Expand Up @@ -460,10 +460,6 @@ def nms_rotated(dets: Tensor,
input_labels = labels
if dets.device.type in ('npu', 'mlu'):
order = scores.new_empty(0, dtype=torch.long)
if dets.device.type == 'npu':
coefficient = 57.29578 # 180 / PI
for i in range(dets.size()[0]):
dets_cw[i][4] *= coefficient # radians to angle
keep_inds = ext_module.nms_rotated(dets_cw, scores, order, dets_cw,
input_labels, iou_threshold,
multi_label)
Expand Down
11 changes: 8 additions & 3 deletions tests/test_ops/test_roipoint_pool3d.py
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,7 @@
import torch

from mmcv.ops import RoIPointPool3d
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.parametrize('device', [
Expand All @@ -14,14 +14,19 @@
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, torch.half,
pytest.param(
torch.double,
marks=pytest.mark.skipif(
IS_MLU_AVAILABLE, reason='MLU does not support for double'))
IS_MLU_AVAILABLE or IS_NPU_AVAILABLE,
reason='MLU and NPU does not support for double'))
])
def test_roipoint(device, dtype):
points = torch.tensor(
Expand Down

0 comments on commit 174b647

Please sign in to comment.