pybind.cpp 46.6 KB
Newer Older
1
// Copyright (c) OpenMMLab. All rights reserved
2
3
#include <torch/extension.h>

4
5
6
7
8
#include "pytorch_cpp_helper.hpp"

std::string get_compiler_version();
std::string get_compiling_cuda_version();

pc's avatar
pc committed
9
10
11
12
13
14
15
16
17
18
19
void assign_score_withk_forward(const Tensor &points, const Tensor &centers,
                                const Tensor &scores, const Tensor &knn_idx,
                                Tensor &output, int B, int N0, int N1, int M,
                                int K, int O, int aggregate);

void assign_score_withk_backward(const Tensor &grad_out, const Tensor &points,
                                 const Tensor &centers, const Tensor &scores,
                                 const Tensor &knn_idx, Tensor &grad_points,
                                 Tensor &grad_centers, Tensor &grad_scores,
                                 int B, int N0, int N1, int M, int K, int O,
                                 int aggregate);
20

21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
void carafe_naive_forward(Tensor features, Tensor masks, Tensor output,
                          int kernel_size, int group_size, int scale_factor);

void carafe_naive_backward(Tensor top_grad, Tensor features, Tensor masks,
                           Tensor bottom_grad, Tensor mask_grad,
                           int kernel_size, int group_size, int scale_factor);

void carafe_forward(Tensor features, Tensor masks, Tensor rfeatures,
                    Tensor routput, Tensor rmasks, Tensor output,
                    int kernel_size, int group_size, int scale_factor);

void carafe_backward(Tensor top_grad, Tensor rfeatures, Tensor masks,
                     Tensor rtop_grad, Tensor rbottom_grad_hs,
                     Tensor rbottom_grad, Tensor rmask_grad, Tensor bottom_grad,
                     Tensor mask_grad, int kernel_size, int group_size,
                     int scale_factor);

void deform_conv_forward(Tensor input, Tensor weight, Tensor offset,
                         Tensor output, Tensor columns, Tensor ones, int kW,
                         int kH, int dW, int dH, int padW, int padH,
                         int dilationW, int dilationH, int group,
                         int deformable_group, int im2col_step);

void deform_conv_backward_input(Tensor input, Tensor offset, Tensor gradOutput,
                                Tensor gradInput, Tensor gradOffset,
                                Tensor weight, Tensor columns, int kW, int kH,
                                int dW, int dH, int padW, int padH,
                                int dilationW, int dilationH, int group,
                                int deformable_group, int im2col_step);

void deform_conv_backward_parameters(Tensor input, Tensor offset,
                                     Tensor gradOutput, Tensor gradWeight,
                                     Tensor columns, Tensor ones, int kW,
                                     int kH, int dW, int dH, int padW, int padH,
                                     int dilationW, int dilationH, int group,
                                     int deformable_group, float scale,
                                     int im2col_step);
58
59
60
61
62
63
64
65
66
67
68
69

void deform_roi_pool_forward(Tensor input, Tensor rois, Tensor offset,
                             Tensor output, int pooled_height, int pooled_width,
                             float spatial_scale, int sampling_ratio,
                             float gamma);

void deform_roi_pool_backward(Tensor grad_output, Tensor input, Tensor rois,
                              Tensor offset, Tensor grad_input,
                              Tensor grad_offset, int pooled_height,
                              int pooled_width, float spatial_scale,
                              int sampling_ratio, float gamma);

70
71
72
void group_points_forward(Tensor points_tensor, Tensor idx_tensor,
                          Tensor out_tensor, int b, int c, int n, int npoints,
                          int nsample);
73

74
75
76
void group_points_backward(Tensor grad_out_tensor, Tensor idx_tensor,
                           Tensor grad_points_tensor, int b, int c, int n,
                           int npoints, int nsample);
77

78
79
80
void roipoint_pool3d_forward(Tensor xyz, Tensor boxes3d, Tensor pts_feature,
                             Tensor pooled_features, Tensor pooled_empty_flag);

pc's avatar
pc committed
81
82
void gather_points_forward(Tensor points_tensor, Tensor idx_tensor,
                           Tensor out_tensor, int b, int c, int n, int npoints);
83

pc's avatar
pc committed
84
85
86
void gather_points_backward(Tensor grad_out_tensor, Tensor idx_tensor,
                            Tensor grad_points_tensor, int b, int c, int n,
                            int npoints);
87

88
89
void sigmoid_focal_loss_forward(Tensor input, Tensor target, Tensor weight,
                                Tensor output, float gamma, float alpha);
90

91
92
void sigmoid_focal_loss_backward(Tensor input, Tensor target, Tensor weight,
                                 Tensor grad_input, float gamma, float alpha);
93

94
95
void softmax_focal_loss_forward(Tensor input, Tensor target, Tensor weight,
                                Tensor output, float gamma, float alpha);
96

97
98
99
void softmax_focal_loss_backward(Tensor input, Tensor target, Tensor weight,
                                 Tensor buff, Tensor grad_input, float gamma,
                                 float alpha);
100

pc's avatar
pc committed
101
102
103
void three_interpolate_forward(Tensor points_tensor, Tensor idx_tensor,
                               Tensor weight_tensor, Tensor out_tensor, int b,
                               int c, int m, int n);
104

pc's avatar
pc committed
105
106
107
void three_interpolate_backward(Tensor grad_out_tensor, Tensor idx_tensor,
                                Tensor weight_tensor, Tensor grad_points_tensor,
                                int b, int c, int n, int m);
108

pc's avatar
pc committed
109
110
111
void three_nn_forward(Tensor unknown_tensor, Tensor known_tensor,
                      Tensor dist2_tensor, Tensor idx_tensor, int b, int n,
                      int m);
112

113
114
115
void bbox_overlaps(const Tensor bboxes1, const Tensor bboxes2, Tensor ious,
                   const int mode, const bool aligned, const int offset);

pc's avatar
pc committed
116
117
void knn_forward(Tensor xyz_tensor, Tensor new_xyz_tensor, Tensor idx_tensor,
                 Tensor dist2_tensor, int b, int n, int m, int nsample);
118
119
120
121
122
123
void iou3d_boxes_overlap_bev_forward(Tensor boxes_a, Tensor boxes_b,
                                     Tensor ans_overlap);

void iou3d_boxes_iou_bev_forward(Tensor boxes_a, Tensor boxes_b,
                                 Tensor ans_iou);

q.yao's avatar
q.yao committed
124
125
void iou3d_nms_forward(Tensor boxes, Tensor keep, Tensor keep_num,
                       float nms_overlap_thresh);
126

q.yao's avatar
q.yao committed
127
128
void iou3d_nms_normal_forward(Tensor boxes, Tensor keep, Tensor keep_num,
                              float nms_overlap_thresh);
