wehub-resource-sync
413 lines
16 KiB
Plaintext
413 lines
16 KiB
Plaintext
/*
|
|
* SPDX-FileCopyrightText: Copyright (c) 1993-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
|
* SPDX-License-Identifier: Apache-2.0
|
|
*
|
|
* Licensed under the Apache License, Version 2.0 (the "License");
|
|
* you may not use this file except in compliance with the License.
|
|
* You may obtain a copy of the License at
|
|
*
|
|
* http://www.apache.org/licenses/LICENSE-2.0
|
|
*
|
|
* Unless required by applicable law or agreed to in writing, software
|
|
* distributed under the License is distributed on an "AS IS" BASIS,
|
|
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
* See the License for the specific language governing permissions and
|
|
* limitations under the License.
|
|
*/
|
|
|
|
#include <iostream>
|
|
#include <cuda_runtime_api.h>
|
|
namespace nvinfer1
|
|
{
|
|
namespace plugin
|
|
{
|
|
__global__ void generateVoxels_kernel(
|
|
int max_num_points,
|
|
float *points, unsigned int* points_size,
|
|
float min_x_range, float max_x_range,
|
|
float min_y_range, float max_y_range,
|
|
float min_z_range, float max_z_range,
|
|
float pillar_x_size, float pillar_y_size, float pillar_z_size,
|
|
int grid_y_size, int grid_x_size, int num_point_values,
|
|
int max_points_per_voxel,
|
|
unsigned int *mask, float *voxels)
|
|
{
|
|
int point_idx = blockIdx.x * blockDim.x + threadIdx.x;
|
|
int batch_idx = point_idx / max_num_points;
|
|
int point_idx_in_frame = point_idx % max_num_points;
|
|
if(point_idx_in_frame >= points_size[batch_idx]) return;
|
|
float px = points[num_point_values * point_idx];
|
|
float py = points[num_point_values * point_idx + 1];
|
|
float pz = points[num_point_values * point_idx + 2];
|
|
float pw = points[num_point_values * point_idx + 3];
|
|
float pt;
|
|
if (num_point_values == 5) {
|
|
pt = points[num_point_values * point_idx + 4];
|
|
}
|
|
if(px<min_x_range||px>=max_x_range
|
|
|| py<min_y_range||py>=max_y_range
|
|
|| pz<min_z_range||pz>=max_z_range) return;
|
|
int voxel_idx = floorf((px - min_x_range)/pillar_x_size);
|
|
int voxel_idy = floorf((py - min_y_range)/pillar_y_size);
|
|
unsigned int voxel_index = (batch_idx * grid_y_size + voxel_idy) * grid_x_size + voxel_idx;
|
|
unsigned int point_id = atomicAdd(&(mask[voxel_index]), 1);
|
|
if(point_id >= max_points_per_voxel) return;
|
|
float *address = voxels + (voxel_index*max_points_per_voxel + point_id)*num_point_values;
|
|
atomicExch(address+0, px);
|
|
atomicExch(address+1, py);
|
|
atomicExch(address+2, pz);
|
|
atomicExch(address+3, pw);
|
|
if (num_point_values == 5) {
|
|
atomicExch(address+4, pt);
|
|
}
|
|
}
|
|
|
|
__global__ void generateBaseFeatures_kernel(
|
|
int batch_size,
|
|
unsigned int *mask, float *voxels,
|
|
int grid_y_size, int grid_x_size,
|
|
unsigned int *pillar_num,
|
|
int max_pillar_num,
|
|
int max_points_per_voxel,
|
|
int num_point_values,
|
|
float *voxel_features,
|
|
unsigned int *voxel_num_points,
|
|
unsigned int *coords)
|
|
{
|
|
int voxel_id = blockIdx.x * blockDim.x + threadIdx.x;
|
|
int voxel_idx = voxel_id % grid_x_size;
|
|
int voxel_idy = (voxel_id / grid_x_size) % grid_y_size;
|
|
int batch_id = voxel_id / (grid_y_size * grid_x_size);
|
|
if (batch_id >= batch_size) return;
|
|
unsigned int count = mask[voxel_id];
|
|
if( !(count>0) ) return;
|
|
count = count<max_points_per_voxel?count:max_points_per_voxel;
|
|
int current_pillarId = 0;
|
|
current_pillarId = atomicAdd(pillar_num + batch_id, 1);
|
|
voxel_num_points[batch_id * grid_y_size * grid_x_size + current_pillarId] = count;
|
|
int4 coord = {0, 0, voxel_idy, voxel_idx};
|
|
((int4*)coords)[batch_id * max_pillar_num + current_pillarId] = coord;
|
|
for (int i=0; i<count; i++){
|
|
int inIndex = voxel_id*max_points_per_voxel + i;
|
|
int outIndex = (batch_id * grid_x_size * grid_y_size + current_pillarId)*max_points_per_voxel + i;
|
|
if (num_point_values == 4) {
|
|
((float4*)voxel_features)[outIndex] = ((float4*)voxels)[inIndex];
|
|
}
|
|
else if (num_point_values == 5) {
|
|
for(int k=0; k<5;k++)
|
|
voxel_features[5 * outIndex + k] = voxels[5 * inIndex + k];
|
|
}
|
|
}
|
|
}
|
|
|
|
void generateVoxels_launch(
|
|
int batch_size, int max_num_points,
|
|
float *points, unsigned int* points_size,
|
|
float min_x_range, float max_x_range,
|
|
float min_y_range, float max_y_range,
|
|
float min_z_range, float max_z_range,
|
|
float pillar_x_size, float pillar_y_size, float pillar_z_size,
|
|
int grid_y_size, int grid_x_size, int num_point_values,
|
|
int max_points_per_voxel,
|
|
unsigned int *mask, float *voxels,
|
|
cudaStream_t stream)
|
|
{
|
|
int threadNum = 256;
|
|
dim3 blocks((batch_size * max_num_points + threadNum - 1) / threadNum);
|
|
dim3 threads(threadNum);
|
|
generateVoxels_kernel<<<blocks, threads, 0, stream>>>
|
|
(max_num_points,
|
|
points, points_size,
|
|
min_x_range, max_x_range,
|
|
min_y_range, max_y_range,
|
|
min_z_range, max_z_range,
|
|
pillar_x_size, pillar_y_size, pillar_z_size,
|
|
grid_y_size, grid_x_size, num_point_values,
|
|
max_points_per_voxel,
|
|
mask, voxels);
|
|
}
|
|
|
|
void generateBaseFeatures_launch(
|
|
int batch_size,
|
|
unsigned int *mask, float *voxels,
|
|
int grid_y_size, int grid_x_size,
|
|
unsigned int *pillar_num,
|
|
int max_pillar_num,
|
|
int max_points_per_voxel,
|
|
int num_point_values,
|
|
float *voxel_features,
|
|
unsigned int *voxel_num_points,
|
|
unsigned int *coords,
|
|
cudaStream_t stream)
|
|
{
|
|
int blockSize = 1024;
|
|
dim3 threads(blockSize);
|
|
dim3 blocks((batch_size * grid_y_size * grid_x_size + blockSize - 1) / blockSize);
|
|
generateBaseFeatures_kernel<<<blocks, threads, 0, stream>>>
|
|
(
|
|
batch_size,
|
|
mask, voxels, grid_y_size, grid_x_size,
|
|
pillar_num,
|
|
max_pillar_num,
|
|
max_points_per_voxel,
|
|
num_point_values,
|
|
voxel_features,
|
|
voxel_num_points,
|
|
coords
|
|
);
|
|
}
|
|
|
|
__global__ void generateFeatures_kernel(
|
|
int batch_size,
|
|
int dense_pillar_num,
|
|
float* voxel_features,
|
|
unsigned int* voxel_num_points,
|
|
unsigned int* coords, unsigned int *params,
|
|
float voxel_x, float voxel_y, float voxel_z,
|
|
float range_min_x, float range_min_y, float range_min_z,
|
|
unsigned int voxel_features_size, unsigned int max_points,
|
|
unsigned int max_voxels,
|
|
float* features)
|
|
{
|
|
int warp_size = max_points;
|
|
int pillar_idx = blockIdx.x * 4 + threadIdx.x/warp_size;
|
|
int point_idx = threadIdx.x % warp_size;
|
|
// In case the actual number of points is less than warp_size
|
|
// E.g., warp_size=32, max_points=20
|
|
if (point_idx >= max_points) return;
|
|
int batch_idx = pillar_idx / max_voxels;
|
|
if (batch_idx >= batch_size) return;
|
|
int pillar_idx_in_frame = pillar_idx % max_voxels;
|
|
int dense_pillar_idx = pillar_idx_in_frame + dense_pillar_num * batch_idx;
|
|
int pillar_idx_inBlock = threadIdx.x/warp_size;
|
|
// Limit number of voxels to max_voxels
|
|
unsigned int num_pillars = params[batch_idx] > max_voxels ? max_voxels : params[batch_idx];
|
|
// Update max_voxel to actual number
|
|
if (pillar_idx_in_frame == 0 && point_idx == 0) {
|
|
params[batch_idx] = num_pillars;
|
|
}
|
|
if (pillar_idx_in_frame >= num_pillars) return;
|
|
|
|
//load src
|
|
__shared__ float pillarSM[4][64][5]; // up to 64 points per pillar
|
|
__shared__ float4 pillarSumSM[4]; //4*4
|
|
__shared__ int4 cordsSM[4]; //4*4
|
|
__shared__ int pointsNumSM[4]; //4
|
|
__shared__ float pillarOutSM[4][64][11]; // up to 11 features per point
|
|
|
|
if (point_idx == 0) {
|
|
pointsNumSM[pillar_idx_inBlock] = voxel_num_points[dense_pillar_idx];
|
|
cordsSM[pillar_idx_inBlock] = ((int4*)coords)[dense_pillar_idx];
|
|
pillarSumSM[pillar_idx_inBlock] = {0,0,0,0};
|
|
}
|
|
for(int k=0; k<5; k++) {
|
|
pillarSM[pillar_idx_inBlock][point_idx][k] = voxel_features[5 * (dense_pillar_idx*max_points + point_idx) + k];
|
|
}
|
|
__syncthreads();
|
|
//calculate sm
|
|
if (point_idx < pointsNumSM[pillar_idx_inBlock]) {
|
|
atomicAdd(&(pillarSumSM[pillar_idx_inBlock].x), pillarSM[pillar_idx_inBlock][point_idx][0]);
|
|
atomicAdd(&(pillarSumSM[pillar_idx_inBlock].y), pillarSM[pillar_idx_inBlock][point_idx][1]);
|
|
atomicAdd(&(pillarSumSM[pillar_idx_inBlock].z), pillarSM[pillar_idx_inBlock][point_idx][2]);
|
|
}
|
|
__syncthreads();
|
|
//feature-mean
|
|
float4 mean;
|
|
float validPoints = pointsNumSM[pillar_idx_inBlock];
|
|
mean.x = pillarSumSM[pillar_idx_inBlock].x / validPoints;
|
|
mean.y = pillarSumSM[pillar_idx_inBlock].y / validPoints;
|
|
mean.z = pillarSumSM[pillar_idx_inBlock].z / validPoints;
|
|
mean.x = pillarSM[pillar_idx_inBlock][point_idx][0] - mean.x;
|
|
mean.y = pillarSM[pillar_idx_inBlock][point_idx][1] - mean.y;
|
|
mean.z = pillarSM[pillar_idx_inBlock][point_idx][2] - mean.z;
|
|
//calculate offset
|
|
float x_offset = voxel_x / 2.0f + cordsSM[pillar_idx_inBlock].w * voxel_x + range_min_x;
|
|
float y_offset = voxel_y / 2.0f + cordsSM[pillar_idx_inBlock].z * voxel_y + range_min_y;
|
|
float z_offset = voxel_z / 2.0f + cordsSM[pillar_idx_inBlock].y * voxel_z + range_min_z;
|
|
//feature-offset
|
|
float4 center;
|
|
center.x = pillarSM[pillar_idx_inBlock][point_idx][0] - x_offset;
|
|
center.y = pillarSM[pillar_idx_inBlock][point_idx][1] - y_offset;
|
|
center.z = pillarSM[pillar_idx_inBlock][point_idx][2] - z_offset;
|
|
//store output
|
|
if (point_idx < pointsNumSM[pillar_idx_inBlock]) {
|
|
for(int k=0; k<5; k++)
|
|
pillarOutSM[pillar_idx_inBlock][point_idx][k] = pillarSM[pillar_idx_inBlock][point_idx][k];
|
|
pillarOutSM[pillar_idx_inBlock][point_idx][5] = mean.x;
|
|
pillarOutSM[pillar_idx_inBlock][point_idx][5 + 1] = mean.y;
|
|
pillarOutSM[pillar_idx_inBlock][point_idx][5 + 2] = mean.z;
|
|
pillarOutSM[pillar_idx_inBlock][point_idx][5 + 3] = center.x;
|
|
pillarOutSM[pillar_idx_inBlock][point_idx][5 + 4] = center.y;
|
|
if (5 + 5 < voxel_features_size)
|
|
pillarOutSM[pillar_idx_inBlock][point_idx][warp_size + 5] = center.z;
|
|
} else {
|
|
for (int k = 0; k < voxel_features_size; k++)
|
|
pillarOutSM[pillar_idx_inBlock][point_idx][k] = 0;
|
|
}
|
|
__syncthreads();
|
|
for(int i = 0; i < voxel_features_size; i ++) {
|
|
int outputSMId = pillar_idx_inBlock*64*11 + point_idx * 11 + i;
|
|
int outputId = pillar_idx*max_points*voxel_features_size + point_idx * voxel_features_size + i;
|
|
features[outputId] = ((float*)pillarOutSM)[outputSMId] ;
|
|
}
|
|
|
|
}
|
|
|
|
__global__ void generateFeatures_kernel_4x(
|
|
int batch_size,
|
|
int dense_pillar_num,
|
|
float* voxel_features,
|
|
unsigned int* voxel_num_points, unsigned int* coords,
|
|
unsigned int *params,
|
|
float voxel_x, float voxel_y, float voxel_z,
|
|
float range_min_x, float range_min_y, float range_min_z,
|
|
unsigned int voxel_features_size, unsigned int max_points,
|
|
unsigned int max_voxels,
|
|
float* features)
|
|
{
|
|
int warp_size = max_points;
|
|
int pillar_idx = blockIdx.x * 4 + threadIdx.x / warp_size;
|
|
int point_idx = threadIdx.x % warp_size;
|
|
// In case the actual number of points is less than warp_size
|
|
// E.g., warp_size=32, max_points=20
|
|
if (point_idx >= max_points) return;
|
|
int batch_idx = pillar_idx / max_voxels;
|
|
if (batch_idx >= batch_size) return;
|
|
int pillar_idx_in_frame = pillar_idx % max_voxels;
|
|
int dense_pillar_idx = pillar_idx_in_frame + dense_pillar_num * batch_idx;
|
|
int pillar_idx_inBlock = threadIdx.x / warp_size;
|
|
// Limit number of voxels to max_voxels
|
|
unsigned int num_pillars = params[batch_idx] > max_voxels ? max_voxels : params[batch_idx];
|
|
// Update max_voxel to actual number
|
|
if (pillar_idx_in_frame == 0 && point_idx == 0) {
|
|
params[batch_idx] = num_pillars;
|
|
}
|
|
if (pillar_idx_in_frame >= num_pillars) return;
|
|
//load src
|
|
__shared__ float4 pillarSM[4][64]; // up to 64 points per pillar
|
|
__shared__ float4 pillarSumSM[4]; //4*4
|
|
__shared__ int4 cordsSM[4]; //4*4
|
|
__shared__ int pointsNumSM[4]; //4
|
|
__shared__ float pillarOutSM[4][64][11]; // up to 11 output features per point
|
|
|
|
if (point_idx == 0) {
|
|
pointsNumSM[pillar_idx_inBlock] = voxel_num_points[dense_pillar_idx];
|
|
cordsSM[pillar_idx_inBlock] = ((int4*)coords)[pillar_idx];
|
|
pillarSumSM[pillar_idx_inBlock] = {0,0,0,0};
|
|
}
|
|
pillarSM[pillar_idx_inBlock][point_idx] = ((float4*)voxel_features)[dense_pillar_idx*max_points + point_idx];
|
|
__syncthreads();
|
|
//calculate sm
|
|
if (point_idx < pointsNumSM[pillar_idx_inBlock]) {
|
|
atomicAdd(&(pillarSumSM[pillar_idx_inBlock].x), pillarSM[pillar_idx_inBlock][point_idx].x);
|
|
atomicAdd(&(pillarSumSM[pillar_idx_inBlock].y), pillarSM[pillar_idx_inBlock][point_idx].y);
|
|
atomicAdd(&(pillarSumSM[pillar_idx_inBlock].z), pillarSM[pillar_idx_inBlock][point_idx].z);
|
|
}
|
|
__syncthreads();
|
|
//feature-mean
|
|
float4 mean;
|
|
float validPoints = pointsNumSM[pillar_idx_inBlock];
|
|
mean.x = pillarSumSM[pillar_idx_inBlock].x / validPoints;
|
|
mean.y = pillarSumSM[pillar_idx_inBlock].y / validPoints;
|
|
mean.z = pillarSumSM[pillar_idx_inBlock].z / validPoints;
|
|
mean.x = pillarSM[pillar_idx_inBlock][point_idx].x - mean.x;
|
|
mean.y = pillarSM[pillar_idx_inBlock][point_idx].y - mean.y;
|
|
mean.z = pillarSM[pillar_idx_inBlock][point_idx].z - mean.z;
|
|
//calculate offset
|
|
float x_offset = voxel_x / 2.0f + cordsSM[pillar_idx_inBlock].w * voxel_x + range_min_x;
|
|
float y_offset = voxel_y / 2.0f + cordsSM[pillar_idx_inBlock].z * voxel_y + range_min_y;
|
|
float z_offset = voxel_z / 2.0f + cordsSM[pillar_idx_inBlock].y * voxel_z + range_min_z;
|
|
//feature-offset
|
|
float4 center;
|
|
center.x = pillarSM[pillar_idx_inBlock][point_idx].x - x_offset;
|
|
center.y = pillarSM[pillar_idx_inBlock][point_idx].y - y_offset;
|
|
center.z = pillarSM[pillar_idx_inBlock][point_idx].z - z_offset;
|
|
//store output
|
|
if (point_idx < pointsNumSM[pillar_idx_inBlock]) {
|
|
pillarOutSM[pillar_idx_inBlock][point_idx][0] = pillarSM[pillar_idx_inBlock][point_idx].x;
|
|
pillarOutSM[pillar_idx_inBlock][point_idx][1] = pillarSM[pillar_idx_inBlock][point_idx].y;
|
|
pillarOutSM[pillar_idx_inBlock][point_idx][2] = pillarSM[pillar_idx_inBlock][point_idx].z;
|
|
pillarOutSM[pillar_idx_inBlock][point_idx][3] = pillarSM[pillar_idx_inBlock][point_idx].w;
|
|
pillarOutSM[pillar_idx_inBlock][point_idx][4] = mean.x;
|
|
pillarOutSM[pillar_idx_inBlock][point_idx][5] = mean.y;
|
|
pillarOutSM[pillar_idx_inBlock][point_idx][6] = mean.z;
|
|
pillarOutSM[pillar_idx_inBlock][point_idx][7] = center.x;
|
|
pillarOutSM[pillar_idx_inBlock][point_idx][8] = center.y;
|
|
pillarOutSM[pillar_idx_inBlock][point_idx][9] = center.z;
|
|
|
|
} else {
|
|
pillarOutSM[pillar_idx_inBlock][point_idx][0] = 0;
|
|
pillarOutSM[pillar_idx_inBlock][point_idx][1] = 0;
|
|
pillarOutSM[pillar_idx_inBlock][point_idx][2] = 0;
|
|
pillarOutSM[pillar_idx_inBlock][point_idx][3] = 0;
|
|
pillarOutSM[pillar_idx_inBlock][point_idx][4] = 0;
|
|
pillarOutSM[pillar_idx_inBlock][point_idx][5] = 0;
|
|
pillarOutSM[pillar_idx_inBlock][point_idx][6] = 0;
|
|
pillarOutSM[pillar_idx_inBlock][point_idx][7] = 0;
|
|
pillarOutSM[pillar_idx_inBlock][point_idx][8] = 0;
|
|
pillarOutSM[pillar_idx_inBlock][point_idx][9] = 0;
|
|
}
|
|
__syncthreads();
|
|
for(int i = 0; i < voxel_features_size; i ++) {
|
|
int outputSMId = pillar_idx_inBlock*64*11 + point_idx * 11 + i;
|
|
int outputId = pillar_idx*max_points*voxel_features_size + point_idx * voxel_features_size + i;
|
|
features[outputId] = ((float*)pillarOutSM)[outputSMId] ;
|
|
}
|
|
|
|
}
|
|
|
|
int generateFeatures_launch(
|
|
int batch_size,
|
|
int dense_pillar_num,
|
|
float* voxel_features,
|
|
unsigned int* voxel_num_points,
|
|
unsigned int* coords,
|
|
unsigned int *params,
|
|
float voxel_x, float voxel_y, float voxel_z,
|
|
float range_min_x, float range_min_y, float range_min_z,
|
|
unsigned int voxel_features_size, unsigned int max_points,
|
|
unsigned int max_voxels, unsigned int num_point_values,
|
|
float* features,
|
|
cudaStream_t stream)
|
|
{
|
|
unsigned int warp_size = max_points;
|
|
dim3 blocks((batch_size * max_voxels + 3) / 4);
|
|
dim3 threads(4*warp_size);
|
|
if (num_point_values == 4) {
|
|
generateFeatures_kernel_4x<<<blocks, threads, 0, stream>>>
|
|
(batch_size,
|
|
dense_pillar_num,
|
|
voxel_features,
|
|
voxel_num_points,
|
|
coords,
|
|
params,
|
|
voxel_x, voxel_y, voxel_z,
|
|
range_min_x, range_min_y, range_min_z,
|
|
voxel_features_size, max_points,
|
|
max_voxels,
|
|
features);
|
|
}
|
|
else {
|
|
generateFeatures_kernel<<<blocks, threads, 0, stream>>>
|
|
(batch_size,
|
|
dense_pillar_num,
|
|
voxel_features,
|
|
voxel_num_points,
|
|
coords,
|
|
params,
|
|
voxel_x, voxel_y, voxel_z,
|
|
range_min_x, range_min_y, range_min_z,
|
|
voxel_features_size, max_points,
|
|
max_voxels,
|
|
features);
|
|
}
|
|
auto err = cudaGetLastError();
|
|
if (cudaSuccess != err) {
|
|
fprintf(stderr, "CUDA kernel failed : %s\n", cudaGetErrorString(err));
|
|
exit(-1);
|
|
}
|
|
return err;
|
|
}
|
|
} // namespace plugin
|
|
} // namespace nvinfer1
|