Open3D (C++ API)  0.16.0
VoxelizeOpKernel.h
Go to the documentation of this file.
1// ----------------------------------------------------------------------------
2// - Open3D: www.open3d.org -
3// ----------------------------------------------------------------------------
4// The MIT License (MIT)
5//
6// Copyright (c) 2018-2021 www.open3d.org
7//
8// Permission is hereby granted, free of charge, to any person obtaining a copy
9// of this software and associated documentation files (the "Software"), to deal
10// in the Software without restriction, including without limitation the rights
11// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
12// copies of the Software, and to permit persons to whom the Software is
13// furnished to do so, subject to the following conditions:
14//
15// The above copyright notice and this permission notice shall be included in
16// all copies or substantial portions of the Software.
17//
18// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
19// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
20// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
21// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
22// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
23// FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS
24// IN THE SOFTWARE.
25// ----------------------------------------------------------------------------
26//
27#pragma once
28
30#include "torch/script.h"
31
32template <class T>
33void VoxelizeCPU(const torch::Tensor& points,
34 const torch::Tensor& row_splits,
35 const torch::Tensor& voxel_size,
36 const torch::Tensor& points_range_min,
37 const torch::Tensor& points_range_max,
38 const int64_t max_points_per_voxel,
39 const int64_t max_voxels,
40 torch::Tensor& voxel_coords,
41 torch::Tensor& voxel_point_indices,
42 torch::Tensor& voxel_point_row_splits,
43 torch::Tensor& voxel_batch_splits);
44
45#ifdef BUILD_CUDA_MODULE
46template <class T>
47void VoxelizeCUDA(const torch::Tensor& points,
48 const torch::Tensor& row_splits,
49 const torch::Tensor& voxel_size,
50 const torch::Tensor& points_range_min,
51 const torch::Tensor& points_range_max,
52 const int64_t max_points_per_voxel,
53 const int64_t max_voxels,
54 torch::Tensor& voxel_coords,
55 torch::Tensor& voxel_point_indices,
56 torch::Tensor& voxel_point_row_splits,
57 torch::Tensor& voxel_batch_splits);
58#endif
59
61public:
62 VoxelizeOutputAllocator(torch::DeviceType device_type, int device_idx)
63 : device_type(device_type), device_idx(device_idx) {}
64
65 void AllocVoxelCoords(int32_t** ptr, int64_t rows, int64_t cols) {
66 voxel_coords = torch::empty({rows, cols},
67 torch::dtype(ToTorchDtype<int32_t>())
68 .device(device_type, device_idx));
69 *ptr = voxel_coords.data_ptr<int32_t>();
70 }
71
72 void AllocVoxelPointIndices(int64_t** ptr, int64_t num) {
73 voxel_point_indices =
74 torch::empty({num}, torch::dtype(ToTorchDtype<int64_t>())
75 .device(device_type, device_idx));
76 *ptr = voxel_point_indices.data_ptr<int64_t>();
77 }
78
79 void AllocVoxelPointRowSplits(int64_t** ptr, int64_t num) {
80 voxel_point_row_splits =
81 torch::empty({num}, torch::dtype(ToTorchDtype<int64_t>())
82 .device(device_type, device_idx));
83 *ptr = voxel_point_row_splits.data_ptr<int64_t>();
84 }
85
86 void AllocVoxelBatchSplits(int64_t** ptr, int64_t num) {
87 voxel_batch_splits =
88 torch::empty({num}, torch::dtype(ToTorchDtype<int64_t>())
89 .device(device_type, device_idx));
90 *ptr = voxel_batch_splits.data_ptr<int64_t>();
91 }
92
93 const torch::Tensor& VoxelCoords() const { return voxel_coords; }
94 const torch::Tensor& VoxelPointIndices() const {
95 return voxel_point_indices;
96 }
97 const torch::Tensor& VoxelPointRowSplits() const {
98 return voxel_point_row_splits;
99 }
100 const torch::Tensor& VoxelBatchSplits() const { return voxel_batch_splits; }
101
102private:
103 torch::Tensor voxel_coords;
104 torch::Tensor voxel_point_indices;
105 torch::Tensor voxel_point_row_splits;
106 torch::Tensor voxel_batch_splits;
107 torch::DeviceType device_type;
108 int device_idx;
109};
TorchDtype_t ToTorchDtype< int64_t >()
Definition: TorchHelper.h:95
TorchDtype_t ToTorchDtype< int32_t >()
Definition: TorchHelper.h:91
Definition: VoxelizeOpKernel.h:60
const torch::Tensor & VoxelPointIndices() const
Definition: VoxelizeOpKernel.h:94
void AllocVoxelPointRowSplits(int64_t **ptr, int64_t num)
Definition: VoxelizeOpKernel.h:79
void AllocVoxelBatchSplits(int64_t **ptr, int64_t num)
Definition: VoxelizeOpKernel.h:86
void AllocVoxelCoords(int32_t **ptr, int64_t rows, int64_t cols)
Definition: VoxelizeOpKernel.h:65
const torch::Tensor & VoxelPointRowSplits() const
Definition: VoxelizeOpKernel.h:97
const torch::Tensor & VoxelBatchSplits() const
Definition: VoxelizeOpKernel.h:100
void AllocVoxelPointIndices(int64_t **ptr, int64_t num)
Definition: VoxelizeOpKernel.h:72
const torch::Tensor & VoxelCoords() const
Definition: VoxelizeOpKernel.h:93
VoxelizeOutputAllocator(torch::DeviceType device_type, int device_idx)
Definition: VoxelizeOpKernel.h:62
int points
Definition: FilePCD.cpp:73
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:414
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:37