129

pc's avatar
pc committed
130
131
void furthest_point_sampling_forward(Tensor points_tensor, Tensor temp_tensor,
                                     Tensor idx_tensor, int b, int n, int m);
132

pc's avatar
pc committed
133
void furthest_point_sampling_with_dist_forward(Tensor points_tensor,
134
                                               Tensor temp_tensor,
pc's avatar
pc committed
135
136
                                               Tensor idx_tensor, int b, int n,
                                               int m);
137

138
139
140
141
void masked_im2col_forward(const Tensor im, const Tensor mask_h_idx,
                           const Tensor mask_w_idx, Tensor col,
                           const int kernel_h, const int kernel_w,
                           const int pad_h, const int pad_w);
142

143
144
145
void masked_col2im_forward(const Tensor col, const Tensor mask_h_idx,
                           const Tensor mask_w_idx, Tensor im, int height,
                           int width, int channels);
146

147
void modulated_deform_conv_forward(
148
149
150
151
152
153
    Tensor input, Tensor weight, Tensor bias, Tensor ones, Tensor offset,
    Tensor mask, Tensor output, Tensor columns, int kernel_h, int kernel_w,
    const int stride_h, const int stride_w, const int pad_h, const int pad_w,
    const int dilation_h, const int dilation_w, const int group,
    const int deformable_group, const bool with_bias);

154
void modulated_deform_conv_backward(
155
156
157
158
159
160
161
    Tensor input, Tensor weight, Tensor bias, Tensor ones, Tensor offset,
    Tensor mask, Tensor columns, Tensor grad_input, Tensor grad_weight,
    Tensor grad_bias, Tensor grad_offset, Tensor grad_mask, Tensor grad_output,
    int kernel_h, int kernel_w, int stride_h, int stride_w, int pad_h,
    int pad_w, int dilation_h, int dilation_w, int group, int deformable_group,
    const bool with_bias);

162
163
164
165
166
Tensor ms_deform_attn_forward(const Tensor &value, const Tensor &spatial_shapes,
                              const Tensor &level_start_index,
                              const Tensor &sampling_loc,
                              const Tensor &attn_weight, const int im2col_step);

pc's avatar
pc committed
167
168
169
170
171
172
173
void ms_deform_attn_backward(const Tensor &value, const Tensor &spatial_shapes,
                             const Tensor &level_start_index,
                             const Tensor &sampling_loc,
                             const Tensor &attn_weight,
                             const Tensor &grad_output, Tensor &grad_value,
                             Tensor &grad_sampling_loc,
                             Tensor &grad_attn_weight, const int im2col_step);
174

175
176
177
178
179
Tensor nms(Tensor boxes, Tensor scores, float iou_threshold, int offset);

Tensor softnms(Tensor boxes, Tensor scores, Tensor dets, float iou_threshold,
               float sigma, float min_score, int method, int offset);

180
std::vector<std::vector<int>> nms_match(Tensor dets, float iou_threshold);
181

182
std::vector<std::vector<float>> pixel_group(
183
184
185
    Tensor score, Tensor mask, Tensor embedding, Tensor kernel_label,
    Tensor kernel_contour, int kernel_region_num, float distance_threshold);

186
187
188
189
std::vector<std::vector<int>> contour_expand(Tensor kernel_mask,
                                             Tensor internal_kernel_label,
                                             int min_kernel_area,
                                             int kernel_num);
190

191
192
void roi_align_forward(Tensor input, Tensor rois, Tensor output,
                       Tensor argmax_y, Tensor argmax_x, int aligned_height,
193
194
195
                       int aligned_width, float spatial_scale,
                       int sampling_ratio, int pool_mode, bool aligned);

196
197
198
199
200
201
202
void roi_align_backward(Tensor grad_output, Tensor rois, Tensor argmax_y,
                        Tensor argmax_x, Tensor grad_input, int aligned_height,
                        int aligned_width, float spatial_scale,
                        int sampling_ratio, int pool_mode, bool aligned);

void roi_pool_forward(Tensor input, Tensor rois, Tensor output, Tensor argmax,
                      int pooled_height, int pooled_width, float spatial_scale);
203

204
205
206
void roi_pool_backward(Tensor grad_output, Tensor rois, Tensor argmax,
                       Tensor grad_input, int pooled_height, int pooled_width,
                       float spatial_scale);
207
208
209
210
211
212

void sync_bn_forward_mean(const Tensor input, Tensor mean);

void sync_bn_forward_var(const Tensor input, const Tensor mean, Tensor var);

void sync_bn_forward_output(const Tensor input, const Tensor mean,
213
214
215
                            const Tensor var, const Tensor weight,
                            const Tensor bias, Tensor running_mean,
                            Tensor running_var, Tensor norm, Tensor std,
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
                            Tensor output, float eps, float momentum,
                            int group_size);

void sync_bn_backward_param(const Tensor grad_output, const Tensor norm,
                            Tensor grad_weight, Tensor grad_bias);

void sync_bn_backward_data(const Tensor grad_output, const Tensor weight,
                           const Tensor grad_weight, const Tensor grad_bias,
                           const Tensor norm, const Tensor std,
                           Tensor grad_input);

void psamask_forward(const Tensor input, Tensor output, const int psa_type,
                     const int num_, const int h_feature, const int w_feature,
                     const int h_mask, const int w_mask, const int half_h_mask,
                     const int half_w_mask);

void psamask_backward(Tensor grad_output, const Tensor grad_input,
                      const int psa_type, const int num_, const int h_feature,
                      const int w_feature, const int h_mask, const int w_mask,
                      const int half_h_mask, const int half_w_mask);

237
void tin_shift_forward(Tensor input, Tensor shift, Tensor output);
Jintao Lin's avatar
Jintao Lin committed
238

239
void tin_shift_backward(Tensor grad_output, Tensor shift, Tensor grad_input);
Jintao Lin's avatar
Jintao Lin committed
240

pc's avatar
pc committed
241
242
243
void ball_query_forward(Tensor new_xyz_tensor, Tensor xyz_tensor,
                        Tensor idx_tensor, int b, int n, int m,
                        float min_radius, float max_radius, int nsample);
244

245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
template <unsigned NDim>
std::vector<torch::Tensor> get_indice_pairs_forward(
    torch::Tensor indices, int64_t batchSize,
    std::vector<int64_t> outSpatialShape, std::vector<int64_t> spatialShape,
    std::vector<int64_t> kernelSize, std::vector<int64_t> stride,
    std::vector<int64_t> padding, std::vector<int64_t> dilation,
    std::vector<int64_t> outPadding, int64_t _subM, int64_t _transpose);

