Open3D (C++ API)  0.19.0
Loading...
Searching...
No Matches
VoxelizeOpKernel.h
Go to the documentation of this file.
1// ----------------------------------------------------------------------------
2// - Open3D: www.open3d.org -
3// ----------------------------------------------------------------------------
4// Copyright (c) 2018-2024 www.open3d.org
5// SPDX-License-Identifier: MIT
6// ----------------------------------------------------------------------------
7//
8#pragma once
9
11#include "torch/script.h"
12
13template <class T>
14void 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
27template <class T>
28void 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
42public:
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
83private:
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
const torch::Tensor & VoxelPointIndices() const
Definition VoxelizeOpKernel.h:75
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
const torch::Tensor & VoxelPointRowSplits() const
Definition VoxelizeOpKernel.h:78
const torch::Tensor & VoxelBatchSplits() const
Definition VoxelizeOpKernel.h:81
void AllocVoxelPointIndices(int64_t **ptr, int64_t num)
Definition VoxelizeOpKernel.h:53
const torch::Tensor & VoxelCoords() const
Definition VoxelizeOpKernel.h:74
VoxelizeOutputAllocator(torch::DeviceType device_type, int device_idx)
Definition VoxelizeOpKernel.h:43
int points
Definition FilePCD.cpp:54
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