Spaces:
Runtime error
Runtime error
File size: 2,739 Bytes
4893ce0 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 |
/*
Ball Query with BatchIdx
Written by Li Jiang
All Rights Reserved 2020.
*/
#include <stdio.h>
#include <stdlib.h>
#include <time.h>
#define TOTAL_THREADS 1024
#define THREADS_PER_BLOCK 512
#define DIVUP(m,n) ((m) / (n) + ((m) % (n) > 0))
/* ================================== ballquery_batch_p ================================== */
__global__ void ballquery_batch_p_cuda_(int n, int meanActive, float radius, const float *xyz, const int *batch_idxs, const int *batch_offsets, int *idx, int *start_len, int *cumsum) {
int pt_idx = blockIdx.x * blockDim.x + threadIdx.x;
if (pt_idx >= n) return;
start_len += (pt_idx * 2);
int idx_temp[1000];
float radius2 = radius * radius;
float o_x = xyz[pt_idx * 3 + 0];
float o_y = xyz[pt_idx * 3 + 1];
float o_z = xyz[pt_idx * 3 + 2];
int batch_idx = batch_idxs[pt_idx];
int start = batch_offsets[batch_idx];
int end = batch_offsets[batch_idx + 1];
int cnt = 0;
for(int k = start; k < end; k++){
float x = xyz[k * 3 + 0];
float y = xyz[k * 3 + 1];
float z = xyz[k * 3 + 2];
float d2 = (o_x - x) * (o_x - x) + (o_y - y) * (o_y - y) + (o_z - z) * (o_z - z);
if(d2 < radius2){
if(cnt < 1000){
idx_temp[cnt] = k;
}
else{
break;
}
++cnt;
}
}
start_len[0] = atomicAdd(cumsum, cnt);
start_len[1] = cnt;
int thre = n * meanActive;
if(start_len[0] >= thre) return;
idx += start_len[0];
if(start_len[0] + cnt >= thre) cnt = thre - start_len[0];
for(int k = 0; k < cnt; k++){
idx[k] = idx_temp[k];
}
}
int ballquery_batch_p_cuda(int n, int meanActive, float radius, const float *xyz, const int *batch_idxs, const int *batch_offsets, int *idx, int *start_len, cudaStream_t stream) {
// param xyz: (n, 3)
// param batch_idxs: (n)
// param batch_offsets: (B + 1)
// output idx: (n * meanActive) dim 0 for number of points in the ball, idx in n
// output start_len: (n, 2), int
cudaError_t err;
dim3 blocks(DIVUP(n, THREADS_PER_BLOCK));
dim3 threads(THREADS_PER_BLOCK);
int cumsum = 0;
int* p_cumsum;
cudaMalloc((void**)&p_cumsum, sizeof(int));
cudaMemcpy(p_cumsum, &cumsum, sizeof(int), cudaMemcpyHostToDevice);
ballquery_batch_p_cuda_<<<blocks, threads, 0, stream>>>(n, meanActive, radius, xyz, batch_idxs, batch_offsets, idx, start_len, p_cumsum);
err = cudaGetLastError();
if (cudaSuccess != err) {
fprintf(stderr, "CUDA kernel failed : %s\n", cudaGetErrorString(err));
exit(-1);
}
cudaMemcpy(&cumsum, p_cumsum, sizeof(int), cudaMemcpyDeviceToHost);
return cumsum;
}
|