template <unsigned NDim>
std::vector<Tensor> get_indice_pairs_backward(
    Tensor indices, Tensor gridOut, int64_t batchSize,
    std::vector<int64_t> outSpatialShape, std::vector<int64_t> spatialShape,
    std::vector<int64_t> kernelSize, std::vector<int64_t> stride,
    std::vector<int64_t> padding, std::vector<int64_t> dilation,
    std::vector<int64_t> outPadding, int64_t _subM, int64_t _transpose);

Tensor indice_conv_forward(Tensor features, Tensor filters, Tensor indicePairs,
                           Tensor indiceNum, int64_t numActOut,
                           int64_t _inverse, int64_t _subM);

std::vector<Tensor> indice_conv_backward(Tensor features, Tensor filters,
                                         Tensor outGrad, Tensor indicePairs,
                                         Tensor indiceNum, int64_t _inverse,
                                         int64_t _subM);

Tensor fused_indice_conv_batchnorm_forward(Tensor features, Tensor filters,
                                           Tensor bias, Tensor indicePairs,
                                           Tensor indiceNum, int64_t numActOut,
                                           int64_t _inverse, int64_t _subM);

Tensor indice_maxpool_forward(Tensor features, Tensor indicePairs,
                              Tensor indiceNum, int64_t numAct);

Tensor indice_maxpool_backward(Tensor features, Tensor outFeatures,
                               Tensor outGrad, Tensor indicePairs,
                               Tensor indiceNum);

282
void box_iou_rotated(const Tensor boxes1, const Tensor boxes2, Tensor ious,
283
                     const int mode_flag, const bool aligned);
284

285
286
Tensor nms_rotated(const Tensor dets, const Tensor scores, const Tensor order,
                   const Tensor dets_sorted, const float iou_threshold,
287
288
                   const int multi_label);

289
Tensor upfirdn2d(const Tensor &input, const Tensor &kernel, int up_x, int up_y,
290
291
292
                 int down_x, int down_y, int pad_x0, int pad_x1, int pad_y0,
                 int pad_y1);

293
294
Tensor fused_bias_leakyrelu(const Tensor &input, const Tensor &bias,
                            const Tensor &refer, int act, int grad, float alpha,
295
296
                            float scale);

297
298
void roi_align_rotated_forward(Tensor input, Tensor rois, Tensor output,
                               int pooled_height, int pooled_width,
299
                               float spatial_scale, int sampling_ratio,
300
301
302
303
304
                               bool aligned, bool clockwise);

void roi_align_rotated_backward(Tensor grad_output, Tensor rois,
                                Tensor grad_input, int pooled_height,
                                int pooled_width, float spatial_scale,
305
306
                                int sampling_ratio, bool aligned,
                                bool clockwise);
307

308
309
310
311
312
313
314
315
316
317
318
319
std::vector<torch::Tensor> dynamic_point_to_voxel_forward(
    const torch::Tensor &feats, const torch::Tensor &coors,
    const std::string &reduce_type);

void dynamic_point_to_voxel_backward(torch::Tensor &grad_feats,
                                     const torch::Tensor &grad_reduced_feats,
                                     const torch::Tensor &feats,
                                     const torch::Tensor &reduced_feats,
                                     const torch::Tensor &coors_idx,
                                     const torch::Tensor &reduce_count,
                                     const std::string &reduce_type);

320
321
322
323
324
void hard_voxelize_forward(const at::Tensor &points,
                           const at::Tensor &voxel_size,
                           const at::Tensor &coors_range, at::Tensor &voxels,
                           at::Tensor &coors, at::Tensor &num_points_per_voxel,
                           at::Tensor &voxel_num, const int max_points,
325
326
                           const int max_voxels, const int NDim,
                           const bool deterministic);
327

328
329
330
void dynamic_voxelize_forward(const at::Tensor &points,
                              const at::Tensor &voxel_size,
                              const at::Tensor &coors_range, at::Tensor &coors,
331
332
                              const int NDim);

333
334
335
336
337
338
339
340
void border_align_forward(const Tensor &input, const Tensor &boxes,
                          Tensor output, Tensor argmax_idx,
                          const int pool_size);

void border_align_backward(const Tensor &grad_output, const Tensor &boxes,
                           const Tensor &argmax_idx, Tensor grad_input,
                           const int pool_size);

341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
void points_in_boxes_cpu_forward(Tensor boxes_tensor, Tensor pts_tensor,
                                 Tensor pts_indices_tensor);

void points_in_boxes_part_forward(Tensor boxes_tensor, Tensor pts_tensor,
                                  Tensor box_idx_of_points_tensor);

void points_in_boxes_all_forward(Tensor boxes_tensor, Tensor pts_tensor,
                                 Tensor box_idx_of_points_tensor);

void roiaware_pool3d_forward(Tensor rois, Tensor pts, Tensor pts_feature,
                             Tensor argmax, Tensor pts_idx_of_voxels,
                             Tensor pooled_features, int pool_method);

void roiaware_pool3d_backward(Tensor pts_idx_of_voxels, Tensor argmax,
                              Tensor grad_out, Tensor grad_in, int pool_method);

357
358
359
360
361
362
363
364
365
366
void correlation_forward(Tensor input1, Tensor input2, Tensor output, int kH,
                         int kW, int patchH, int patchW, int padH, int padW,
                         int dilationH, int dilationW, int dilation_patchH,
                         int dilation_patchW, int dH, int dW);

void correlation_backward(Tensor grad_output, Tensor input1, Tensor input2,
                          Tensor grad_input1, Tensor grad_input2, int kH,
                          int kW, int patchH, int patchW, int padH, int padW,
                          int dilationH, int dilationW, int dilation_patchH,
                          int dilation_patchW, int dH, int dW);
367

368
369
370
371
372
373
374
375
376
377
void rotated_feature_align_forward(const Tensor features,
                                   const Tensor best_bboxes, Tensor output,
                                   const float spatial_scale, const int points);

void rotated_feature_align_backward(const Tensor top_grad,
                                    const Tensor best_bboxes,
                                    Tensor bottom_grad,
                                    const float spatial_scale,
                                    const int points);

378
379
380
381
382
383
384
385
386
387
388
void riroi_align_rotated_forward(Tensor features, Tensor rois, Tensor output,
                                 int pooled_height, int pooled_width,
                                 float spatial_scale, int num_samples,
                                 int num_orientations, bool clockwise);

void riroi_align_rotated_backward(Tensor top_grad, Tensor rois,
                                  Tensor bottom_grad, int pooled_height,
                                  int pooled_width, float spatial_scale,
                                  int num_samples, int num_orientations,
                                  bool clockwise);

389
390
void points_in_polygons_forward(Tensor points, Tensor polygons, Tensor output);

