Open3D (C++ API)  0.18.0+252c867
VoxelizeOpKernel.h
Go to the documentation of this file.
1 // ----------------------------------------------------------------------------
2 // - Open3D: www.open3d.org -
3 // ----------------------------------------------------------------------------
4 // Copyright (c) 2018-2023 www.open3d.org
5 // SPDX-License-Identifier: MIT
6 // ----------------------------------------------------------------------------
7 //
8 #pragma once
9 
11 #include "torch/script.h"
12 
13 template <class T>
14 void VoxelizeCPU(const torch::Tensor& points,
15  const torch::Tensor& row_splits,
16  const torch::Tensor& voxel_size,
17  const torch::Tensor& points_range_min,
18  const torch::Tensor& points_range_max,
19  const int64_t max_points_per_voxel,
20  const int64_t max_voxels,
21  torch::Tensor& voxel_coords,
22  torch::Tensor& voxel_point_indices,
23  torch::Tensor& voxel_point_row_splits,
24  torch::Tensor& voxel_batch_splits);
25 
26 #ifdef BUILD_CUDA_MODULE
27 template <class T>
28 void VoxelizeCUDA(const torch::Tensor& points,
29  const torch::Tensor& row_splits,
30  const torch::Tensor& voxel_size,
31  const torch::Tensor& points_range_min,
32  const torch::Tensor& points_range_max,
33  const int64_t max_points_per_voxel,
34  const int64_t max_voxels,
35  torch::Tensor& voxel_coords,
36  torch::Tensor& voxel_point_indices,
37  torch::Tensor& voxel_point_row_splits,
38  torch::Tensor& voxel_batch_splits);
39 #endif
40 
42 public:
43  VoxelizeOutputAllocator(torch::DeviceType device_type, int device_idx)
44  : device_type(device_type), device_idx(device_idx) {}
45 
46  void AllocVoxelCoords(int32_t** ptr, int64_t rows, int64_t cols) {
47  voxel_coords = torch::empty({rows, cols},
48  torch::dtype(ToTorchDtype<int32_t>())
49  .device(device_type, device_idx));
50  *ptr = voxel_coords.data_ptr<int32_t>();
51  }
52 
53  void AllocVoxelPointIndices(int64_t** ptr, int64_t num) {
54  voxel_point_indices =
55  torch::empty({num}, torch::dtype(ToTorchDtype<int64_t>())
56  .device(device_type, device_idx));
57  *ptr = voxel_point_indices.data_ptr<int64_t>();
58  }
59 
60  void AllocVoxelPointRowSplits(int64_t** ptr, int64_t num) {
61  voxel_point_row_splits =
62  torch::empty({num}, torch::dtype(ToTorchDtype<int64_t>())
63  .device(device_type, device_idx));
64  *ptr = voxel_point_row_splits.data_ptr<int64_t>();
65  }
66 
67  void AllocVoxelBatchSplits(int64_t** ptr, int64_t num) {
68  voxel_batch_splits =
69  torch::empty({num}, torch::dtype(ToTorchDtype<int64_t>())
70  .device(device_type, device_idx));
71  *ptr = voxel_batch_splits.data_ptr<int64_t>();
72  }
73 
74  const torch::Tensor& VoxelCoords() const { return voxel_coords; }
75  const torch::Tensor& VoxelPointIndices() const {
76  return voxel_point_indices;
77  }
78  const torch::Tensor& VoxelPointRowSplits() const {
79  return voxel_point_row_splits;
80  }
81  const torch::Tensor& VoxelBatchSplits() const { return voxel_batch_splits; }
82 
83 private:
84  torch::Tensor voxel_coords;
85  torch::Tensor voxel_point_indices;
86  torch::Tensor voxel_point_row_splits;
87  torch::Tensor voxel_batch_splits;
88  torch::DeviceType device_type;
89  int device_idx;
90 };
TorchDtype_t ToTorchDtype< int64_t >()
Definition: TorchHelper.h:76
TorchDtype_t ToTorchDtype< int32_t >()
Definition: TorchHelper.h:72
Definition: VoxelizeOpKernel.h:41
void AllocVoxelPointRowSplits(int64_t **ptr, int64_t num)
Definition: VoxelizeOpKernel.h:60
void AllocVoxelBatchSplits(int64_t **ptr, int64_t num)
Definition: VoxelizeOpKernel.h:67
void AllocVoxelCoords(int32_t **ptr, int64_t rows, int64_t cols)
Definition: VoxelizeOpKernel.h:46
void AllocVoxelPointIndices(int64_t **ptr, int64_t num)
Definition: VoxelizeOpKernel.h:53
const torch::Tensor & VoxelPointRowSplits() const
Definition: VoxelizeOpKernel.h:78
const torch::Tensor & VoxelPointIndices() const
Definition: VoxelizeOpKernel.h:75
const torch::Tensor & VoxelBatchSplits() const
Definition: VoxelizeOpKernel.h:81
VoxelizeOutputAllocator(torch::DeviceType device_type, int device_idx)
Definition: VoxelizeOpKernel.h:43
const torch::Tensor & VoxelCoords() const
Definition: VoxelizeOpKernel.h:74
int points
Definition: FilePCD.cpp:54
const char const char value recording_handle imu_sample recording_handle uint8_t size_t data_size k4a_record_configuration_t config target_format k4a_capture_t capture_handle k4a_imu_sample_t imu_sample playback_handle k4a_logging_message_cb_t void min_level device_handle k4a_imu_sample_t int32_t
Definition: K4aPlugin.cpp:395
void VoxelizeCPU(const torch::Tensor &points, const torch::Tensor &row_splits, const torch::Tensor &voxel_size, const torch::Tensor &points_range_min, const torch::Tensor &points_range_max, const int64_t max_points_per_voxel, const int64_t max_voxels, torch::Tensor &voxel_coords, torch::Tensor &voxel_point_indices, torch::Tensor &voxel_point_row_splits, torch::Tensor &voxel_batch_splits)
Definition: VoxelizeOpKernel.cpp:18