voxelization.h 5.54 KB
Newer Older
zhangwenwei's avatar
zhangwenwei committed
1
2
3
#pragma once
#include <torch/extension.h>

4
typedef enum { SUM = 0, MEAN = 1, MAX = 2 } reduce_t;
5

zhangwenwei's avatar
zhangwenwei committed
6
7
namespace voxelization {

8
9
int hard_voxelize_cpu(const at::Tensor &points, at::Tensor &voxels,
                      at::Tensor &coors, at::Tensor &num_points_per_voxel,
zhangwenwei's avatar
zhangwenwei committed
10
11
12
13
14
                      const std::vector<float> voxel_size,
                      const std::vector<float> coors_range,
                      const int max_points, const int max_voxels,
                      const int NDim = 3);

15
void dynamic_voxelize_cpu(const at::Tensor &points, at::Tensor &coors,
zhangwenwei's avatar
zhangwenwei committed
16
17
18
19
20
                          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(
21
    const at::Tensor &points, const at::Tensor &voxel_mapping,
zhangwenwei's avatar
zhangwenwei committed
22
23
24
    const std::vector<float> voxel_size, const std::vector<float> coors_range);

#ifdef WITH_CUDA
25
26
int hard_voxelize_gpu(const at::Tensor &points, at::Tensor &voxels,
                      at::Tensor &coors, at::Tensor &num_points_per_voxel,
zhangwenwei's avatar
zhangwenwei committed
27
28
29
30
31
                      const std::vector<float> voxel_size,
                      const std::vector<float> coors_range,
                      const int max_points, const int max_voxels,
                      const int NDim = 3);

32
void dynamic_voxelize_gpu(const at::Tensor &points, at::Tensor &coors,
zhangwenwei's avatar
zhangwenwei committed
33
34
35
36
                          const std::vector<float> voxel_size,
                          const std::vector<float> coors_range,
                          const int NDim = 3);

37
38
39
std::vector<torch::Tensor> dynamic_point_to_voxel_forward_gpu(const torch::Tensor &feats,
                                                              const torch::Tensor &coors,
                                                              const reduce_t reduce_type);
zhangwenwei's avatar
zhangwenwei committed
40

41
42
43
44
45
46
47
void dynamic_point_to_voxel_backward_gpu(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 reduce_t reduce_type);
zhangwenwei's avatar
zhangwenwei committed
48
49
50
#endif

// Interface for Python
51
52
inline int hard_voxelize(const at::Tensor &points, at::Tensor &voxels,
                         at::Tensor &coors, at::Tensor &num_points_per_voxel,
zhangwenwei's avatar
zhangwenwei committed
53
54
55
56
                         const std::vector<float> voxel_size,
                         const std::vector<float> coors_range,
                         const int max_points, const int max_voxels,
                         const int NDim = 3) {
57
  if (points.device().is_cuda()) {
zhangwenwei's avatar
zhangwenwei committed
58
59
60
61
62
63
64
65
66
67
68
69
70
#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);
}

71
inline void dynamic_voxelize(const at::Tensor &points, at::Tensor &coors,
zhangwenwei's avatar
zhangwenwei committed
72
73
74
                             const std::vector<float> voxel_size,
                             const std::vector<float> coors_range,
                             const int NDim = 3) {
75
  if (points.device().is_cuda()) {
zhangwenwei's avatar
zhangwenwei committed
76
77
78
79
80
81
82
83
84
#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);
}

85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
inline reduce_t convert_reduce_type(const std::string &reduce_type) {
  if (reduce_type == "max")
    return reduce_t::MAX;
  else if (reduce_type == "sum")
    return reduce_t::SUM;
  else if (reduce_type == "mean")
    return reduce_t::MEAN;
  else TORCH_CHECK(false, "do not support reduce type " + reduce_type)
  return reduce_t::SUM;
}

inline std::vector<torch::Tensor> dynamic_point_to_voxel_forward(const torch::Tensor &feats,
                                                                 const torch::Tensor &coors,
                                                                 const std::string &reduce_type) {
  if (feats.device().is_cuda()) {
zhangwenwei's avatar
zhangwenwei committed
100
#ifdef WITH_CUDA
101
    return dynamic_point_to_voxel_forward_gpu(feats, coors, convert_reduce_type(reduce_type));
zhangwenwei's avatar
zhangwenwei committed
102
#else
103
    TORCH_CHECK(false, "Not compiled with GPU support");
zhangwenwei's avatar
zhangwenwei committed
104
105
#endif
  }
106
107
  TORCH_CHECK(false, "do not support cpu yet");
  return std::vector<torch::Tensor>();
zhangwenwei's avatar
zhangwenwei committed
108
109
}

110
111
112
113
114
115
116
117
inline 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) {
  if (grad_feats.device().is_cuda()) {
zhangwenwei's avatar
zhangwenwei committed
118
#ifdef WITH_CUDA
119
120
121
122
    dynamic_point_to_voxel_backward_gpu(
        grad_feats, grad_reduced_feats, feats, reduced_feats, coors_idx, reduce_count,
        convert_reduce_type(reduce_type));
    return;
zhangwenwei's avatar
zhangwenwei committed
123
#else
124
    TORCH_CHECK(false, "Not compiled with GPU support");
zhangwenwei's avatar
zhangwenwei committed
125
126
#endif
  }
127
  TORCH_CHECK(false, "do not support cpu yet");
zhangwenwei's avatar
zhangwenwei committed
128
129
130
}

}  // namespace voxelization