391
392
void min_area_polygons(const Tensor pointsets, Tensor polygons);

393
394
395
396
397
398
void active_rotated_filter_forward(const Tensor input, const Tensor indices,
                                   Tensor output);

void active_rotated_filter_backward(const Tensor grad_out, const Tensor indices,
                                    Tensor grad_in);

399
400
401
402
void convex_iou(const Tensor pointsets, const Tensor polygons, Tensor ious);

void convex_giou(const Tensor pointsets, const Tensor polygons, Tensor output);

403
404
405
406
at::Tensor diff_iou_rotated_sort_vertices_forward(at::Tensor vertices,
                                                  at::Tensor mask,
                                                  at::Tensor num_valid);

407
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
408
409
410
411
  m.def("upfirdn2d", &upfirdn2d, "upfirdn2d (CUDA)", py::arg("input"),
        py::arg("kernel"), py::arg("up_x"), py::arg("up_y"), py::arg("down_x"),
        py::arg("down_y"), py::arg("pad_x0"), py::arg("pad_x1"),
        py::arg("pad_y0"), py::arg("pad_y1"));
412
  m.def("fused_bias_leakyrelu", &fused_bias_leakyrelu,
413
414
415
        "fused_bias_leakyrelu (CUDA)", py::arg("input"), py::arg("bias"),
        py::arg("empty"), py::arg("act"), py::arg("grad"), py::arg("alpha"),
        py::arg("scale"));
416
  m.def("gather_points_forward", &gather_points_forward,
pc's avatar
pc committed
417
418
419
        "gather_points_forward", py::arg("points_tensor"),
        py::arg("idx_tensor"), py::arg("out_tensor"), py::arg("b"),
        py::arg("c"), py::arg("n"), py::arg("npoints"));
420
  m.def("gather_points_backward", &gather_points_backward,
pc's avatar
pc committed
421
422
423
        "gather_points_backward", py::arg("grad_out_tensor"),
        py::arg("idx_tensor"), py::arg("grad_points_tensor"), py::arg("b"),
        py::arg("c"), py::arg("n"), py::arg("npoints"));
424
425
426
  m.def("get_compiler_version", &get_compiler_version, "get_compiler_version");
  m.def("get_compiling_cuda_version", &get_compiling_cuda_version,
        "get_compiling_cuda_version");
427
  m.def("assign_score_withk_forward", &assign_score_withk_forward,
pc's avatar
pc committed
428
429
430
431
        "assign_score_withk_forward", py::arg("points"), py::arg("centers"),
        py::arg("scores"), py::arg("knn_idx"), py::arg("output"), py::arg("B"),
        py::arg("N0"), py::arg("N1"), py::arg("M"), py::arg("K"), py::arg("O"),
        py::arg("aggregate"));
432
  m.def("assign_score_withk_backward", &assign_score_withk_backward,
pc's avatar
pc committed
433
        "assign_score_withk_backward", py::arg("grad_out"), py::arg("points"),
434
        py::arg("centers"), py::arg("scores"), py::arg("knn_idx"),
pc's avatar
pc committed
435
436
437
438
        py::arg("grad_points"), py::arg("grad_centers"), py::arg("grad_scores"),
        py::arg("B"), py::arg("N0"), py::arg("N1"), py::arg("M"), py::arg("K"),
        py::arg("O"), py::arg("aggregate"));
  m.def("knn_forward", &knn_forward, "knn_forward", py::arg("xyz_tensor"),
439
        py::arg("new_xyz_tensor"), py::arg("idx_tensor"),
pc's avatar
pc committed
440
441
        py::arg("dist2_tensor"), py::arg("b"), py::arg("n"), py::arg("m"),
        py::arg("nsample"));
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
  m.def("carafe_naive_forward", &carafe_naive_forward, "carafe_naive_forward",
        py::arg("features"), py::arg("masks"), py::arg("output"),
        py::arg("kernel_size"), py::arg("group_size"), py::arg("scale_factor"));
  m.def("carafe_naive_backward", &carafe_naive_backward,
        "carafe_naive_backward", py::arg("top_grad"), py::arg("features"),
        py::arg("masks"), py::arg("bottom_grad"), py::arg("mask_grad"),
        py::arg("kernel_size"), py::arg("group_size"), py::arg("scale_factor"));
  m.def("carafe_forward", &carafe_forward, "carafe_forward",
        py::arg("features"), py::arg("masks"), py::arg("rfeatures"),
        py::arg("routput"), py::arg("rmasks"), py::arg("output"),
        py::arg("kernel_size"), py::arg("group_size"), py::arg("scale_factor"));
  m.def("carafe_backward", &carafe_backward, "carafe_backward",
        py::arg("top_grad"), py::arg("rfeatures"), py::arg("masks"),
        py::arg("rtop_grad"), py::arg("rbottom_grad_hs"),
        py::arg("rbottom_grad"), py::arg("rmask_grad"), py::arg("bottom_grad"),
        py::arg("mask_grad"), py::arg("kernel_size"), py::arg("group_size"),
        py::arg("scale_factor"));
  m.def("deform_conv_forward", &deform_conv_forward, "deform_conv_forward",
        py::arg("input"), py::arg("weight"), py::arg("offset"),
        py::arg("output"), py::arg("columns"), py::arg("ones"), py::arg("kW"),
462
463
        py::arg("kH"), py::arg("dW"), py::arg("dH"), py::arg("padW"),
        py::arg("padH"), py::arg("dilationW"), py::arg("dilationH"),
464
465
466
467
468
        py::arg("group"), py::arg("deformable_group"), py::arg("im2col_step"));
  m.def("deform_conv_backward_input", &deform_conv_backward_input,
        "deform_conv_backward_input", py::arg("input"), py::arg("offset"),
        py::arg("gradOutput"), py::arg("gradInput"), py::arg("gradOffset"),
        py::arg("weight"), py::arg("columns"), py::arg("kW"), py::arg("kH"),
469
        py::arg("dW"), py::arg("dH"), py::arg("padW"), py::arg("padH"),
470
471
472
473
474
475
        py::arg("dilationW"), py::arg("dilationH"), py::arg("group"),
        py::arg("deformable_group"), py::arg("im2col_step"));
  m.def("deform_conv_backward_parameters", &deform_conv_backward_parameters,
        "deform_conv_backward_parameters", py::arg("input"), py::arg("offset"),
        py::arg("gradOutput"), py::arg("gradWeight"), py::arg("columns"),
        py::arg("ones"), py::arg("kW"), py::arg("kH"), py::arg("dW"),
476
        py::arg("dH"), py::arg("padW"), py::arg("padH"), py::arg("dilationW"),
477
478
479
480
481
482
483
484
485
486
487
488
489
        py::arg("dilationH"), py::arg("group"), py::arg("deformable_group"),
        py::arg("scale"), py::arg("im2col_step"));
  m.def("deform_roi_pool_forward", &deform_roi_pool_forward,
        "deform roi pool forward", py::arg("input"), py::arg("rois"),
        py::arg("offset"), py::arg("output"), py::arg("pooled_height"),
        py::arg("pooled_width"), py::arg("spatial_scale"),
        py::arg("sampling_ratio"), py::arg("gamma"));
  m.def("deform_roi_pool_backward", &deform_roi_pool_backward,
        "deform roi pool backward", py::arg("grad_output"), py::arg("input"),
        py::arg("rois"), py::arg("offset"), py::arg("grad_input"),
        py::arg("grad_offset"), py::arg("pooled_height"),
        py::arg("pooled_width"), py::arg("spatial_scale"),
        py::arg("sampling_ratio"), py::arg("gamma"));
