962 lines
39 KiB
Plaintext
962 lines
39 KiB
Plaintext
|
// ----------------------------------------------------------------------------
|
||
|
|
// - Open3D: www.open3d.org -
|
||
|
|
// ----------------------------------------------------------------------------
|
||
|
|
// Copyright (c) 2018-2023 www.open3d.org
|
||
|
|
// SPDX-License-Identifier: MIT
|
||
|
|
// ----------------------------------------------------------------------------
|
||
|
|
|
||
|
|
#pragma once
|
||
|
|
|
||
|
|
#include <thrust/sequence.h>
|
||
|
|
|
||
|
|
#include <cub/cub.cuh>
|
||
|
|
|
||
|
|
#include "open3d/ml/impl/misc/MemoryAllocation.h"
|
||
|
|
#include "open3d/utility/Helper.h"
|
||
|
|
#include "open3d/utility/MiniVec.h"
|
||
|
|
|
||
|
|
namespace open3d {
|
||
|
|
namespace ml {
|
||
|
|
namespace impl {
|
||
|
|
|
||
|
|
namespace {
|
||
|
|
|
||
|
|
using namespace open3d::utility;
|
||
|
|
|
||
|
|
template <class T, bool LARGE_ARRAY>
|
||
|
|
__global__ void IotaCUDAKernel(T* first, int64_t len, T value) {
|
||
|
|
int64_t linear_idx;
|
||
|
|
const int64_t x = blockDim.x * blockIdx.x + threadIdx.x;
|
||
|
|
if (LARGE_ARRAY) {
|
||
|
|
const int64_t y = blockDim.y * blockIdx.y + threadIdx.y;
|
||
|
|
const int64_t z = blockDim.z * blockIdx.z + threadIdx.z;
|
||
|
|
linear_idx = z * gridDim.x * blockDim.x * gridDim.y +
|
||
|
|
y * gridDim.x * blockDim.x + x;
|
||
|
|
} else {
|
||
|
|
linear_idx = x;
|
||
|
|
}
|
||
|
|
if (linear_idx < len) {
|
||
|
|
T* ptr = first + linear_idx;
|
||
|
|
value += linear_idx;
|
||
|
|
*ptr = value;
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
/// Iota function for CUDA
|
||
|
|
template <class T>
|
||
|
|
void IotaCUDA(const cudaStream_t& stream, T* first, T* last, T value) {
|
||
|
|
ptrdiff_t len = last - first;
|
||
|
|
if (len) {
|
||
|
|
const int BLOCKSIZE = 128;
|
||
|
|
dim3 block(BLOCKSIZE, 1, 1);
|
||
|
|
dim3 grid;
|
||
|
|
if (len > block.x * INT32_MAX) {
|
||
|
|
grid.y = std::ceil(std::cbrt(len));
|
||
|
|
grid.z = grid.y;
|
||
|
|
grid.x = DivUp(len, int64_t(grid.z) * grid.y * block.x);
|
||
|
|
IotaCUDAKernel<T, true>
|
||
|
|
<<<grid, block, 0, stream>>>(first, len, value);
|
||
|
|
} else {
|
||
|
|
grid = dim3(DivUp(len, block.x), 1, 1);
|
||
|
|
IotaCUDAKernel<T, false>
|
||
|
|
<<<grid, block, 0, stream>>>(first, len, value);
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
__global__ void ComputeBatchIdKernel(int64_t* hashes,
|
||
|
|
const int64_t num_voxels,
|
||
|
|
const int64_t batch_hash) {
|
||
|
|
int64_t linear_idx;
|
||
|
|
const int64_t x = blockDim.x * blockIdx.x + threadIdx.x;
|
||
|
|
const int64_t y = blockDim.y * blockIdx.y + threadIdx.y;
|
||
|
|
const int64_t z = blockDim.z * blockIdx.z + threadIdx.z;
|
||
|
|
linear_idx = z * gridDim.x * blockDim.x * gridDim.y +
|
||
|
|
y * gridDim.x * blockDim.x + x;
|
||
|
|
|
||
|
|
if (linear_idx >= num_voxels) return;
|
||
|
|
|
||
|
|
hashes[linear_idx] /= batch_hash;
|
||
|
|
}
|
||
|
|
|
||
|
|
/// This function computes batch_id from hash value.
|
||
|
|
///
|
||
|
|
/// \param hashes Input and output array.
|
||
|
|
/// \param num_voxels Number of valid voxels.
|
||
|
|
/// \param batch_hash The value used to hash batch dimension.
|
||
|
|
///
|
||
|
|
void ComputeBatchId(const cudaStream_t& stream,
|
||
|
|
int64_t* hashes,
|
||
|
|
const int64_t num_voxels,
|
||
|
|
const int64_t batch_hash) {
|
||
|
|
if (num_voxels) {
|
||
|
|
const int BLOCKSIZE = 128;
|
||
|
|
dim3 block(BLOCKSIZE, 1, 1);
|
||
|
|
dim3 grid;
|
||
|
|
grid.y = std::ceil(std::cbrt(num_voxels));
|
||
|
|
grid.z = grid.y;
|
||
|
|
grid.x = DivUp(num_voxels, int64_t(grid.z) * grid.y * block.x);
|
||
|
|
ComputeBatchIdKernel<<<grid, block, 0, stream>>>(hashes, num_voxels,
|
||
|
|
batch_hash);
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
__global__ void ComputeVoxelPerBatchKernel(int64_t* num_voxels_per_batch,
|
||
|
|
int64_t* unique_batches_count,
|
||
|
|
int64_t* unique_batches,
|
||
|
|
const int64_t num_batches) {
|
||
|
|
int64_t linear_idx;
|
||
|
|
const int64_t x = blockDim.x * blockIdx.x + threadIdx.x;
|
||
|
|
const int64_t y = blockDim.y * blockIdx.y + threadIdx.y;
|
||
|
|
const int64_t z = blockDim.z * blockIdx.z + threadIdx.z;
|
||
|
|
linear_idx = z * gridDim.x * blockDim.x * gridDim.y +
|
||
|
|
y * gridDim.x * blockDim.x + x;
|
||
|
|
|
||
|
|
if (linear_idx >= num_batches) return;
|
||
|
|
|
||
|
|
int64_t out_idx = unique_batches[linear_idx];
|
||
|
|
num_voxels_per_batch[out_idx] = unique_batches_count[linear_idx];
|
||
|
|
}
|
||
|
|
|
||
|
|
/// This function computes number of voxels per batch element.
|
||
|
|
///
|
||
|
|
/// \param num_voxels_per_batch The output array.
|
||
|
|
/// \param unique_batches_count Counts for unique batch_id.
|
||
|
|
/// \param unique_batches Unique batch_id.
|
||
|
|
/// \param num_batches Number of non empty batches (<= batch_size).
|
||
|
|
///
|
||
|
|
void ComputeVoxelPerBatch(const cudaStream_t& stream,
|
||
|
|
int64_t* num_voxels_per_batch,
|
||
|
|
int64_t* unique_batches_count,
|
||
|
|
int64_t* unique_batches,
|
||
|
|
const int64_t num_batches) {
|
||
|
|
if (num_batches) {
|
||
|
|
const int BLOCKSIZE = 128;
|
||
|
|
dim3 block(BLOCKSIZE, 1, 1);
|
||
|
|
dim3 grid;
|
||
|
|
grid.y = std::ceil(std::cbrt(num_batches));
|
||
|
|
grid.z = grid.y;
|
||
|
|
grid.x = DivUp(num_batches, int64_t(grid.z) * grid.y * block.x);
|
||
|
|
ComputeVoxelPerBatchKernel<<<grid, block, 0, stream>>>(
|
||
|
|
num_voxels_per_batch, unique_batches_count, unique_batches,
|
||
|
|
num_batches);
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
__global__ void ComputeIndicesBatchesKernel(int64_t* indices_batches,
|
||
|
|
const int64_t* row_splits,
|
||
|
|
const int64_t batch_size) {
|
||
|
|
int64_t linear_idx;
|
||
|
|
const int64_t x = blockDim.x * blockIdx.x + threadIdx.x;
|
||
|
|
const int64_t y = blockDim.y * blockIdx.y + threadIdx.y;
|
||
|
|
const int64_t z = blockDim.z * blockIdx.z + threadIdx.z;
|
||
|
|
linear_idx = z * gridDim.x * blockDim.x * gridDim.y +
|
||
|
|
y * gridDim.x * blockDim.x + x;
|
||
|
|
|
||
|
|
if (linear_idx >= batch_size) return;
|
||
|
|
|
||
|
|
for (int64_t i = row_splits[linear_idx]; i < row_splits[linear_idx + 1];
|
||
|
|
++i) {
|
||
|
|
indices_batches[i] = linear_idx;
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
/// This function computes mapping of index to batch_id.
|
||
|
|
///
|
||
|
|
/// \param indices_batches The output array.
|
||
|
|
/// \param row_splits The row_splits for defining batches.
|
||
|
|
/// \param batch_size The batch_size of given points.
|
||
|
|
///
|
||
|
|
void ComputeIndicesBatches(const cudaStream_t& stream,
|
||
|
|
int64_t* indices_batches,
|
||
|
|
const int64_t* row_splits,
|
||
|
|
const int64_t batch_size) {
|
||
|
|
if (batch_size) {
|
||
|
|
const int BLOCKSIZE = 128;
|
||
|
|
dim3 block(BLOCKSIZE, 1, 1);
|
||
|
|
dim3 grid;
|
||
|
|
grid.y = std::ceil(std::cbrt(batch_size));
|
||
|
|
grid.z = grid.y;
|
||
|
|
grid.x = DivUp(batch_size, int64_t(grid.z) * grid.y * block.x);
|
||
|
|
ComputeIndicesBatchesKernel<<<grid, block, 0, stream>>>(
|
||
|
|
indices_batches, row_splits, batch_size);
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
template <class T, int NDIM>
|
||
|
|
__global__ void ComputeHashKernel(
|
||
|
|
int64_t* __restrict__ hashes,
|
||
|
|
int64_t num_points,
|
||
|
|
const T* const __restrict__ points,
|
||
|
|
const int64_t batch_size,
|
||
|
|
const int64_t* row_splits,
|
||
|
|
const int64_t* indices_batches,
|
||
|
|
const open3d::utility::MiniVec<T, NDIM> points_range_min_vec,
|
||
|
|
const open3d::utility::MiniVec<T, NDIM> points_range_max_vec,
|
||
|
|
const open3d::utility::MiniVec<T, NDIM> inv_voxel_size,
|
||
|
|
const open3d::utility::MiniVec<int64_t, NDIM> strides,
|
||
|
|
const int64_t batch_hash,
|
||
|
|
const int64_t invalid_hash) {
|
||
|
|
typedef MiniVec<T, NDIM> Vec_t;
|
||
|
|
int64_t linear_idx;
|
||
|
|
const int64_t x = blockDim.x * blockIdx.x + threadIdx.x;
|
||
|
|
const int64_t y = blockDim.y * blockIdx.y + threadIdx.y;
|
||
|
|
const int64_t z = blockDim.z * blockIdx.z + threadIdx.z;
|
||
|
|
linear_idx = z * gridDim.x * blockDim.x * gridDim.y +
|
||
|
|
y * gridDim.x * blockDim.x + x;
|
||
|
|
|
||
|
|
if (linear_idx >= num_points) return;
|
||
|
|
|
||
|
|
Vec_t point(points + linear_idx * NDIM);
|
||
|
|
|
||
|
|
if ((point >= points_range_min_vec && point <= points_range_max_vec)
|
||
|
|
.all()) {
|
||
|
|
auto coords = ((point - points_range_min_vec) * inv_voxel_size)
|
||
|
|
.template cast<int64_t>();
|
||
|
|
int64_t h = coords.dot(strides);
|
||
|
|
h += indices_batches[linear_idx] * batch_hash; // add hash for batch_id
|
||
|
|
hashes[linear_idx] = h;
|
||
|
|
} else {
|
||
|
|
hashes[linear_idx] = invalid_hash; // max hash value used as invalid
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
/// This function computes the hash (linear index) for each point.
|
||
|
|
/// Points outside the range will get a specific hash value.
|
||
|
|
///
|
||
|
|
/// \tparam T The floating point type for the points
|
||
|
|
/// \tparam NDIM The number of dimensions, e.g., 3.
|
||
|
|
///
|
||
|
|
/// \param hashes The output vector with the hashes/linear indexes.
|
||
|
|
/// \param num_points The number of points.
|
||
|
|
/// \param points The array with the point coordinates. The shape is
|
||
|
|
/// [num_points,NDIM] and the storage order is row-major.
|
||
|
|
/// \param batch_size The batch size of points.
|
||
|
|
/// \param row_splits row_splits for defining batches.
|
||
|
|
/// \param indices_batches Mapping of index to batch_id.
|
||
|
|
/// \param points_range_min_vec The minimum range for a point to be valid.
|
||
|
|
/// \param points_range_max_vec The maximum range for a point to be valid.
|
||
|
|
/// \param inv_voxel_size The reciprocal of the voxel edge lengths in each
|
||
|
|
/// dimension
|
||
|
|
/// \param strides The strides for computing the linear index.
|
||
|
|
/// \param batch_hash The value for hashing batch dimension.
|
||
|
|
/// \param invalid_hash The value to use for points outside the range.
|
||
|
|
template <class T, int NDIM>
|
||
|
|
void ComputeHash(const cudaStream_t& stream,
|
||
|
|
int64_t* hashes,
|
||
|
|
int64_t num_points,
|
||
|
|
const T* const points,
|
||
|
|
const int64_t batch_size,
|
||
|
|
const int64_t* row_splits,
|
||
|
|
const int64_t* indices_batches,
|
||
|
|
const MiniVec<T, NDIM> points_range_min_vec,
|
||
|
|
const MiniVec<T, NDIM> points_range_max_vec,
|
||
|
|
const MiniVec<T, NDIM> inv_voxel_size,
|
||
|
|
const MiniVec<int64_t, NDIM> strides,
|
||
|
|
const int64_t batch_hash,
|
||
|
|
const int64_t invalid_hash) {
|
||
|
|
if (num_points) {
|
||
|
|
const int BLOCKSIZE = 128;
|
||
|
|
dim3 block(BLOCKSIZE, 1, 1);
|
||
|
|
dim3 grid;
|
||
|
|
grid.y = std::ceil(std::cbrt(num_points));
|
||
|
|
grid.z = grid.y;
|
||
|
|
grid.x = DivUp(num_points, int64_t(grid.z) * grid.y * block.x);
|
||
|
|
ComputeHashKernel<T, NDIM><<<grid, block, 0, stream>>>(
|
||
|
|
hashes, num_points, points, batch_size, row_splits,
|
||
|
|
indices_batches, points_range_min_vec, points_range_max_vec,
|
||
|
|
inv_voxel_size, strides, batch_hash, invalid_hash);
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
template <class T>
|
||
|
|
__global__ void LimitCountsKernel(T* counts, int64_t num, T limit) {
|
||
|
|
int64_t linear_idx;
|
||
|
|
const int64_t x = blockDim.x * blockIdx.x + threadIdx.x;
|
||
|
|
const int64_t y = blockDim.y * blockIdx.y + threadIdx.y;
|
||
|
|
const int64_t z = blockDim.z * blockIdx.z + threadIdx.z;
|
||
|
|
linear_idx = z * gridDim.x * blockDim.x * gridDim.y +
|
||
|
|
y * gridDim.x * blockDim.x + x;
|
||
|
|
|
||
|
|
if (linear_idx >= num) return;
|
||
|
|
|
||
|
|
if (counts[linear_idx] > limit) {
|
||
|
|
counts[linear_idx] = limit;
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
/// This function performs an element-wise minimum operation.
|
||
|
|
///
|
||
|
|
/// \param counts The input and output array.
|
||
|
|
/// \param num Number of input elements.
|
||
|
|
/// \param limit The second operator for the minimum operation.
|
||
|
|
template <class T>
|
||
|
|
void LimitCounts(const cudaStream_t& stream, T* counts, int64_t num, T limit) {
|
||
|
|
if (num) {
|
||
|
|
const int BLOCKSIZE = 128;
|
||
|
|
dim3 block(BLOCKSIZE, 1, 1);
|
||
|
|
dim3 grid;
|
||
|
|
grid.y = std::ceil(std::cbrt(num));
|
||
|
|
grid.z = grid.y;
|
||
|
|
grid.x = DivUp(num, int64_t(grid.z) * grid.y * block.x);
|
||
|
|
LimitCountsKernel<<<grid, block, 0, stream>>>(counts, num, limit);
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
__global__ void ComputeStartIdxKernel(
|
||
|
|
int64_t* start_idx,
|
||
|
|
int64_t* points_count,
|
||
|
|
const int64_t* num_voxels_prefix_sum,
|
||
|
|
const int64_t* unique_hashes_count_prefix_sum,
|
||
|
|
const int64_t* out_batch_splits,
|
||
|
|
const int64_t batch_size,
|
||
|
|
const int64_t max_voxels,
|
||
|
|
const int64_t max_points_per_voxel) {
|
||
|
|
int64_t linear_idx;
|
||
|
|
const int64_t x = blockDim.x * blockIdx.x + threadIdx.x;
|
||
|
|
const int64_t y = blockDim.y * blockIdx.y + threadIdx.y;
|
||
|
|
const int64_t z = blockDim.z * blockIdx.z + threadIdx.z;
|
||
|
|
linear_idx = z * gridDim.x * blockDim.x * gridDim.y +
|
||
|
|
y * gridDim.x * blockDim.x + x;
|
||
|
|
|
||
|
|
if (linear_idx >= batch_size) return;
|
||
|
|
|
||
|
|
int64_t voxel_idx;
|
||
|
|
if (0 == linear_idx) {
|
||
|
|
voxel_idx = 0;
|
||
|
|
} else {
|
||
|
|
voxel_idx = num_voxels_prefix_sum[linear_idx - 1];
|
||
|
|
}
|
||
|
|
|
||
|
|
int64_t begin_out = out_batch_splits[linear_idx];
|
||
|
|
int64_t end_out = out_batch_splits[linear_idx + 1];
|
||
|
|
|
||
|
|
for (int64_t out_idx = begin_out; out_idx < end_out;
|
||
|
|
out_idx++, voxel_idx++) {
|
||
|
|
if (voxel_idx == 0) {
|
||
|
|
start_idx[out_idx] = 0;
|
||
|
|
points_count[out_idx] = min(max_points_per_voxel,
|
||
|
|
unique_hashes_count_prefix_sum[0]);
|
||
|
|
} else {
|
||
|
|
start_idx[out_idx] = unique_hashes_count_prefix_sum[voxel_idx - 1];
|
||
|
|
points_count[out_idx] =
|
||
|
|
min(max_points_per_voxel,
|
||
|
|
unique_hashes_count_prefix_sum[voxel_idx] -
|
||
|
|
unique_hashes_count_prefix_sum[voxel_idx - 1]);
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
/// Computes the starting index of each voxel.
|
||
|
|
///
|
||
|
|
/// \param start_idx The output array for storing starting index.
|
||
|
|
/// \param points_count The output array for storing points count.
|
||
|
|
/// \param num_voxels_prefix_sum The Inclusive prefix sum which gives
|
||
|
|
/// the index of starting voxel for each batch.
|
||
|
|
/// \param unique_hashes_count_prefix_sum Inclusive prefix sum defining
|
||
|
|
/// where point indices for each voxel ends.
|
||
|
|
/// \param out_batch_splits Defines starting and ending voxels for
|
||
|
|
/// each batch element.
|
||
|
|
/// \param batch_size The batch size.
|
||
|
|
/// \param max_voxels Maximum voxels per batch.
|
||
|
|
/// \param max_points_per_voxel Maximum points per voxel.
|
||
|
|
///
|
||
|
|
void ComputeStartIdx(const cudaStream_t& stream,
|
||
|
|
int64_t* start_idx,
|
||
|
|
int64_t* points_count,
|
||
|
|
const int64_t* num_voxels_prefix_sum,
|
||
|
|
const int64_t* unique_hashes_count_prefix_sum,
|
||
|
|
const int64_t* out_batch_splits,
|
||
|
|
const int64_t batch_size,
|
||
|
|
const int64_t max_voxels,
|
||
|
|
const int64_t max_points_per_voxel) {
|
||
|
|
if (batch_size) {
|
||
|
|
const int BLOCKSIZE = 128;
|
||
|
|
dim3 block(BLOCKSIZE, 1, 1);
|
||
|
|
dim3 grid;
|
||
|
|
grid.y = std::ceil(std::cbrt(batch_size));
|
||
|
|
grid.z = grid.y;
|
||
|
|
grid.x = DivUp(batch_size, int64_t(grid.z) * grid.y * block.x);
|
||
|
|
ComputeStartIdxKernel<<<grid, block, 0, stream>>>(
|
||
|
|
start_idx, points_count, num_voxels_prefix_sum,
|
||
|
|
unique_hashes_count_prefix_sum, out_batch_splits, batch_size,
|
||
|
|
max_voxels, max_points_per_voxel);
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
template <class T, int NDIM>
|
||
|
|
__global__ void ComputeVoxelCoordsKernel(
|
||
|
|
int32_t* __restrict__ voxel_coords,
|
||
|
|
const T* const __restrict__ points,
|
||
|
|
const int64_t* const __restrict__ point_indices,
|
||
|
|
const int64_t* const __restrict__ prefix_sum,
|
||
|
|
const MiniVec<T, NDIM> points_range_min_vec,
|
||
|
|
const MiniVec<T, NDIM> inv_voxel_size,
|
||
|
|
int64_t num_voxels) {
|
||
|
|
typedef MiniVec<T, NDIM> Vec_t;
|
||
|
|
int64_t linear_idx;
|
||
|
|
const int64_t x = blockDim.x * blockIdx.x + threadIdx.x;
|
||
|
|
const int64_t y = blockDim.y * blockIdx.y + threadIdx.y;
|
||
|
|
const int64_t z = blockDim.z * blockIdx.z + threadIdx.z;
|
||
|
|
linear_idx = z * gridDim.x * blockDim.x * gridDim.y +
|
||
|
|
y * gridDim.x * blockDim.x + x;
|
||
|
|
|
||
|
|
if (linear_idx >= num_voxels) return;
|
||
|
|
|
||
|
|
int64_t point_idx = point_indices[prefix_sum[linear_idx]];
|
||
|
|
|
||
|
|
Vec_t point(points + point_idx * NDIM);
|
||
|
|
auto coords = ((point - points_range_min_vec) * inv_voxel_size)
|
||
|
|
.template cast<int32_t>();
|
||
|
|
|
||
|
|
for (int i = 0; i < NDIM; ++i) {
|
||
|
|
voxel_coords[linear_idx * NDIM + i] = coords[i];
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
/// Computes the coordinates for each voxel
|
||
|
|
///
|
||
|
|
/// \param voxel_coords The output array with shape [num_voxels, NDIM].
|
||
|
|
/// \param points The array with the point coordinates.
|
||
|
|
/// \param point_indices The array with the point indices for all voxels.
|
||
|
|
/// \param prefix_sum Inclusive prefix sum defining where the point indices
|
||
|
|
/// for each voxels end.
|
||
|
|
/// \param points_range_min The lower bound of the domain to be
|
||
|
|
/// voxelized.
|
||
|
|
/// \param points_range_max The upper bound of the domain to be
|
||
|
|
/// voxelized.
|
||
|
|
/// \param inv_voxel_size The reciprocal of the voxel edge lengths for each
|
||
|
|
/// dimension.
|
||
|
|
/// \param num_voxels The number of voxels.
|
||
|
|
template <class T, int NDIM>
|
||
|
|
void ComputeVoxelCoords(const cudaStream_t& stream,
|
||
|
|
int32_t* voxel_coords,
|
||
|
|
const T* const points,
|
||
|
|
const int64_t* const point_indices,
|
||
|
|
const int64_t* const prefix_sum,
|
||
|
|
const MiniVec<T, NDIM> points_range_min_vec,
|
||
|
|
const MiniVec<T, NDIM> inv_voxel_size,
|
||
|
|
int64_t num_voxels) {
|
||
|
|
if (num_voxels) {
|
||
|
|
const int BLOCKSIZE = 128;
|
||
|
|
dim3 block(BLOCKSIZE, 1, 1);
|
||
|
|
dim3 grid;
|
||
|
|
grid.y = std::ceil(std::cbrt(num_voxels));
|
||
|
|
grid.z = grid.y;
|
||
|
|
grid.x = DivUp(num_voxels, int64_t(grid.z) * grid.y * block.x);
|
||
|
|
ComputeVoxelCoordsKernel<<<grid, block, 0, stream>>>(
|
||
|
|
voxel_coords, points, point_indices, prefix_sum,
|
||
|
|
points_range_min_vec, inv_voxel_size, num_voxels);
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
__global__ void CopyPointIndicesKernel(
|
||
|
|
int64_t* __restrict__ out,
|
||
|
|
const int64_t* const __restrict__ point_indices,
|
||
|
|
const int64_t* const __restrict__ prefix_sum_in,
|
||
|
|
const int64_t* const __restrict__ prefix_sum_out,
|
||
|
|
const int64_t num_voxels) {
|
||
|
|
// TODO data coalescing can be optimized
|
||
|
|
int64_t linear_idx;
|
||
|
|
const int64_t x = blockDim.x * blockIdx.x + threadIdx.x;
|
||
|
|
const int64_t y = blockDim.y * blockIdx.y + threadIdx.y;
|
||
|
|
const int64_t z = blockDim.z * blockIdx.z + threadIdx.z;
|
||
|
|
linear_idx = z * gridDim.x * blockDim.x * gridDim.y +
|
||
|
|
y * gridDim.x * blockDim.x + x;
|
||
|
|
|
||
|
|
if (linear_idx >= num_voxels) return;
|
||
|
|
|
||
|
|
int64_t begin_out;
|
||
|
|
if (0 == linear_idx) {
|
||
|
|
begin_out = 0;
|
||
|
|
} else {
|
||
|
|
begin_out = prefix_sum_out[linear_idx - 1];
|
||
|
|
}
|
||
|
|
int64_t end_out = prefix_sum_out[linear_idx];
|
||
|
|
|
||
|
|
int64_t num_points = end_out - begin_out;
|
||
|
|
|
||
|
|
int64_t in_idx = prefix_sum_in[linear_idx];
|
||
|
|
|
||
|
|
for (int64_t out_idx = begin_out; out_idx < end_out; ++out_idx, ++in_idx) {
|
||
|
|
out[out_idx] = point_indices[in_idx];
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
/// Copies the point indices for each voxel to the output.
|
||
|
|
///
|
||
|
|
/// \param out The output array with the point indices for all voxels.
|
||
|
|
/// \param point_indices The array with the point indices for all voxels.
|
||
|
|
/// \param prefix_sum_in Inclusive prefix sum defining where the point
|
||
|
|
/// indices for each voxels end.
|
||
|
|
/// \param prefix_sum_out Inclusive prefix sum defining where the point
|
||
|
|
/// indices for each voxels end.
|
||
|
|
/// \param num_voxels The number of voxels.
|
||
|
|
///
|
||
|
|
void CopyPointIndices(const cudaStream_t& stream,
|
||
|
|
int64_t* out,
|
||
|
|
const int64_t* const point_indices,
|
||
|
|
const int64_t* const prefix_sum_in,
|
||
|
|
const int64_t* const prefix_sum_out,
|
||
|
|
const int64_t num_voxels) {
|
||
|
|
if (num_voxels) {
|
||
|
|
const int BLOCKSIZE = 128;
|
||
|
|
dim3 block(BLOCKSIZE, 1, 1);
|
||
|
|
dim3 grid;
|
||
|
|
grid.y = std::ceil(std::cbrt(num_voxels));
|
||
|
|
grid.z = grid.y;
|
||
|
|
grid.x = DivUp(num_voxels, int64_t(grid.z) * grid.y * block.x);
|
||
|
|
CopyPointIndicesKernel<<<grid, block, 0, stream>>>(
|
||
|
|
out, point_indices, prefix_sum_in, prefix_sum_out, num_voxels);
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
} // namespace
|
||
|
|
|
||
|
|
/// This function voxelizes a point cloud.
|
||
|
|
/// The function returns the integer coordinates of the voxels that
|
||
|
|
/// contain points and a compact list of the indices that associate the
|
||
|
|
/// voxels to the points.
|
||
|
|
///
|
||
|
|
/// All pointer arguments point to device memory unless stated
|
||
|
|
/// otherwise.
|
||
|
|
///
|
||
|
|
/// \tparam T Floating-point data type for the point positions.
|
||
|
|
///
|
||
|
|
/// \tparam NDIM The number of dimensions of the points.
|
||
|
|
///
|
||
|
|
/// \tparam OUTPUT_ALLOCATOR Type of the output_allocator. See
|
||
|
|
/// \p output_allocator for more information.
|
||
|
|
///
|
||
|
|
/// \param stream The cuda stream for all kernel launches.
|
||
|
|
///
|
||
|
|
/// \param temp Pointer to temporary memory. If nullptr then the required
|
||
|
|
/// size of temporary memory will be written to \p temp_size and no
|
||
|
|
/// work is done.
|
||
|
|
///
|
||
|
|
/// \param temp_size The size of the temporary memory in bytes. This is
|
||
|
|
/// used as an output if temp is nullptr
|
||
|
|
///
|
||
|
|
/// \param texture_alignment The texture alignment in bytes. This is used
|
||
|
|
/// for allocating segments within the temporary memory.
|
||
|
|
///
|
||
|
|
/// \param num_points The number of points.
|
||
|
|
///
|
||
|
|
/// \param points Array with the point positions. The shape is
|
||
|
|
/// [num_points,NDIM].
|
||
|
|
///
|
||
|
|
/// \param batch_size The batch size of points.
|
||
|
|
///
|
||
|
|
/// \param row_splits row_splits for defining batches.
|
||
|
|
///
|
||
|
|
/// \param voxel_size The edge lengths of the voxel. The shape is
|
||
|
|
/// [NDIM]. This pointer points to host memory!
|
||
|
|
///
|
||
|
|
/// \param points_range_min The lower bound of the domain to be
|
||
|
|
/// voxelized.
|
||
|
|
/// The shape is [NDIM].
|
||
|
|
/// This pointer points to host memory!
|
||
|
|
///
|
||
|
|
/// \param points_range_max The upper bound of the domain to be
|
||
|
|
/// voxelized.
|
||
|
|
/// The shape is [NDIM].
|
||
|
|
/// This pointer points to host memory!
|
||
|
|
///
|
||
|
|
/// \param max_points_per_voxel This parameter limits the number of
|
||
|
|
/// points
|
||
|
|
/// that are recorderd for each voxel.
|
||
|
|
///
|
||
|
|
/// \param max_voxels This parameter limits the number of voxels that
|
||
|
|
/// will be generated.
|
||
|
|
///
|
||
|
|
/// \param output_allocator An object that implements functions for
|
||
|
|
/// allocating the output arrays. The object must implement
|
||
|
|
/// functions AllocVoxelCoords(int32_t** ptr, int64_t rows,
|
||
|
|
/// int64_t cols), AllocVoxelPointIndices(int64_t** ptr, int64_t
|
||
|
|
/// size), AllocVoxelPointRowSplits(int64_t** ptr, int64_t
|
||
|
|
/// size) and AllocVoxelBatchSplits(int64_t** ptr, int64_t size).
|
||
|
|
/// All functions should allocate memory and return a pointer
|
||
|
|
/// to that memory in ptr. The arguments size, rows, and cols
|
||
|
|
/// define the size of the array as the number of elements.
|
||
|
|
/// All functions must accept zero size arguments. In this case
|
||
|
|
/// ptr does not need to be set.
|
||
|
|
///
|
||
|
|
template <class T, int NDIM, class OUTPUT_ALLOCATOR>
|
||
|
|
void VoxelizeCUDA(const cudaStream_t& stream,
|
||
|
|
void* temp,
|
||
|
|
size_t& temp_size,
|
||
|
|
int texture_alignment,
|
||
|
|
size_t num_points,
|
||
|
|
const T* const points,
|
||
|
|
const size_t batch_size,
|
||
|
|
const int64_t* const row_splits,
|
||
|
|
const T* const voxel_size,
|
||
|
|
const T* const points_range_min,
|
||
|
|
const T* const points_range_max,
|
||
|
|
const int64_t max_points_per_voxel,
|
||
|
|
const int64_t max_voxels,
|
||
|
|
OUTPUT_ALLOCATOR& output_allocator) {
|
||
|
|
using namespace open3d::utility;
|
||
|
|
typedef MiniVec<T, NDIM> Vec_t;
|
||
|
|
|
||
|
|
const bool get_temp_size = !temp;
|
||
|
|
|
||
|
|
if (get_temp_size) {
|
||
|
|
temp = (char*)1; // worst case pointer alignment
|
||
|
|
temp_size = std::numeric_limits<int64_t>::max();
|
||
|
|
}
|
||
|
|
|
||
|
|
MemoryAllocation mem_temp(temp, temp_size, texture_alignment);
|
||
|
|
|
||
|
|
const Vec_t inv_voxel_size = T(1) / Vec_t(voxel_size);
|
||
|
|
const Vec_t points_range_min_vec(points_range_min);
|
||
|
|
const Vec_t points_range_max_vec(points_range_max);
|
||
|
|
MiniVec<int32_t, NDIM> extents =
|
||
|
|
ceil((points_range_max_vec - points_range_min_vec) * inv_voxel_size)
|
||
|
|
.template cast<int32_t>();
|
||
|
|
MiniVec<int64_t, NDIM> strides;
|
||
|
|
for (int i = 0; i < NDIM; ++i) {
|
||
|
|
strides[i] = 1;
|
||
|
|
for (int j = 0; j < i; ++j) {
|
||
|
|
strides[i] *= extents[j];
|
||
|
|
}
|
||
|
|
}
|
||
|
|
const int64_t batch_hash = strides[NDIM - 1] * extents[NDIM - 1];
|
||
|
|
const int64_t invalid_hash = batch_hash * batch_size;
|
||
|
|
|
||
|
|
/// store batch_id for each point
|
||
|
|
std::pair<int64_t*, size_t> indices_batches =
|
||
|
|
mem_temp.Alloc<int64_t>(num_points);
|
||
|
|
if (!get_temp_size) {
|
||
|
|
ComputeIndicesBatches(stream, indices_batches.first, row_splits,
|
||
|
|
batch_size);
|
||
|
|
}
|
||
|
|
|
||
|
|
// use double buffers for the sorting
|
||
|
|
std::pair<int64_t*, size_t> point_indices =
|
||
|
|
mem_temp.Alloc<int64_t>(num_points);
|
||
|
|
std::pair<int64_t*, size_t> point_indices_alt =
|
||
|
|
mem_temp.Alloc<int64_t>(num_points);
|
||
|
|
std::pair<int64_t*, size_t> hashes = mem_temp.Alloc<int64_t>(num_points);
|
||
|
|
std::pair<int64_t*, size_t> hashes_alt =
|
||
|
|
mem_temp.Alloc<int64_t>(num_points);
|
||
|
|
|
||
|
|
cub::DoubleBuffer<int64_t> point_indices_dbuf(point_indices.first,
|
||
|
|
point_indices_alt.first);
|
||
|
|
cub::DoubleBuffer<int64_t> hashes_dbuf(hashes.first, hashes_alt.first);
|
||
|
|
|
||
|
|
if (!get_temp_size) {
|
||
|
|
IotaCUDA(stream, point_indices.first,
|
||
|
|
point_indices.first + point_indices.second, int64_t(0));
|
||
|
|
ComputeHash(stream, hashes.first, num_points, points, batch_size,
|
||
|
|
row_splits, indices_batches.first, points_range_min_vec,
|
||
|
|
points_range_max_vec, inv_voxel_size, strides, batch_hash,
|
||
|
|
invalid_hash);
|
||
|
|
}
|
||
|
|
|
||
|
|
{
|
||
|
|
// TODO compute end_bit for radix sort
|
||
|
|
std::pair<void*, size_t> sort_pairs_temp(nullptr, 0);
|
||
|
|
cub::DeviceRadixSort::SortPairs(
|
||
|
|
sort_pairs_temp.first, sort_pairs_temp.second, hashes_dbuf,
|
||
|
|
point_indices_dbuf, num_points, 0, sizeof(int64_t) * 8, stream);
|
||
|
|
sort_pairs_temp = mem_temp.Alloc(sort_pairs_temp.second);
|
||
|
|
if (!get_temp_size) {
|
||
|
|
cub::DeviceRadixSort::SortPairs(sort_pairs_temp.first,
|
||
|
|
sort_pairs_temp.second, hashes_dbuf,
|
||
|
|
point_indices_dbuf, num_points, 0,
|
||
|
|
sizeof(int64_t) * 8, stream);
|
||
|
|
}
|
||
|
|
mem_temp.Free(sort_pairs_temp);
|
||
|
|
}
|
||
|
|
|
||
|
|
// reuse the alternate buffers
|
||
|
|
std::pair<int64_t*, size_t> unique_hashes(hashes_dbuf.Alternate(),
|
||
|
|
hashes.second);
|
||
|
|
std::pair<int64_t*, size_t> unique_hashes_count(
|
||
|
|
point_indices_dbuf.Alternate(), point_indices.second);
|
||
|
|
|
||
|
|
// encode unique hashes(voxels) and their counts(points per voxel)
|
||
|
|
int64_t num_voxels = 0;
|
||
|
|
int64_t last_hash = 0; // 0 is a valid hash value
|
||
|
|
{
|
||
|
|
std::pair<void*, size_t> encode_temp(nullptr, 0);
|
||
|
|
std::pair<int64_t*, size_t> num_voxels_mem = mem_temp.Alloc<int64_t>(1);
|
||
|
|
|
||
|
|
cub::DeviceRunLengthEncode::Encode(
|
||
|
|
encode_temp.first, encode_temp.second, hashes_dbuf.Current(),
|
||
|
|
unique_hashes.first, unique_hashes_count.first,
|
||
|
|
num_voxels_mem.first, num_points, stream);
|
||
|
|
|
||
|
|
encode_temp = mem_temp.Alloc(encode_temp.second);
|
||
|
|
if (!get_temp_size) {
|
||
|
|
cub::DeviceRunLengthEncode::Encode(
|
||
|
|
encode_temp.first, encode_temp.second,
|
||
|
|
hashes_dbuf.Current(), unique_hashes.first,
|
||
|
|
unique_hashes_count.first, num_voxels_mem.first, num_points,
|
||
|
|
stream);
|
||
|
|
|
||
|
|
// get the number of voxels
|
||
|
|
cudaMemcpyAsync(&num_voxels, num_voxels_mem.first, sizeof(int64_t),
|
||
|
|
cudaMemcpyDeviceToHost, stream);
|
||
|
|
// get the last hash value
|
||
|
|
cudaMemcpyAsync(&last_hash,
|
||
|
|
hashes_dbuf.Current() + hashes.second - 1,
|
||
|
|
sizeof(int64_t), cudaMemcpyDeviceToHost, stream);
|
||
|
|
// wait for the async copies
|
||
|
|
while (cudaErrorNotReady == cudaStreamQuery(stream)) { /*empty*/
|
||
|
|
}
|
||
|
|
}
|
||
|
|
mem_temp.Free(encode_temp);
|
||
|
|
}
|
||
|
|
if (invalid_hash == last_hash) {
|
||
|
|
// the last hash is invalid we have one voxel less
|
||
|
|
--num_voxels;
|
||
|
|
}
|
||
|
|
|
||
|
|
// reuse the hashes buffer
|
||
|
|
std::pair<int64_t*, size_t> unique_hashes_count_prefix_sum(
|
||
|
|
hashes_dbuf.Current(), hashes.second);
|
||
|
|
|
||
|
|
// compute the prefix sum for unique_hashes_count
|
||
|
|
// gives starting index of each voxel
|
||
|
|
{
|
||
|
|
std::pair<void*, size_t> inclusive_scan_temp(nullptr, 0);
|
||
|
|
|
||
|
|
cub::DeviceScan::InclusiveSum(
|
||
|
|
inclusive_scan_temp.first, inclusive_scan_temp.second,
|
||
|
|
unique_hashes_count.first, unique_hashes_count_prefix_sum.first,
|
||
|
|
unique_hashes_count.second, stream);
|
||
|
|
|
||
|
|
inclusive_scan_temp = mem_temp.Alloc(inclusive_scan_temp.second);
|
||
|
|
if (!get_temp_size) {
|
||
|
|
// We only need the prefix sum for the first num_voxels.
|
||
|
|
cub::DeviceScan::InclusiveSum(
|
||
|
|
inclusive_scan_temp.first, inclusive_scan_temp.second,
|
||
|
|
unique_hashes_count.first,
|
||
|
|
unique_hashes_count_prefix_sum.first, num_voxels, stream);
|
||
|
|
}
|
||
|
|
mem_temp.Free(inclusive_scan_temp);
|
||
|
|
}
|
||
|
|
|
||
|
|
// Limit the number of output points to max_points_per_voxel by
|
||
|
|
// limiting the unique_hashes_count.
|
||
|
|
if (!get_temp_size) {
|
||
|
|
if (max_points_per_voxel < num_points) {
|
||
|
|
LimitCounts(stream, unique_hashes_count.first, num_voxels,
|
||
|
|
max_points_per_voxel);
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
// Convert unique_hashes to batch_id (divide with batch_hash)
|
||
|
|
int64_t* unique_hashes_batch_id = unique_hashes.first;
|
||
|
|
if (!get_temp_size) {
|
||
|
|
ComputeBatchId(stream, unique_hashes_batch_id, num_voxels, batch_hash);
|
||
|
|
}
|
||
|
|
|
||
|
|
std::pair<int64_t*, size_t> unique_batches =
|
||
|
|
mem_temp.Alloc<int64_t>(batch_size);
|
||
|
|
std::pair<int64_t*, size_t> unique_batches_count =
|
||
|
|
mem_temp.Alloc<int64_t>(batch_size);
|
||
|
|
int64_t num_batches = 0; // Store non empty batches
|
||
|
|
|
||
|
|
// Convert batch_id to counts (array of num_voxels per batch)
|
||
|
|
{
|
||
|
|
std::pair<void*, size_t> encode_temp(nullptr, 0);
|
||
|
|
std::pair<int64_t*, size_t> num_batches_mem =
|
||
|
|
mem_temp.Alloc<int64_t>(1);
|
||
|
|
|
||
|
|
cub::DeviceRunLengthEncode::Encode(
|
||
|
|
encode_temp.first, encode_temp.second, unique_hashes_batch_id,
|
||
|
|
unique_batches.first, unique_batches_count.first,
|
||
|
|
num_batches_mem.first, num_voxels, stream);
|
||
|
|
encode_temp = mem_temp.Alloc(encode_temp.second);
|
||
|
|
if (!get_temp_size) {
|
||
|
|
cub::DeviceRunLengthEncode::Encode(
|
||
|
|
encode_temp.first, encode_temp.second,
|
||
|
|
unique_hashes_batch_id, unique_batches.first,
|
||
|
|
unique_batches_count.first, num_batches_mem.first,
|
||
|
|
num_voxels, stream);
|
||
|
|
|
||
|
|
// get the number of non empty batches.
|
||
|
|
cudaMemcpyAsync(&num_batches, num_batches_mem.first,
|
||
|
|
sizeof(int64_t), cudaMemcpyDeviceToHost, stream);
|
||
|
|
// wait for the async copies
|
||
|
|
while (cudaErrorNotReady == cudaStreamQuery(stream)) { /*empty*/
|
||
|
|
}
|
||
|
|
}
|
||
|
|
mem_temp.Free(encode_temp);
|
||
|
|
}
|
||
|
|
|
||
|
|
// Insert count(0) for empty batches
|
||
|
|
std::pair<int64_t*, size_t> num_voxels_per_batch =
|
||
|
|
mem_temp.Alloc<int64_t>(batch_size);
|
||
|
|
if (!get_temp_size) {
|
||
|
|
cudaMemset(num_voxels_per_batch.first, 0, batch_size * sizeof(int64_t));
|
||
|
|
ComputeVoxelPerBatch(stream, num_voxels_per_batch.first,
|
||
|
|
unique_batches_count.first, unique_batches.first,
|
||
|
|
num_batches);
|
||
|
|
}
|
||
|
|
|
||
|
|
std::pair<int64_t*, size_t> num_voxels_prefix_sum(unique_batches.first,
|
||
|
|
batch_size);
|
||
|
|
|
||
|
|
// compute the prefix sum for number of voxels per batch
|
||
|
|
// gives starting voxel index for each batch
|
||
|
|
// used only when voxel count exceeds max_voxels
|
||
|
|
{
|
||
|
|
std::pair<void*, size_t> inclusive_scan_temp(nullptr, 0);
|
||
|
|
|
||
|
|
cub::DeviceScan::InclusiveSum(
|
||
|
|
inclusive_scan_temp.first, inclusive_scan_temp.second,
|
||
|
|
num_voxels_per_batch.first, num_voxels_prefix_sum.first,
|
||
|
|
num_voxels_per_batch.second, stream);
|
||
|
|
|
||
|
|
inclusive_scan_temp = mem_temp.Alloc(inclusive_scan_temp.second);
|
||
|
|
if (!get_temp_size) {
|
||
|
|
if (num_voxels > max_voxels) {
|
||
|
|
cub::DeviceScan::InclusiveSum(
|
||
|
|
inclusive_scan_temp.first, inclusive_scan_temp.second,
|
||
|
|
num_voxels_per_batch.first, num_voxels_prefix_sum.first,
|
||
|
|
num_voxels_per_batch.second, stream);
|
||
|
|
}
|
||
|
|
}
|
||
|
|
mem_temp.Free(inclusive_scan_temp);
|
||
|
|
}
|
||
|
|
|
||
|
|
// Limit the number of voxels per batch to max_voxels
|
||
|
|
if (!get_temp_size) {
|
||
|
|
if (num_voxels >= max_voxels)
|
||
|
|
LimitCounts(stream, num_voxels_per_batch.first, batch_size,
|
||
|
|
max_voxels);
|
||
|
|
}
|
||
|
|
|
||
|
|
// Prefix sum of limited counts to get batch splits.
|
||
|
|
int64_t* out_batch_splits = nullptr;
|
||
|
|
if (!get_temp_size) {
|
||
|
|
output_allocator.AllocVoxelBatchSplits(&out_batch_splits,
|
||
|
|
batch_size + 1);
|
||
|
|
cudaMemsetAsync(out_batch_splits, 0, sizeof(int64_t), stream);
|
||
|
|
}
|
||
|
|
|
||
|
|
// Prefix sum of counts to get batch splits
|
||
|
|
{
|
||
|
|
std::pair<void*, size_t> inclusive_scan_temp(nullptr, 0);
|
||
|
|
|
||
|
|
cub::DeviceScan::InclusiveSum(
|
||
|
|
inclusive_scan_temp.first, inclusive_scan_temp.second,
|
||
|
|
num_voxels_per_batch.first, out_batch_splits + 1,
|
||
|
|
num_voxels_per_batch.second, stream);
|
||
|
|
|
||
|
|
inclusive_scan_temp = mem_temp.Alloc(inclusive_scan_temp.second);
|
||
|
|
|
||
|
|
if (!get_temp_size) {
|
||
|
|
cub::DeviceScan::InclusiveSum(
|
||
|
|
inclusive_scan_temp.first, inclusive_scan_temp.second,
|
||
|
|
num_voxels_per_batch.first, out_batch_splits + 1,
|
||
|
|
batch_size, stream);
|
||
|
|
}
|
||
|
|
mem_temp.Free(inclusive_scan_temp);
|
||
|
|
}
|
||
|
|
|
||
|
|
// num_valid_voxels excludes voxels exceeding max_voxels
|
||
|
|
int64_t num_valid_voxels = num_points;
|
||
|
|
if (!get_temp_size) {
|
||
|
|
// get the number of valid voxels.
|
||
|
|
cudaMemcpyAsync(&num_valid_voxels, out_batch_splits + batch_size,
|
||
|
|
sizeof(int64_t), cudaMemcpyDeviceToHost, stream);
|
||
|
|
// wait for the async copies
|
||
|
|
while (cudaErrorNotReady == cudaStreamQuery(stream)) { /*empty*/
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
// start_idx stores starting index of each valid voxel.
|
||
|
|
// points_count stores number of valid points in respective voxel.
|
||
|
|
std::pair<int64_t*, size_t> start_idx(indices_batches.first,
|
||
|
|
num_valid_voxels);
|
||
|
|
std::pair<int64_t*, size_t> points_count =
|
||
|
|
mem_temp.Alloc<int64_t>(num_valid_voxels);
|
||
|
|
|
||
|
|
if (!get_temp_size) {
|
||
|
|
if (num_voxels <= max_voxels) {
|
||
|
|
// starting index and points_count will be same as
|
||
|
|
// unique_hashes_count_prefix_sum and unique_hashes_count when voxel
|
||
|
|
// count doesn't exceeds max_voxels
|
||
|
|
cudaMemsetAsync(start_idx.first, 0, sizeof(int64_t), stream);
|
||
|
|
cudaMemcpyAsync(start_idx.first + 1,
|
||
|
|
unique_hashes_count_prefix_sum.first,
|
||
|
|
(num_voxels - 1) * sizeof(int64_t),
|
||
|
|
cudaMemcpyDeviceToDevice, stream);
|
||
|
|
mem_temp.Free(points_count);
|
||
|
|
points_count.first = unique_hashes_count.first;
|
||
|
|
} else {
|
||
|
|
ComputeStartIdx(stream, start_idx.first, points_count.first,
|
||
|
|
num_voxels_prefix_sum.first,
|
||
|
|
unique_hashes_count_prefix_sum.first,
|
||
|
|
out_batch_splits, batch_size, max_voxels,
|
||
|
|
max_points_per_voxel);
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
int64_t* out_voxel_row_splits = nullptr;
|
||
|
|
if (!get_temp_size) {
|
||
|
|
output_allocator.AllocVoxelPointRowSplits(&out_voxel_row_splits,
|
||
|
|
num_valid_voxels + 1);
|
||
|
|
}
|
||
|
|
|
||
|
|
if (!get_temp_size) {
|
||
|
|
// set first element to 0
|
||
|
|
cudaMemsetAsync(out_voxel_row_splits, 0, sizeof(int64_t), stream);
|
||
|
|
}
|
||
|
|
{
|
||
|
|
std::pair<void*, size_t> inclusive_scan_temp(nullptr, 0);
|
||
|
|
|
||
|
|
cub::DeviceScan::InclusiveSum(
|
||
|
|
inclusive_scan_temp.first, inclusive_scan_temp.second,
|
||
|
|
points_count.first, out_voxel_row_splits + 1, num_valid_voxels,
|
||
|
|
stream);
|
||
|
|
|
||
|
|
inclusive_scan_temp = mem_temp.Alloc(inclusive_scan_temp.second);
|
||
|
|
if (!get_temp_size) {
|
||
|
|
cub::DeviceScan::InclusiveSum(
|
||
|
|
inclusive_scan_temp.first, inclusive_scan_temp.second,
|
||
|
|
points_count.first, out_voxel_row_splits + 1,
|
||
|
|
num_valid_voxels, stream);
|
||
|
|
}
|
||
|
|
mem_temp.Free(inclusive_scan_temp);
|
||
|
|
}
|
||
|
|
|
||
|
|
if (get_temp_size) {
|
||
|
|
// return the memory peak as the required temporary memory size.
|
||
|
|
temp_size = mem_temp.MaxUsed();
|
||
|
|
return;
|
||
|
|
}
|
||
|
|
|
||
|
|
int32_t* out_voxel_coords = nullptr;
|
||
|
|
output_allocator.AllocVoxelCoords(&out_voxel_coords, num_valid_voxels,
|
||
|
|
NDIM);
|
||
|
|
ComputeVoxelCoords(stream, out_voxel_coords, points,
|
||
|
|
point_indices_dbuf.Current(), start_idx.first,
|
||
|
|
points_range_min_vec, inv_voxel_size, num_valid_voxels);
|
||
|
|
|
||
|
|
int64_t num_valid_points = 0;
|
||
|
|
{
|
||
|
|
cudaMemcpyAsync(&num_valid_points,
|
||
|
|
out_voxel_row_splits + num_valid_voxels,
|
||
|
|
sizeof(int64_t), cudaMemcpyDeviceToHost, stream);
|
||
|
|
// wait for the async copies
|
||
|
|
while (cudaErrorNotReady == cudaStreamQuery(stream)) { /*empty*/
|
||
|
|
}
|
||
|
|
}
|
||
|
|
int64_t* out_point_indices = nullptr;
|
||
|
|
output_allocator.AllocVoxelPointIndices(&out_point_indices,
|
||
|
|
num_valid_points);
|
||
|
|
CopyPointIndices(stream, out_point_indices, point_indices_dbuf.Current(),
|
||
|
|
start_idx.first, out_voxel_row_splits + 1,
|
||
|
|
num_valid_voxels);
|
||
|
|
}
|
||
|
|
|
||
|
|
} // namespace impl
|
||
|
|
} // namespace ml
|
||
|
|
} // namespace open3d
|