add points_in_box

momo609 · Feb 5, 2024 · a8b0971 · a8b0971
1 parent ae81a80
commit a8b0971
Show file tree

Hide file tree

Showing 2 changed files with 18 additions and 42 deletions.
diff --git a/mmcv/ops/csrc/pytorch/npu/ms_deform_attn.cpp b/mmcv/ops/csrc/pytorch/npu/ms_deform_attn.cpp
diff --git a/mmcv/ops/csrc/pytorch/npu/points_in_box_npu.cpp b/mmcv/ops/csrc/pytorch/npu/points_in_box_npu.cpp
@@ -0,0 +1,18 @@
+#include "pytorch_npu_helper.hpp"
+
+using namespace NPU_NAME_SPACE;
+using namespace std;
+
+void points_in_boxes_part_forward_impl_npu(int batch_size, int boxes_num,
+                                       int pts_num, const Tensor boxes,
+                                       const Tensor pts,
+                                       Tensor box_idx_of_points) {
+    c10::SmallVector<int64_t, 8> output_size = {pts.size(0), pts.size(1)};
+    auto boxes_trans = boxes.transpose(1, 2).contiguous();
+    EXEC_NPU_CMD(aclnnPointsInBox, boxes_trans, pts, box_idx_of_points);
+}
+void points_in_boxes_part_forward_impl(int batch_size, int boxes_num,
+                                       int pts_num, const Tensor boxes,
+                                       const Tensor pts,
+                                       Tensor box_idx_of_points);
+REGISTER_NPU_IMPL(points_in_boxes_part_forward_impl, points_in_boxes_part_forward_impl_npu);