490
491
492
493
  m.def("roipoint_pool3d_forward", &roipoint_pool3d_forward,
        "roipoint_pool3d_forward", py::arg("xyz"), py::arg("boxes3d"),
        py::arg("pts_feature"), py::arg("pooled_features"),
        py::arg("pooled_empty_flag"));
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
  m.def("sigmoid_focal_loss_forward", &sigmoid_focal_loss_forward,
        "sigmoid_focal_loss_forward ", py::arg("input"), py::arg("target"),
        py::arg("weight"), py::arg("output"), py::arg("gamma"),
        py::arg("alpha"));
  m.def("sigmoid_focal_loss_backward", &sigmoid_focal_loss_backward,
        "sigmoid_focal_loss_backward", py::arg("input"), py::arg("target"),
        py::arg("weight"), py::arg("grad_input"), py::arg("gamma"),
        py::arg("alpha"));
  m.def("softmax_focal_loss_forward", &softmax_focal_loss_forward,
        "softmax_focal_loss_forward", py::arg("input"), py::arg("target"),
        py::arg("weight"), py::arg("output"), py::arg("gamma"),
        py::arg("alpha"));
  m.def("softmax_focal_loss_backward", &softmax_focal_loss_backward,
        "softmax_focal_loss_backward", py::arg("input"), py::arg("target"),
        py::arg("weight"), py::arg("buff"), py::arg("grad_input"),
        py::arg("gamma"), py::arg("alpha"));
510
  m.def("three_interpolate_forward", &three_interpolate_forward,
pc's avatar
pc committed
511
512
513
        "three_interpolate_forward", py::arg("points_tensor"),
        py::arg("idx_tensor"), py::arg("weight_tensor"), py::arg("out_tensor"),
        py::arg("b"), py::arg("c"), py::arg("m"), py::arg("n"));
514
  m.def("three_interpolate_backward", &three_interpolate_backward,
pc's avatar
pc committed
515
516
517
518
519
520
521
522
        "three_interpolate_backward", py::arg("grad_out_tensor"),
        py::arg("idx_tensor"), py::arg("weight_tensor"),
        py::arg("grad_points_tensor"), py::arg("b"), py::arg("c"), py::arg("n"),
        py::arg("m"));
  m.def("three_nn_forward", &three_nn_forward, "three_nn_forward",
        py::arg("unknown_tensor"), py::arg("known_tensor"),
        py::arg("dist2_tensor"), py::arg("idx_tensor"), py::arg("b"),
        py::arg("n"), py::arg("m"));
523
524
525
  m.def("bbox_overlaps", &bbox_overlaps, "bbox_overlaps", py::arg("bboxes1"),
        py::arg("bboxes2"), py::arg("ious"), py::arg("mode"),
        py::arg("aligned"), py::arg("offset"));
526
  m.def("group_points_forward", &group_points_forward, "group_points_forward",
527
        py::arg("points_tensor"), py::arg("idx_tensor"), py::arg("out_tensor"),
528
        py::arg("b"), py::arg("c"), py::arg("n"), py::arg("npoints"),
529
        py::arg("nsample"));
530
  m.def("group_points_backward", &group_points_backward,
531
532
533
        "group_points_backward", py::arg("grad_out_tensor"),
        py::arg("idx_tensor"), py::arg("grad_points_tensor"), py::arg("b"),
        py::arg("c"), py::arg("n"), py::arg("npoints"), py::arg("nsample"));
534
535
536
537
  m.def("knn_forward", &knn_forward, "knn_forward", py::arg("b"), py::arg("n"),
        py::arg("m"), py::arg("nsample"), py::arg("xyz_tensor"),
        py::arg("new_xyz_tensor"), py::arg("idx_tensor"),
        py::arg("dist2_tensor"));
538
539
540
541
542
543
544
  m.def("iou3d_boxes_overlap_bev_forward", &iou3d_boxes_overlap_bev_forward,
        "iou3d_boxes_overlap_bev_forward", py::arg("boxes_a"),
        py::arg("boxes_b"), py::arg("ans_overlap"));
  m.def("iou3d_boxes_iou_bev_forward", &iou3d_boxes_iou_bev_forward,
        "iou3d_boxes_iou_bev_forward", py::arg("boxes_a"), py::arg("boxes_b"),
        py::arg("ans_iou"));
  m.def("iou3d_nms_forward", &iou3d_nms_forward, "iou3d_nms_forward",
545
546
        py::arg("boxes"), py::arg("keep"), py::arg("num_out"),
        py::arg("nms_overlap_thresh"));
547
548
  m.def("iou3d_nms_normal_forward", &iou3d_nms_normal_forward,
        "iou3d_nms_normal_forward", py::arg("boxes"), py::arg("keep"),
549
        py::arg("num_out"), py::arg("nms_overlap_thresh"));
550
  m.def("furthest_point_sampling_forward", &furthest_point_sampling_forward,
pc's avatar
pc committed
551
552
553
        "furthest_point_sampling_forward", py::arg("points_tensor"),
        py::arg("temp_tensor"), py::arg("idx_tensor"), py::arg("b"),
        py::arg("n"), py::arg("m"));
