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

4
5
6
typedef enum { SUM, MEAN, MAX } reduce_t;
#define DIVUP(m, n) ((m) / (n) + ((m) % (n) > 0))

zhangwenwei's avatar
zhangwenwei committed
7
8
namespace voxelization {

9
10
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
11
12
13
14
15
                      const std::vector<float> voxel_size,
                      const std::vector<float> coors_range,
                      const int max_points, const int max_voxels,
                      const int NDim = 3);

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

#ifdef WITH_CUDA
26
27
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
28
29
30
31
32
                      const std::vector<float> voxel_size,
                      const std::vector<float> coors_range,
                      const int max_points, const int max_voxels,
                      const int NDim = 3);

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

38
39
40
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
41

42
43
44
45
46
47
48
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
49
50
51
#endif

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

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

86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
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
101
#ifdef WITH_CUDA
102
    return dynamic_point_to_voxel_forward_gpu(feats, coors, convert_reduce_type(reduce_type));
zhangwenwei's avatar
zhangwenwei committed
103
#else
104
    TORCH_CHECK(false, "Not compiled with GPU support");
zhangwenwei's avatar
zhangwenwei committed
105
106
#endif
  }
107
108
  TORCH_CHECK(false, "do not support cpu yet");
  return std::vector<torch::Tensor>();
zhangwenwei's avatar
zhangwenwei committed
109
110
}

111
112
113
114
115
116
117
118
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
119
#ifdef WITH_CUDA
120
121
122
123
    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
124
#else
125
    TORCH_CHECK(false, "Not compiled with GPU support");
zhangwenwei's avatar
zhangwenwei committed
126
127
#endif
  }
128
  TORCH_CHECK(false, "do not support cpu yet");
zhangwenwei's avatar
zhangwenwei committed
129
130
131
}

}  // namespace voxelization