[Feature] Add rotated_feature_align_cpu and update iod3d in parrots (#2027)

mmcv/ops/csrc/parrots/cudabind.cpp

@@ -564,73 +564,58 @@ REGISTER_DEVICE_IMPL(group_points_forward_impl, CUDA,
 REGISTER_DEVICE_IMPL(group_points_backward_impl, CUDA,
                      group_points_backward_cuda);
 
-void IoU3DBoxesOverlapBevForwardCUDAKernelLauncher(const int num_a,
-                                                   const Tensor boxes_a,
-                                                   const int num_b,
-                                                   const Tensor boxes_b,
-                                                   Tensor ans_overlap);
+void IoU3DBoxesIoU3DForwardCUDAKernelLauncher(const int num_a,
+                                              const Tensor boxes_a,
+                                              const int num_b,
+                                              const Tensor boxes_b,
+                                              Tensor ans_iou);
 
-void IoU3DBoxesIoUBevForwardCUDAKernelLauncher(const int num_a,
-                                               const Tensor boxes_a,
-                                               const int num_b,
-                                               const Tensor boxes_b,
-                                               Tensor ans_iou);
+void IoU3DNMS3DForwardCUDAKernelLauncher(const Tensor boxes,
+                                         unsigned long long* mask,
+                                         int boxes_num,
+                                         float nms_overlap_thresh);
 
-void IoU3DNMSForwardCUDAKernelLauncher(const Tensor boxes,
-                                       unsigned long long* mask, int boxes_num,
-                                       float nms_overlap_thresh);
+void IoU3DNMS3DNormalForwardCUDAKernelLauncher(const Tensor boxes,
+                                               unsigned long long* mask,
+                                               int boxes_num,
+                                               float nms_overlap_thresh);
 
-void IoU3DNMSNormalForwardCUDAKernelLauncher(const Tensor boxes,
-                                             unsigned long long* mask,
-                                             int boxes_num,
-                                             float nms_overlap_thresh);
-
-void iou3d_boxes_overlap_bev_forward_cuda(const int num_a, const Tensor boxes_a,
-                                          const int num_b, const Tensor boxes_b,
-                                          Tensor ans_overlap) {
-  IoU3DBoxesOverlapBevForwardCUDAKernelLauncher(num_a, boxes_a, num_b, boxes_b,
-                                                ans_overlap);
+void iou3d_boxes_iou3d_forward_cuda(const int num_a, const Tensor boxes_a,
+                                    const int num_b, const Tensor boxes_b,
+                                    Tensor ans_iou) {
+  IoU3DBoxesIoU3DForwardCUDAKernelLauncher(num_a, boxes_a, num_b, boxes_b,
+                                           ans_iou);
 };
 
-void iou3d_boxes_iou_bev_forward_cuda(const int num_a, const Tensor boxes_a,
-                                      const int num_b, const Tensor boxes_b,
-                                      Tensor ans_iou) {
-  IoU3DBoxesIoUBevForwardCUDAKernelLauncher(num_a, boxes_a, num_b, boxes_b,
-                                            ans_iou);
+void iou3d_nms3d_forward_cuda(const Tensor boxes, unsigned long long* mask,
+                              int boxes_num, float nms_overlap_thresh) {
+  IoU3DNMS3DForwardCUDAKernelLauncher(boxes, mask, boxes_num,
+                                      nms_overlap_thresh);
 };
 
-void iou3d_nms_forward_cuda(const Tensor boxes, unsigned long long* mask,
-                            int boxes_num, float nms_overlap_thresh) {
-  IoU3DNMSForwardCUDAKernelLauncher(boxes, mask, boxes_num, nms_overlap_thresh);
+void iou3d_nms3d_normal_forward_cuda(const Tensor boxes,
+                                     unsigned long long* mask, int boxes_num,
+                                     float nms_overlap_thresh) {
+  IoU3DNMS3DNormalForwardCUDAKernelLauncher(boxes, mask, boxes_num,
+                                            nms_overlap_thresh);
 };
 
-void iou3d_nms_normal_forward_cuda(const Tensor boxes, unsigned long long* mask,
-                                   int boxes_num, float nms_overlap_thresh) {
-  IoU3DNMSNormalForwardCUDAKernelLauncher(boxes, mask, boxes_num,
-                                          nms_overlap_thresh);
-};
+void iou3d_boxes_iou3d_forward_impl(const int num_a, const Tensor boxes_a,
+                                    const int num_b, const Tensor boxes_b,
+                                    Tensor ans_iou);
 
-void iou3d_boxes_overlap_bev_forward_impl(const int num_a, const Tensor boxes_a,
-                                          const int num_b, const Tensor boxes_b,
-                                          Tensor ans_overlap);
+void iou3d_nms3d_forward_impl(const Tensor boxes, unsigned long long* mask,
+                              int boxes_num, float nms_overlap_thresh);
 
-void iou3d_boxes_iou_bev_forward_impl(const int num_a, const Tensor boxes_a,
-                                      const int num_b, const Tensor boxes_b,
-                                      Tensor ans_iou);
+void iou3d_nms3d_normal_forward_impl(const Tensor boxes,
+                                     unsigned long long* mask, int boxes_num,
+                                     float nms_overlap_thresh);
 
-void iou3d_nms_forward_impl(const Tensor boxes, unsigned long long* mask,
-                            int boxes_num, float nms_overlap_thresh);
-
-void iou3d_nms_normal_forward_impl(const Tensor boxes, unsigned long long* mask,
-                                   int boxes_num, float nms_overlap_thresh);
-
-REGISTER_DEVICE_IMPL(iou3d_boxes_overlap_bev_forward_impl, CUDA,
-                     iou3d_boxes_overlap_bev_forward_cuda);
-REGISTER_DEVICE_IMPL(iou3d_boxes_iou_bev_forward_impl, CUDA,
-                     iou3d_boxes_iou_bev_forward_cuda);
-REGISTER_DEVICE_IMPL(iou3d_nms_forward_impl, CUDA, iou3d_nms_forward_cuda);
-REGISTER_DEVICE_IMPL(iou3d_nms_normal_forward_impl, CUDA,
-                     iou3d_nms_normal_forward_cuda);
+REGISTER_DEVICE_IMPL(iou3d_boxes_iou3d_forward_impl, CUDA,
+                     iou3d_boxes_iou3d_forward_cuda);
+REGISTER_DEVICE_IMPL(iou3d_nms3d_forward_impl, CUDA, iou3d_nms3d_forward_cuda);
+REGISTER_DEVICE_IMPL(iou3d_nms3d_normal_forward_impl, CUDA,
+                     iou3d_nms3d_normal_forward_cuda);
 
 void KNNForwardCUDAKernelLauncher(int b, int n, int m, int nsample,
                                   const Tensor xyz, const Tensor new_xyz,

mmcv/ops/csrc/parrots/iou3d.cpp

@@ -12,59 +12,39 @@ All Rights Reserved 2019-2020.
 
 const int THREADS_PER_BLOCK_NMS = sizeof(unsigned long long) * 8;
 
-void iou3d_boxes_overlap_bev_forward_impl(const int num_a, const Tensor boxes_a,
-                                          const int num_b, const Tensor boxes_b,
-                                          Tensor ans_overlap) {
-  DISPATCH_DEVICE_IMPL(iou3d_boxes_overlap_bev_forward_impl, num_a, boxes_a,
-                       num_b, boxes_b, ans_overlap);
-}
-
-void iou3d_boxes_iou_bev_forward_impl(const int num_a, const Tensor boxes_a,
-                                      const int num_b, const Tensor boxes_b,
-                                      Tensor ans_iou) {
-  DISPATCH_DEVICE_IMPL(iou3d_boxes_iou_bev_forward_impl, num_a, boxes_a, num_b,
+void iou3d_boxes_iou3d_forward_impl(const int num_a, const Tensor boxes_a,
+                                    const int num_b, const Tensor boxes_b,
+                                    Tensor ans_iou) {
+  DISPATCH_DEVICE_IMPL(iou3d_boxes_iou3d_forward_impl, num_a, boxes_a, num_b,
                        boxes_b, ans_iou);
 }
 
-void iou3d_nms_forward_impl(const Tensor boxes, unsigned long long *mask,
-                            int boxes_num, float nms_overlap_thresh) {
-  DISPATCH_DEVICE_IMPL(iou3d_nms_forward_impl, boxes, mask, boxes_num,
+void iou3d_nms3d_forward_impl(const Tensor boxes, unsigned long long *mask,
+                              int boxes_num, float nms_overlap_thresh) {
+  DISPATCH_DEVICE_IMPL(iou3d_nms3d_forward_impl, boxes, mask, boxes_num,
                        nms_overlap_thresh);
 }
 
-void iou3d_nms_normal_forward_impl(const Tensor boxes, unsigned long long *mask,
-                                   int boxes_num, float nms_overlap_thresh) {
-  DISPATCH_DEVICE_IMPL(iou3d_nms_normal_forward_impl, boxes, mask, boxes_num,
+void iou3d_nms3d_normal_forward_impl(const Tensor boxes,
+                                     unsigned long long *mask, int boxes_num,
+                                     float nms_overlap_thresh) {
+  DISPATCH_DEVICE_IMPL(iou3d_nms3d_normal_forward_impl, boxes, mask, boxes_num,
                        nms_overlap_thresh);
 }
 
-void iou3d_boxes_overlap_bev_forward(Tensor boxes_a, Tensor boxes_b,
-                                     Tensor ans_overlap) {
-  // params boxes_a: (N, 5) [x1, y1, x2, y2, ry]
-  // params boxes_b: (M, 5)
-  // params ans_overlap: (N, M)
-
-  int num_a = boxes_a.size(0);
-  int num_b = boxes_b.size(0);
-
-  iou3d_boxes_overlap_bev_forward_impl(num_a, boxes_a, num_b, boxes_b,
-                                       ans_overlap);
-}
-
-void iou3d_boxes_iou_bev_forward(Tensor boxes_a, Tensor boxes_b,
-                                 Tensor ans_iou) {
-  // params boxes_a: (N, 5) [x1, y1, x2, y2, ry]
+void iou3d_boxes_iou3d_forward(Tensor boxes_a, Tensor boxes_b, Tensor ans_iou) {
+  // params boxes: (N, 7) [x, y, z, dx, dy, dz, heading]
   // params boxes_b: (M, 5)
   // params ans_overlap: (N, M)
   int num_a = boxes_a.size(0);
   int num_b = boxes_b.size(0);
 
-  iou3d_boxes_iou_bev_forward_impl(num_a, boxes_a, num_b, boxes_b, ans_iou);
+  iou3d_boxes_iou3d_forward_impl(num_a, boxes_a, num_b, boxes_b, ans_iou);
 }
 
-void iou3d_nms_forward(Tensor boxes, Tensor keep, Tensor keep_num,
-                       float nms_overlap_thresh) {
-  // params boxes: (N, 5) [x1, y1, x2, y2, ry]
+void iou3d_nms3d_forward(Tensor boxes, Tensor keep, Tensor keep_num,
+                         float nms_overlap_thresh) {
+  // params boxes: (N, 7) [x, y, z, dx, dy, dz, heading]
   // params keep: (N)
   CHECK_CONTIGUOUS(boxes);
   CHECK_CONTIGUOUS(keep);
@@ -80,7 +60,7 @@ void iou3d_nms_forward(Tensor boxes, Tensor keep, Tensor keep_num,
       at::empty({boxes_num, col_blocks}, boxes.options().dtype(at::kLong));
   unsigned long long *mask_data =
       (unsigned long long *)mask.data_ptr<int64_t>();
-  iou3d_nms_forward_impl(boxes, mask_data, boxes_num, nms_overlap_thresh);
+  iou3d_nms3d_forward_impl(boxes, mask_data, boxes_num, nms_overlap_thresh);
 
   at::Tensor mask_cpu = mask.to(at::kCPU);
   unsigned long long *mask_host =
@@ -106,9 +86,9 @@ void iou3d_nms_forward(Tensor boxes, Tensor keep, Tensor keep_num,
   }
 }
 
-void iou3d_nms_normal_forward(Tensor boxes, Tensor keep, Tensor keep_num,
-                              float nms_overlap_thresh) {
-  // params boxes: (N, 5) [x1, y1, x2, y2, ry]
+void iou3d_nms3d_normal_forward(Tensor boxes, Tensor keep, Tensor keep_num,
+                                float nms_overlap_thresh) {
+  // params boxes: (N, 7) [x, y, z, dx, dy, dz, heading]
   // params keep: (N)
 
   CHECK_CONTIGUOUS(boxes);
@@ -125,8 +105,8 @@ void iou3d_nms_normal_forward(Tensor boxes, Tensor keep, Tensor keep_num,
       at::empty({boxes_num, col_blocks}, boxes.options().dtype(at::kLong));
   unsigned long long *mask_data =
       (unsigned long long *)mask.data_ptr<int64_t>();
-  iou3d_nms_normal_forward_impl(boxes, mask_data, boxes_num,
-                                nms_overlap_thresh);
+  iou3d_nms3d_normal_forward_impl(boxes, mask_data, boxes_num,
+                                  nms_overlap_thresh);
 
   at::Tensor mask_cpu = mask.to(at::kCPU);
   unsigned long long *mask_host =

mmcv/ops/csrc/parrots/iou3d_parrots.cpp

@@ -8,20 +8,21 @@
 using namespace parrots;
 
 #ifdef MMCV_WITH_CUDA
-void iou3d_boxes_iou_bev_forward_cuda_parrots(
-    CudaContext& ctx, const SSElement& attr, const OperatorBase::in_list_t& ins,
-    OperatorBase::out_list_t& outs) {
+void iou3d_boxes_iou3d_forward_cuda_parrots(CudaContext& ctx,
+                                            const SSElement& attr,
+                                            const OperatorBase::in_list_t& ins,
+                                            OperatorBase::out_list_t& outs) {
   auto boxes_a = buildATensor(ctx, ins[0]);
   auto boxes_b = buildATensor(ctx, ins[1]);
 
   auto ans_iou = buildATensor(ctx, outs[0]);
 
-  iou3d_boxes_iou_bev_forward(boxes_a, boxes_b, ans_iou);
+  iou3d_boxes_iou3d_forward(boxes_a, boxes_b, ans_iou);
 }
 
-void iou3d_nms_forward_cuda_parrots(CudaContext& ctx, const SSElement& attr,
-                                    const OperatorBase::in_list_t& ins,
-                                    OperatorBase::out_list_t& outs) {
+void iou3d_nms3d_forward_cuda_parrots(CudaContext& ctx, const SSElement& attr,
+                                      const OperatorBase::in_list_t& ins,
+                                      OperatorBase::out_list_t& outs) {
   float nms_overlap_thresh;
   SSAttrs(attr).get<float>("nms_overlap_thresh", nms_overlap_thresh).done();
 
@@ -30,13 +31,13 @@ void iou3d_nms_forward_cuda_parrots(CudaContext& ctx, const SSElement& attr,
   auto keep = buildATensor(ctx, outs[0]);
   auto keep_num = buildATensor(ctx, outs[1]);
 
-  iou3d_nms_forward(boxes, keep, keep_num, nms_overlap_thresh);
+  iou3d_nms3d_forward(boxes, keep, keep_num, nms_overlap_thresh);
 }
 
-void iou3d_nms_normal_forward_cuda_parrots(CudaContext& ctx,
-                                           const SSElement& attr,
-                                           const OperatorBase::in_list_t& ins,
-                                           OperatorBase::out_list_t& outs) {
+void iou3d_nms3d_normal_forward_cuda_parrots(CudaContext& ctx,
+                                             const SSElement& attr,
+                                             const OperatorBase::in_list_t& ins,
+                                             OperatorBase::out_list_t& outs) {
   float nms_overlap_thresh;
   SSAttrs(attr).get<float>("nms_overlap_thresh", nms_overlap_thresh).done();
 
@@ -45,26 +46,26 @@ void iou3d_nms_normal_forward_cuda_parrots(CudaContext& ctx,
   auto keep = buildATensor(ctx, outs[0]);
   auto keep_num = buildATensor(ctx, outs[1]);
 
-  iou3d_nms_normal_forward(boxes, keep, keep_num, nms_overlap_thresh);
+  iou3d_nms3d_normal_forward(boxes, keep, keep_num, nms_overlap_thresh);
 }
 
-PARROTS_EXTENSION_REGISTER(iou3d_boxes_iou_bev_forward)
+PARROTS_EXTENSION_REGISTER(iou3d_boxes_iou3d_forward)
     .input(2)
     .output(1)
-    .apply(iou3d_boxes_iou_bev_forward_cuda_parrots)
+    .apply(iou3d_boxes_iou3d_forward_cuda_parrots)
     .done();
 
-PARROTS_EXTENSION_REGISTER(iou3d_nms_forward)
+PARROTS_EXTENSION_REGISTER(iou3d_nms3d_forward)
     .attr("nms_overlap_thresh")
     .input(1)
     .output(2)
-    .apply(iou3d_nms_forward_cuda_parrots)
+    .apply(iou3d_nms3d_forward_cuda_parrots)
     .done();
 
-PARROTS_EXTENSION_REGISTER(iou3d_nms_normal_forward)
+PARROTS_EXTENSION_REGISTER(iou3d_nms3d_normal_forward)
     .attr("nms_overlap_thresh")
     .input(1)
     .output(2)
-    .apply(iou3d_nms_normal_forward_cuda_parrots)
+    .apply(iou3d_nms3d_normal_forward_cuda_parrots)
     .done();
 #endif

mmcv/ops/csrc/parrots/iou3d_pytorch.h

@@ -4,13 +4,12 @@
 #include <torch/extension.h>
 using namespace at;
 
-void iou3d_boxes_iou_bev_forward(Tensor boxes_a, Tensor boxes_b,
-                                 Tensor ans_iou);
+void iou3d_boxes_iou3d_forward(Tensor boxes_a, Tensor boxes_b, Tensor ans_iou);
 
-void iou3d_nms_forward(Tensor boxes, Tensor keep, Tensor keep_num,
-                       float nms_overlap_thresh);
+void iou3d_nms3d_forward(Tensor boxes, Tensor keep, Tensor keep_num,
+                         float nms_overlap_thresh);
 
-void iou3d_nms_normal_forward(Tensor boxes, Tensor keep, Tensor keep_num,
-                              float nms_overlap_thresh);
+void iou3d_nms3d_normal_forward(Tensor boxes, Tensor keep, Tensor keep_num,
+                                float nms_overlap_thresh);
 
 #endif  // IOU_3D_PYTORCH_H

mmcv/ops/csrc/parrots/rotated_feature_align_parrots.cpp

@@ -41,12 +41,50 @@ void rotated_feature_align_backward_cuda_parrots(
                                  spatial_scale, points);
 }
 
+void rotated_feature_align_forward_cpu_parrots(
+    HostContext& ctx, const SSElement& attr, const OperatorBase::in_list_t& ins,
+    OperatorBase::out_list_t& outs) {
+  float spatial_scale;
+  int points;
+  SSAttrs(attr)
+      .get<float>("spatial_scale", spatial_scale)
+      .get<int>("points", points)
+      .done();
+
+  auto features = buildATensor(ctx, ins[0]);
+  auto best_bboxes = buildATensor(ctx, ins[1]);
+  auto output = buildATensor(ctx, outs[0]);
+  rotated_feature_align_forward(features, best_bboxes, output, spatial_scale,
+                                points);
+}
+#endif
+
+void rotated_feature_align_backward_cpu_parrots(
+    HostContext& ctx, const SSElement& attr, const OperatorBase::in_list_t& ins,
+    OperatorBase::out_list_t& outs) {
+  float spatial_scale;
+  int points;
+  SSAttrs(attr)
+      .get<float>("spatial_scale", spatial_scale)
+      .get<int>("points", points)
+      .done();
+
+  auto grad_output = buildATensor(ctx, ins[0]);
+  auto best_bboxes = buildATensor(ctx, ins[1]);
+  auto grad_input = buildATensor(ctx, outs[0]);
+  rotated_feature_align_backward(grad_output, best_bboxes, grad_input,
+                                 spatial_scale, points);
+}
+
 PARROTS_EXTENSION_REGISTER(rotated_feature_align_forward)
     .attr("spatial_scale")
     .attr("points")
     .input(2)
     .output(1)
+    .apply(rotated_feature_align_forward_cpu_parrots)
+#ifdef MMCV_WITH_CUDA
     .apply(rotated_feature_align_forward_cuda_parrots)
+#endif
     .done();
 
 PARROTS_EXTENSION_REGISTER(rotated_feature_align_backward)
@@ -54,7 +92,8 @@ PARROTS_EXTENSION_REGISTER(rotated_feature_align_backward)
     .attr("points")
     .input(2)
     .output(1)
+    .apply(rotated_feature_align_forward_cpu_parrots)
+#ifdef MMCV_WITH_CUDA
     .apply(rotated_feature_align_backward_cuda_parrots)
-    .done();
-
 #endif
+    .done();