554
555
  m.def("furthest_point_sampling_with_dist_forward",
        &furthest_point_sampling_with_dist_forward,
pc's avatar
pc committed
556
557
558
        "furthest_point_sampling_with_dist_forward", py::arg("points_tensor"),
        py::arg("temp_tensor"), py::arg("idx_tensor"), py::arg("b"),
        py::arg("n"), py::arg("m"));
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
  m.def("masked_im2col_forward", &masked_im2col_forward,
        "masked_im2col_forward", py::arg("im"), py::arg("mask_h_idx"),
        py::arg("mask_w_idx"), py::arg("col"), py::arg("kernel_h"),
        py::arg("kernel_w"), py::arg("pad_h"), py::arg("pad_w"));
  m.def("masked_col2im_forward", &masked_col2im_forward,
        "masked_col2im_forward", py::arg("col"), py::arg("mask_h_idx"),
        py::arg("mask_w_idx"), py::arg("im"), py::arg("height"),
        py::arg("width"), py::arg("channels"));
  m.def("modulated_deform_conv_forward", &modulated_deform_conv_forward,
        "modulated deform conv forward", py::arg("input"), py::arg("weight"),
        py::arg("bias"), py::arg("ones"), py::arg("offset"), py::arg("mask"),
        py::arg("output"), py::arg("columns"), py::arg("kernel_h"),
        py::arg("kernel_w"), py::arg("stride_h"), py::arg("stride_w"),
        py::arg("pad_h"), py::arg("pad_w"), py::arg("dilation_h"),
        py::arg("dilation_w"), py::arg("group"), py::arg("deformable_group"),
        py::arg("with_bias"));
  m.def("modulated_deform_conv_backward", &modulated_deform_conv_backward,
        "modulated deform conv backward", py::arg("input"), py::arg("weight"),
        py::arg("bias"), py::arg("ones"), py::arg("offset"), py::arg("mask"),
        py::arg("columns"), py::arg("grad_input"), py::arg("grad_weight"),
        py::arg("grad_bias"), py::arg("grad_offset"), py::arg("grad_mask"),
        py::arg("grad_output"), py::arg("kernel_h"), py::arg("kernel_w"),
        py::arg("stride_h"), py::arg("stride_w"), py::arg("pad_h"),
        py::arg("pad_w"), py::arg("dilation_h"), py::arg("dilation_w"),
        py::arg("group"), py::arg("deformable_group"), py::arg("with_bias"));
  m.def("nms", &nms, "nms (CPU/CUDA) ", py::arg("boxes"), py::arg("scores"),
        py::arg("iou_threshold"), py::arg("offset"));
  m.def("softnms", &softnms, "softnms (CPU) ", py::arg("boxes"),
        py::arg("scores"), py::arg("dets"), py::arg("iou_threshold"),
        py::arg("sigma"), py::arg("min_score"), py::arg("method"),
        py::arg("offset"));
  m.def("nms_match", &nms_match, "nms_match (CPU) ", py::arg("dets"),
        py::arg("iou_threshold"));
592
593
594
595
596
597
598
  m.def("pixel_group", &pixel_group, "pixel group (CPU) ", py::arg("score"),
        py::arg("mask"), py::arg("embedding"), py::arg("kernel_label"),
        py::arg("kernel_contour"), py::arg("kernel_region_label"),
        py::arg("distance_threshold"));
  m.def("contour_expand", &contour_expand, "contour exapnd (CPU) ",
        py::arg("kernel_mask"), py::arg("internal_kernel_label"),
        py::arg("min_kernel_area"), py::arg("kernel_num"));
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
  m.def("roi_align_forward", &roi_align_forward, "roi_align forward",
        py::arg("input"), py::arg("rois"), py::arg("output"),
        py::arg("argmax_y"), py::arg("argmax_x"), py::arg("aligned_height"),
        py::arg("aligned_width"), py::arg("spatial_scale"),
        py::arg("sampling_ratio"), py::arg("pool_mode"), py::arg("aligned"));
  m.def("roi_align_backward", &roi_align_backward, "roi_align backward",
        py::arg("grad_output"), py::arg("rois"), py::arg("argmax_y"),
        py::arg("argmax_x"), py::arg("grad_input"), py::arg("aligned_height"),
        py::arg("aligned_width"), py::arg("spatial_scale"),
        py::arg("sampling_ratio"), py::arg("pool_mode"), py::arg("aligned"));
  m.def("roi_pool_forward", &roi_pool_forward, "roi_pool forward",
        py::arg("input"), py::arg("rois"), py::arg("output"), py::arg("argmax"),
        py::arg("pooled_height"), py::arg("pooled_width"),
        py::arg("spatial_scale"));
  m.def("roi_pool_backward", &roi_pool_backward, "roi_pool backward",
        py::arg("grad_output"), py::arg("rois"), py::arg("argmax"),
        py::arg("grad_input"), py::arg("pooled_height"),
        py::arg("pooled_width"), py::arg("spatial_scale"));
  m.def("sync_bn_forward_mean", &sync_bn_forward_mean, "sync_bn forward_mean",
        py::arg("input"), py::arg("mean"));
  m.def("sync_bn_forward_var", &sync_bn_forward_var, "sync_bn forward_var",
        py::arg("input"), py::arg("mean"), py::arg("var"));
  m.def("sync_bn_forward_output", &sync_bn_forward_output,
        "sync_bn forward_output", py::arg("input"), py::arg("mean"),
623
624
625
        py::arg("var"), py::arg("weight"), py::arg("bias"),
        py::arg("running_mean"), py::arg("running_var"), py::arg("norm"),
        py::arg("std"), py::arg("output"), py::arg("eps"), py::arg("momentum"),
626
627
628
629
630
631
632
633
        py::arg("group_size"));
  m.def("sync_bn_backward_param", &sync_bn_backward_param,
        "sync_bn backward_param", py::arg("grad_output"), py::arg("norm"),
        py::arg("grad_weight"), py::arg("grad_bias"));
  m.def("sync_bn_backward_data", &sync_bn_backward_data,
        "sync_bn backward_data", py::arg("grad_output"), py::arg("weight"),
        py::arg("grad_weight"), py::arg("grad_bias"), py::arg("norm"),
        py::arg("std"), py::arg("grad_input"));
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
  m.def("get_indice_pairs_2d_forward", &get_indice_pairs_forward<2>,
        "get_indice_pairs_2d_forward", py::arg("indices"), py::arg("batchSize"),
        py::arg("outSpatialShape"), py::arg("spatialShape"),
        py::arg("kernelSize"), py::arg("stride"), py::arg("padding"),
        py::arg("dilation"), py::arg("outPadding"), py::arg("_subM"),
        py::arg("_transpose"));
  m.def("get_indice_pairs_3d_forward", &get_indice_pairs_forward<3>,
        "get_indice_pairs_3d_forward", py::arg("indices"), py::arg("batchSize"),
        py::arg("outSpatialShape"), py::arg("spatialShape"),
        py::arg("kernelSize"), py::arg("stride"), py::arg("padding"),
        py::arg("dilation"), py::arg("outPadding"), py::arg("_subM"),
        py::arg("_transpose"));
  m.def("get_indice_pairs_4d_forward", &get_indice_pairs_forward<4>,
        "get_indice_pairs_4d_forward", py::arg("indices"), py::arg("batchSize"),
        py::arg("outSpatialShape"), py::arg("spatialShape"),
        py::arg("kernelSize"), py::arg("stride"), py::arg("padding"),
        py::arg("dilation"), py::arg("outPadding"), py::arg("_subM"),
        py::arg("_transpose"));
  m.def("get_indice_pairs_2d_backward", &get_indice_pairs_backward<2>,
        "get_indice_pairs_2d_backward", py::arg("indices"), py::arg("gridOut"),
        py::arg("batchSize"), py::arg("outSpatialShape"),
        py::arg("spatialShape"), py::arg("kernelSize"), py::arg("stride"),
        py::arg("padding"), py::arg("dilation"), py::arg("outPadding"),
        py::arg("_subM"), py::arg("_transpose"));
  m.def("get_indice_pairs_3d_backward", &get_indice_pairs_backward<3>,
        "get_indice_pairs_3d_backward", py::arg("indices"), py::arg("gridOut"),
        py::arg("batchSize"), py::arg("outSpatialShape"),
        py::arg("spatialShape"), py::arg("kernelSize"), py::arg("stride"),
        py::arg("padding"), py::arg("dilation"), py::arg("outPadding"),
        py::arg("_subM"), py::arg("_transpose"));
  m.def("indice_conv_forward", &indice_conv_forward, "indice_conv_forward",
        py::arg("features"), py::arg("filters"), py::arg("indicePairs"),
        py::arg("indiceNum"), py::arg("numActOut"), py::arg("_inverse"),
        py::arg("_subM"));
  m.def("indice_conv_backward", &indice_conv_backward, "indice_conv_backward",
        py::arg("features"), py::arg("filters"), py::arg("outGrad"),
        py::arg("indicePairs"), py::arg("indiceNum"), py::arg("_inverse"),
        py::arg("_subM"));
  m.def("fused_indice_conv_forward", &fused_indice_conv_batchnorm_forward,
        "fused_indice_conv_forward", py::arg("features"), py::arg("filters"),
        py::arg("bias"), py::arg("indicePairs"), py::arg("indiceNum"),
        py::arg("numActOut"), py::arg("_inverse"), py::arg("_subM"));
  m.def("indice_maxpool_forward", &indice_maxpool_forward,
        "indice_maxpool_forward", py::arg("features"), py::arg("indicePairs"),
        py::arg("indiceNum"), py::arg("numAct"));
  m.def("indice_maxpool_backward", &indice_maxpool_backward,
        "indice_maxpool_backward", py::arg("features"), py::arg("outFeatures"),
        py::arg("outGrad"), py::arg("indicePairs"), py::arg("indiceNum"));
