Skip to content

GitLab

Switch branch/tag
Find file Normal viewHistoryPermalink
voxelization.h 4.65 KB
Newer Older
zhangwenwei's avatar
Initial commit  
zhangwenwei committed
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
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 
#pragma once
#include <torch/extension.h>

namespace voxelization {

int hard_voxelize_cpu(const at::Tensor& points, at::Tensor& voxels,
                      at::Tensor& coors, at::Tensor& num_points_per_voxel,
                      const std::vector<float> voxel_size,
                      const std::vector<float> coors_range,
                      const int max_points, const int max_voxels,
                      const int NDim = 3);

void dynamic_voxelize_cpu(const at::Tensor& points, at::Tensor& coors,
                          const std::vector<float> voxel_size,
                          const std::vector<float> coors_range,
                          const int NDim = 3);

std::vector<at::Tensor> dynamic_point_to_voxel_cpu(
    const at::Tensor& points, const at::Tensor& voxel_mapping,
    const std::vector<float> voxel_size, const std::vector<float> coors_range);

#ifdef WITH_CUDA
int hard_voxelize_gpu(const at::Tensor& points, at::Tensor& voxels,
                      at::Tensor& coors, at::Tensor& num_points_per_voxel,
                      const std::vector<float> voxel_size,
                      const std::vector<float> coors_range,
                      const int max_points, const int max_voxels,
                      const int NDim = 3);

void dynamic_voxelize_gpu(const at::Tensor& points, at::Tensor& coors,
                          const std::vector<float> voxel_size,
                          const std::vector<float> coors_range,
                          const int NDim = 3);

std::vector<at::Tensor> dynamic_point_to_voxel_forward_gpu(
    const at::Tensor& points, const at::Tensor& voxel_mapping,
    const std::vector<float> voxel_size, const std::vector<float> coors_range);

void dynamic_point_to_voxel_backward_gpu(at::Tensor& grad_input_points,
                                         const at::Tensor& grad_output_voxels,
                                         const at::Tensor& point_to_voxelidx,
                                         const at::Tensor& coor_to_voxelidx);
#endif

// Interface for Python
inline int hard_voxelize(const at::Tensor& points, at::Tensor& voxels,
                         at::Tensor& coors, at::Tensor& num_points_per_voxel,
                         const std::vector<float> voxel_size,
                         const std::vector<float> coors_range,
                         const int max_points, const int max_voxels,
                         const int NDim = 3) {
zhangwenwei's avatar
Reformat cpp code to supress warning  
zhangwenwei committed
52 
  if (points.device().is_cuda()) {
zhangwenwei's avatar
Initial commit  
zhangwenwei committed
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69 
#ifdef WITH_CUDA
    return hard_voxelize_gpu(points, voxels, coors, num_points_per_voxel,
                             voxel_size, coors_range, max_points, max_voxels,
                             NDim);
#else
    AT_ERROR("Not compiled with GPU support");
#endif
  }
  return hard_voxelize_cpu(points, voxels, coors, num_points_per_voxel,
                           voxel_size, coors_range, max_points, max_voxels,
                           NDim);
}

inline void dynamic_voxelize(const at::Tensor& points, at::Tensor& coors,
                             const std::vector<float> voxel_size,
                             const std::vector<float> coors_range,
                             const int NDim = 3) {
zhangwenwei's avatar
Reformat cpp code to supress warning  
zhangwenwei committed
70 
  if (points.device().is_cuda()) {
zhangwenwei's avatar
Initial commit  
zhangwenwei committed
71
72
73
74
75
76
77
78
79
80
81
82 
#ifdef WITH_CUDA
    return dynamic_voxelize_gpu(points, coors, voxel_size, coors_range, NDim);
#else
    AT_ERROR("Not compiled with GPU support");
#endif
  }
  return dynamic_voxelize_cpu(points, coors, voxel_size, coors_range, NDim);
}

inline std::vector<torch::Tensor> dynamic_point_to_voxel_forward(
    const at::Tensor& points, const at::Tensor& voxel_mapping,
    const std::vector<float> voxel_size, const std::vector<float> coors_range) {
zhangwenwei's avatar
Reformat cpp code to supress warning  
zhangwenwei committed
83 
  if (points.device().is_cuda()) {
zhangwenwei's avatar
Initial commit  
zhangwenwei committed
84
85
86
87
88
89
90
91
92
93
94
95
96
97 
#ifdef WITH_CUDA
    return dynamic_point_to_voxel_forward_gpu(points, voxel_mapping, voxel_size,
                                              coors_range);
#else
    AT_ERROR("Not compiled with GPU support");
#endif
  }
  return dynamic_point_to_voxel_cpu(points, voxel_mapping, voxel_size,
                                    coors_range);
}

inline void dynamic_point_to_voxel_backward(
    at::Tensor& grad_input_points, const at::Tensor& grad_output_voxels,
    const at::Tensor& point_to_voxelidx, const at::Tensor& coor_to_voxelidx) {
zhangwenwei's avatar
Reformat cpp code to supress warning  
zhangwenwei committed
98 
  if (grad_input_points.device().is_cuda()) {
zhangwenwei's avatar
Initial commit  
zhangwenwei committed
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113 
#ifdef WITH_CUDA
    return dynamic_point_to_voxel_backward_gpu(
        grad_input_points, grad_output_voxels, point_to_voxelidx,
        coor_to_voxelidx);
#else
    AT_ERROR("Not compiled with GPU support");
#endif
  }
  // return dynamic_point_to_voxel_cpu(points,
  //                                  voxel_mapping,
  //                                  voxel_size,
  //                                  coors_range);
}

}  // namespace voxelization