682
683
684
685
686
687
688
689
690
691
  m.def("psamask_forward", &psamask_forward, "PSAMASK forward (CPU/CUDA)",
        py::arg("input"), py::arg("output"), py::arg("psa_type"),
        py::arg("num_"), py::arg("h_feature"), py::arg("w_feature"),
        py::arg("h_mask"), py::arg("w_mask"), py::arg("half_h_mask"),
        py::arg("half_w_mask"));
  m.def("psamask_backward", &psamask_backward, "PSAMASK backward (CPU/CUDA)",
        py::arg("grad_output"), py::arg("grad_input"), py::arg("psa_type"),
        py::arg("num_"), py::arg("h_feature"), py::arg("w_feature"),
        py::arg("h_mask"), py::arg("w_mask"), py::arg("half_h_mask"),
        py::arg("half_w_mask"));
Jintao Lin's avatar
Jintao Lin committed
692
693
694
695
  m.def("tin_shift_forward", &tin_shift_forward, "tin_shift forward",
        py::arg("input"), py::arg("shift"), py::arg("output"));
  m.def("tin_shift_backward", &tin_shift_backward, "tin_shift backward",
        py::arg("grad_output"), py::arg("shift"), py::arg("grad_input"));
696
  m.def("box_iou_rotated", &box_iou_rotated, "IoU for rotated boxes",
697
        py::arg("boxes1"), py::arg("boxes2"), py::arg("ious"),
698
        py::arg("mode_flag"), py::arg("aligned"));
699
700
701
  m.def("nms_rotated", &nms_rotated, "NMS for rotated boxes", py::arg("dets"),
        py::arg("scores"), py::arg("order"), py::arg("dets_sorted"),
        py::arg("iou_threshold"), py::arg("multi_label"));
702
  m.def("ball_query_forward", &ball_query_forward, "ball_query_forward",
pc's avatar
pc committed
703
        py::arg("new_xyz_tensor"), py::arg("xyz_tensor"), py::arg("idx_tensor"),
704
        py::arg("b"), py::arg("n"), py::arg("m"), py::arg("min_radius"),
pc's avatar
pc committed
705
        py::arg("max_radius"), py::arg("nsample"));
706
707
708
  m.def("roi_align_rotated_forward", &roi_align_rotated_forward,
        "roi_align_rotated forward", py::arg("input"), py::arg("rois"),
        py::arg("output"), py::arg("pooled_height"), py::arg("pooled_width"),
709
        py::arg("spatial_scale"), py::arg("sampling_ratio"), py::arg("aligned"),
710
711
        py::arg("clockwise"));
  m.def("roi_align_rotated_backward", &roi_align_rotated_backward,
712
713
        "roi_align_rotated backward", py::arg("rois"), py::arg("grad_input"),
        py::arg("grad_output"), py::arg("pooled_height"),
714
        py::arg("pooled_width"), py::arg("spatial_scale"),
715
        py::arg("sampling_ratio"), py::arg("aligned"), py::arg("clockwise"));
716
717
718
719
720
721
722
723
724
  m.def("dynamic_point_to_voxel_forward", &dynamic_point_to_voxel_forward,
        "dynamic_point_to_voxel_forward", py::arg("feats"), py::arg("coors"),
        py::arg("reduce_type"));
  m.def("dynamic_point_to_voxel_backward", &dynamic_point_to_voxel_backward,
        "dynamic_point_to_voxel_backward", py::arg("grad_feats"),
        py::arg("grad_reduced_feats"), py::arg("feats"),
        py::arg("reduced_feats"), py::arg("coors_idx"), py::arg("reduce_count"),
        py::arg("reduce_type"));
  m.def("hard_voxelize_forward", &hard_voxelize_forward,
725
726
727
        "hard_voxelize_forward", py::arg("points"), py::arg("voxel_size"),
        py::arg("coors_range"), py::arg("voxels"), py::arg("coors"),
        py::arg("num_points_per_voxel"), py::arg("voxel_num"),
728
729
        py::arg("max_points"), py::arg("max_voxels"), py::arg("NDim"),
        py::arg("deterministic"));
730
  m.def("dynamic_voxelize_forward", &dynamic_voxelize_forward,
731
732
        "dynamic_voxelize_forward", py::arg("points"), py::arg("voxel_size"),
        py::arg("coors_range"), py::arg("coors"), py::arg("NDim"));
733
734
735
736
737
738
739
740
741
742
  m.def("ms_deform_attn_forward", &ms_deform_attn_forward,
        "forward function of multi-scale deformable attention",
        py::arg("value"), py::arg("value_spatial_shapes"),
        py::arg("value_level_start_index"), py::arg("sampling_locations"),
        py::arg("attention_weights"), py::arg("im2col_step"));
  m.def("ms_deform_attn_backward", &ms_deform_attn_backward,
        "backward function of multi-scale deformable attention",
        py::arg("value"), py::arg("value_spatial_shapes"),
        py::arg("value_level_start_index"), py::arg("sampling_locations"),
        py::arg("attention_weights"), py::arg("grad_output"),
pc's avatar
pc committed
743
744
        py::arg("grad_value"), py::arg("grad_sampling_loc"),
        py::arg("grad_attn_weight"), py::arg("im2col_step"));
745
746
747
748
749
750
751
  m.def("border_align_forward", &border_align_forward,
        "forward function of border_align", py::arg("input"), py::arg("boxes"),
        py::arg("output"), py::arg("argmax_idx"), py::arg("pool_size"));
  m.def("border_align_backward", &border_align_backward,
        "backward function of border_align", py::arg("grad_output"),
        py::arg("boxes"), py::arg("argmax_idx"), py::arg("grad_input"),
        py::arg("pool_size"));
pc's avatar
pc committed
752
753
754
755
756
757
758
759
760
761
762
763
764
  m.def("correlation_forward", &correlation_forward, "Correlation forward",
        py::arg("input1"), py::arg("input2"), py::arg("output"), py::arg("kH"),
        py::arg("kW"), py::arg("patchH"), py::arg("patchW"), py::arg("padH"),
        py::arg("padW"), py::arg("dilationH"), py::arg("dilationW"),
        py::arg("dilation_patchH"), py::arg("dilation_patchW"), py::arg("dH"),
        py::arg("dW"));
  m.def("correlation_backward", &correlation_backward, "Correlation backward",
        py::arg("grad_output"), py::arg("input1"), py::arg("input2"),
        py::arg("grad_input1"), py::arg("grad_input2"), py::arg("kH"),
        py::arg("kW"), py::arg("patchH"), py::arg("patchW"), py::arg("padH"),
        py::arg("padW"), py::arg("dilationH"), py::arg("dilationW"),
        py::arg("dilation_patchH"), py::arg("dilation_patchW"), py::arg("dH"),
        py::arg("dW"));
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
  m.def("points_in_boxes_cpu_forward", &points_in_boxes_cpu_forward,
        "points_in_boxes_cpu_forward", py::arg("boxes_tensor"),
        py::arg("pts_tensor"), py::arg("pts_indices_tensor"));
  m.def("points_in_boxes_part_forward", &points_in_boxes_part_forward,
        "points_in_boxes_part_forward", py::arg("boxes_tensor"),
        py::arg("pts_tensor"), py::arg("box_idx_of_points_tensor"));
  m.def("points_in_boxes_all_forward", &points_in_boxes_all_forward,
        "points_in_boxes_all_forward", py::arg("boxes_tensor"),
        py::arg("pts_tensor"), py::arg("box_idx_of_points_tensor"));
  m.def("roiaware_pool3d_forward", &roiaware_pool3d_forward,
        "roiaware_pool3d_forward", py::arg("rois"), py::arg("pts"),
        py::arg("pts_feature"), py::arg("argmax"), py::arg("pts_idx_of_voxels"),
        py::arg("pooled_features"), py::arg("pool_method"));
  m.def("roiaware_pool3d_backward", &roiaware_pool3d_backward,
        "roiaware_pool3d_backward", py::arg("pts_idx_of_voxels"),
        py::arg("argmax"), py::arg("grad_out"), py::arg("grad_in"),
        py::arg("pool_method"));
782
783
784
785
786
787
788
789
  m.def("rotated_feature_align_forward", &rotated_feature_align_forward,
        "Feature Refine forward (CUDA)", py::arg("features"),
        py::arg("best_bboxes"), py::arg("output"), py::arg("spatial_scale"),
        py::arg("points"));
  m.def("rotated_feature_align_backward", &rotated_feature_align_backward,
        "Feature Refine backward (CUDA)", py::arg("top_grad"),
        py::arg("best_bboxes"), py::arg("bottom_grad"),
        py::arg("spatial_scale"), py::arg("points"));
790
791
792
793
794
795
796
797
798
799
800
  m.def("riroi_align_rotated_forward", &riroi_align_rotated_forward,
        "riroi_align_rotated forward", py::arg("features"), py::arg("rois"),
        py::arg("output"), py::arg("pooled_height"), py::arg("pooled_width"),
        py::arg("spatial_scale"), py::arg("num_samples"),
        py::arg("num_orientations"), py::arg("clockwise"));
  m.def("riroi_align_rotated_backward", &riroi_align_rotated_backward,
        "riroi_align_rotated backward", py::arg("top_grad"), py::arg("rois"),
        py::arg("bottom_grad"), py::arg("pooled_height"),
        py::arg("pooled_width"), py::arg("spatial_scale"),
        py::arg("num_samples"), py::arg("num_orientations"),
        py::arg("clockwise"));
801
802
803
  m.def("points_in_polygons_forward", &points_in_polygons_forward,
        "points_in_polygons_forward", py::arg("points"), py::arg("polygons"),
        py::arg("output"));
804
805
  m.def("min_area_polygons", &min_area_polygons, "min_area_polygons",
        py::arg("pointsets"), py::arg("polygons"));
806
807
808
809
810
811
  m.def("active_rotated_filter_forward", &active_rotated_filter_forward,
        "active_rotated_filter_forward", py::arg("input"), py::arg("indices"),
        py::arg("output"));
  m.def("active_rotated_filter_backward", &active_rotated_filter_backward,
        "active_rotated_filter_backward", py::arg("grad_out"),
        py::arg("indices"), py::arg("grad_in"));
812
813
814
815
  m.def("convex_iou", &convex_iou, "convex_iou", py::arg("pointsets"),
        py::arg("polygons"), py::arg("ious"));
  m.def("convex_giou", &convex_giou, "convex_giou", py::arg("pointsets"),
        py::arg("polygons"), py::arg("output"));
816
817
818
819
  m.def("diff_iou_rotated_sort_vertices_forward",
        &diff_iou_rotated_sort_vertices_forward,
        "diff_iou_rotated_sort_vertices_forward", py::arg("vertices"),
        py::arg("mask"), py::arg("num_valid"));
820
}