mirror of
https://github.com/wassname/Pointnet2_PyTorch.git
synced 2026-07-15 01:11:03 +08:00
Initial commit
This commit is contained in:
erikwijmans
@@ -0,0 +1,20 @@
|
||||
#include <THC/THC.h>
|
||||
|
||||
#include "ball_query_gpu.h"
|
||||
|
||||
extern THCState *state;
|
||||
|
||||
int ball_query_wrapper(int b, int n, int m, float radius, int nsample,
|
||||
THCudaTensor *new_xyz_tensor, THCudaTensor *xyz_tensor,
|
||||
THCudaIntTensor *idx_tensor) {
|
||||
|
||||
const float *new_xyz = THCudaTensor_data(state, new_xyz_tensor);
|
||||
const float *xyz = THCudaTensor_data(state, xyz_tensor);
|
||||
int *idx = THCudaIntTensor_data(state, idx_tensor);
|
||||
|
||||
cudaStream_t stream = THCState_getCurrentStream(state);
|
||||
|
||||
query_ball_point_kernel_wrapper(b, n, m, radius, nsample, new_xyz, xyz,
|
||||
idx, stream);
|
||||
return 1;
|
||||
}
|
||||
@@ -0,0 +1,63 @@
|
||||
#include <math.h>
|
||||
#include <stdio.h>
|
||||
#include <stdlib.h>
|
||||
|
||||
#include "ball_query_gpu.h"
|
||||
#include "cuda_utils.h"
|
||||
|
||||
// input: new_xyz(b, m, 3) xyz(b, n, 3)
|
||||
// output: idx(b, m, nsample)
|
||||
__global__ void query_ball_point_kernel(int b, int n, int m, float radius,
|
||||
int nsample,
|
||||
const float *__restrict__ new_xyz,
|
||||
const float *__restrict__ xyz,
|
||||
int * __restrict__ idx) {
|
||||
int batch_index = blockIdx.x;
|
||||
xyz += batch_index * n * 3;
|
||||
new_xyz += batch_index * m * 3;
|
||||
idx += m * nsample * batch_index;
|
||||
|
||||
int index = threadIdx.x;
|
||||
int stride = blockDim.x;
|
||||
|
||||
float radius2 = radius * radius;
|
||||
for (int j = index; j < m; j += stride) {
|
||||
float new_x = new_xyz[j * 3 + 0];
|
||||
float new_y = new_xyz[j * 3 + 1];
|
||||
float new_z = new_xyz[j * 3 + 2];
|
||||
for (int k = 0, cnt = 0; k < n && cnt < nsample; ++k) {
|
||||
float x = xyz[k * 3 + 0];
|
||||
float y = xyz[k * 3 + 1];
|
||||
float z = xyz[k * 3 + 2];
|
||||
float d2 = (new_x - x) * (new_x - x) +
|
||||
(new_y - y) * (new_y - y) +
|
||||
(new_z - z) * (new_z - z);
|
||||
if (d2 < radius2) {
|
||||
if (cnt == 0) {
|
||||
for (int l = 0; l < nsample; ++l) {
|
||||
idx[j * nsample + l] = k;
|
||||
}
|
||||
}
|
||||
idx[j * nsample + cnt] = k;
|
||||
++cnt;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
void query_ball_point_kernel_wrapper(int b, int n, int m, float radius,
|
||||
int nsample, const float *new_xyz,
|
||||
const float *xyz, int *idx,
|
||||
cudaStream_t stream) {
|
||||
|
||||
cudaError_t err;
|
||||
query_ball_point_kernel<<<b, opt_n_threads(m), 0, stream>>>(
|
||||
b, n, m, radius, nsample, new_xyz, xyz, idx);
|
||||
|
||||
err = cudaGetLastError();
|
||||
if (cudaSuccess != err) {
|
||||
fprintf(stderr, "CUDA kernel failed : %s\n",
|
||||
cudaGetErrorString(err));
|
||||
exit(-1);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,37 @@
|
||||
#include <THC/THC.h>
|
||||
|
||||
#include "group_points_gpu.h"
|
||||
|
||||
extern THCState *state;
|
||||
|
||||
int group_points_wrapper(int b, int n, int c, int npoints, int nsample,
|
||||
THCudaTensor *points_tensor,
|
||||
THCudaIntTensor *idx_tensor,
|
||||
THCudaTensor *out_tensor) {
|
||||
|
||||
const float *points = THCudaTensor_data(state, points_tensor);
|
||||
const int *idx = THCudaIntTensor_data(state, idx_tensor);
|
||||
float *out = THCudaTensor_data(state, out_tensor);
|
||||
|
||||
cudaStream_t stream = THCState_getCurrentStream(state);
|
||||
|
||||
group_points_kernel_wrapper(b, n, c, npoints, nsample, points, idx, out,
|
||||
stream);
|
||||
return 1;
|
||||
}
|
||||
|
||||
int group_points_grad_wrapper(int b, int n, int c, int npoints, int nsample,
|
||||
THCudaTensor *grad_out_tensor,
|
||||
THCudaIntTensor *idx_tensor,
|
||||
THCudaTensor *grad_points_tensor) {
|
||||
|
||||
float *grad_points = THCudaTensor_data(state, grad_points_tensor);
|
||||
const int *idx = THCudaIntTensor_data(state, idx_tensor);
|
||||
const float *grad_out = THCudaTensor_data(state, grad_out_tensor);
|
||||
|
||||
cudaStream_t stream = THCState_getCurrentStream(state);
|
||||
|
||||
group_points_grad_kernel_wrapper(b, n, c, npoints, nsample, grad_out,
|
||||
idx, grad_points, stream);
|
||||
return 1;
|
||||
}
|
||||
@@ -0,0 +1,86 @@
|
||||
#include <stdio.h>
|
||||
#include <stdlib.h>
|
||||
|
||||
#include "group_points_gpu.h"
|
||||
#include "cuda_utils.h"
|
||||
|
||||
// input: points(b, n, c) idx(b, npoints, nsample)
|
||||
// output: out(b, npoints, nsample, c)
|
||||
__global__ void group_points_kernel(int b, int n, int c, int npoints,
|
||||
int nsample,
|
||||
const float *__restrict__ points,
|
||||
const int *__restrict__ idx,
|
||||
float *__restrict__ out) {
|
||||
int batch_index = blockIdx.x;
|
||||
points += batch_index * n * c;
|
||||
idx += batch_index * npoints * nsample;
|
||||
out += batch_index * npoints * nsample * c;
|
||||
|
||||
int index = threadIdx.x;
|
||||
int stride = blockDim.x;
|
||||
for (int j = index; j < npoints; j += stride) {
|
||||
for (int k = 0; k < nsample; ++k) {
|
||||
int ii = idx[j * nsample + k];
|
||||
memcpy(out + j * nsample * c + k * c, points + ii * c,
|
||||
sizeof(float) * c);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
void group_points_kernel_wrapper(int b, int n, int c, int npoints, int nsample,
|
||||
const float *points, const int *idx,
|
||||
float *out, cudaStream_t stream) {
|
||||
|
||||
cudaError_t err;
|
||||
group_points_kernel<<<b, opt_n_threads(npoints), 0, stream>>>(
|
||||
b, n, c, npoints, nsample, points, idx, out);
|
||||
|
||||
err = cudaGetLastError();
|
||||
if (cudaSuccess != err) {
|
||||
fprintf(stderr, "CUDA kernel failed : %s\n",
|
||||
cudaGetErrorString(err));
|
||||
exit(-1);
|
||||
}
|
||||
}
|
||||
|
||||
// input: grad_out(b, npoints, nsample, c), idx(b, npoints, nsample)
|
||||
// output: grad_points(b, n, c)
|
||||
__global__ void group_points_grad_kernel(int b, int n, int c, int npoints,
|
||||
int nsample,
|
||||
const float *__restrict__ grad_out,
|
||||
const int *__restrict__ idx,
|
||||
float *__restrict__ grad_points) {
|
||||
int batch_index = blockIdx.x;
|
||||
grad_points += batch_index * n * c;
|
||||
idx += batch_index * npoints * nsample;
|
||||
grad_out += batch_index * npoints * nsample * c;
|
||||
|
||||
int index = threadIdx.x;
|
||||
int stride = blockDim.x;
|
||||
for (int j = index; j < npoints; j += stride) {
|
||||
for (int k = 0; k < nsample; ++k) {
|
||||
int ii = idx[j * nsample + k];
|
||||
for (int l = 0; l < c; ++l) {
|
||||
atomicAdd(
|
||||
grad_points + ii * c + l,
|
||||
grad_out[j * nsample * c + k * c + l]);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
void group_points_grad_kernel_wrapper(int b, int n, int c, int npoints,
|
||||
int nsample, const float *grad_out,
|
||||
const int *idx, float *grad_points,
|
||||
cudaStream_t stream) {
|
||||
cudaError_t err;
|
||||
group_points_grad_kernel<<<b, opt_n_threads(npoints), 0, stream>>>(
|
||||
b, n, c, npoints, nsample, grad_out, idx, grad_points);
|
||||
|
||||
err = cudaGetLastError();
|
||||
if (cudaSuccess != err) {
|
||||
fprintf(stderr, "CUDA kernel failed : %s\n",
|
||||
cudaGetErrorString(err));
|
||||
exit(-1);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,52 @@
|
||||
#include <THC/THC.h>
|
||||
#include <math.h>
|
||||
#include <stdio.h>
|
||||
#include <stdlib.h>
|
||||
|
||||
#include "interpolate_gpu.h"
|
||||
|
||||
extern THCState *state;
|
||||
|
||||
void three_nn_wrapper(int b, int n, int m, THCudaTensor *unknown_tensor,
|
||||
THCudaTensor *known_tensor, THCudaTensor *dist2_tensor,
|
||||
THCudaIntTensor *idx_tensor) {
|
||||
const float *unknown = THCudaTensor_data(state, unknown_tensor);
|
||||
const float *known = THCudaTensor_data(state, known_tensor);
|
||||
float *dist2 = THCudaTensor_data(state, dist2_tensor);
|
||||
int *idx = THCudaIntTensor_data(state, idx_tensor);
|
||||
|
||||
cudaStream_t stream = THCState_getCurrentStream(state);
|
||||
three_nn_kernel_wrapper(b, n, m, unknown, known, dist2, idx, stream);
|
||||
}
|
||||
|
||||
void three_interpolate_wrapper(int b, int m, int c, int n,
|
||||
THCudaTensor *points_tensor,
|
||||
THCudaIntTensor *idx_tensor,
|
||||
THCudaTensor *weight_tensor,
|
||||
THCudaTensor *out_tensor) {
|
||||
|
||||
const float *points = THCudaTensor_data(state, points_tensor);
|
||||
const float *weight = THCudaTensor_data(state, weight_tensor);
|
||||
float *out = THCudaTensor_data(state, out_tensor);
|
||||
const int *idx = THCudaIntTensor_data(state, idx_tensor);
|
||||
|
||||
cudaStream_t stream = THCState_getCurrentStream(state);
|
||||
three_interpolate_kernel_wrapper(b, m, c, n, points, idx, weight, out,
|
||||
stream);
|
||||
}
|
||||
|
||||
void three_interpolate_grad_wrapper(int b, int n, int c, int m,
|
||||
THCudaTensor *grad_out_tensor,
|
||||
THCudaIntTensor *idx_tensor,
|
||||
THCudaTensor *weight_tensor,
|
||||
THCudaTensor *grad_points_tensor) {
|
||||
|
||||
const float *grad_out = THCudaTensor_data(state, grad_out_tensor);
|
||||
const float *weight = THCudaTensor_data(state, weight_tensor);
|
||||
float *grad_points = THCudaTensor_data(state, grad_points_tensor);
|
||||
const int *idx = THCudaIntTensor_data(state, idx_tensor);
|
||||
|
||||
cudaStream_t stream = THCState_getCurrentStream(state);
|
||||
three_interpolate_grad_kernel_wrapper(b, n, c, m, grad_out, idx, weight,
|
||||
grad_points, stream);
|
||||
}
|
||||
@@ -0,0 +1,180 @@
|
||||
#include <math.h>
|
||||
#include <stdio.h>
|
||||
#include <stdlib.h>
|
||||
|
||||
#include "interpolate_gpu.h"
|
||||
#include "cuda_utils.h"
|
||||
|
||||
// input: unknown(b, n, 3) known(b, m, 3)
|
||||
// output: dist2(b, n, 3), idx(b, n, 3)
|
||||
__global__ void three_nn_kernel(int b, int n, int m,
|
||||
const float *__restrict__ unknown,
|
||||
const float *__restrict__ known,
|
||||
float *__restrict__ dist2,
|
||||
int *__restrict__ idx) {
|
||||
int batch_index = blockIdx.x;
|
||||
unknown += batch_index * n * 3;
|
||||
known += batch_index * m * 3;
|
||||
dist2 += batch_index * n * 3;
|
||||
idx += batch_index * n * 3;
|
||||
|
||||
int index = threadIdx.x;
|
||||
int stride = blockDim.x;
|
||||
for (int j = index; j < n; j += stride) {
|
||||
float ux = unknown[j * 3 + 0];
|
||||
float uy = unknown[j * 3 + 1];
|
||||
float uz = unknown[j * 3 + 2];
|
||||
|
||||
double best1 = 1e40, best2 = 1e40, best3 = 1e40;
|
||||
int besti1 = 0, besti2 = 0, besti3 = 0;
|
||||
for (int k = 0; k < m; ++k) {
|
||||
float x = known[k * 3 + 0];
|
||||
float y = known[k * 3 + 1];
|
||||
float z = known[k * 3 + 2];
|
||||
float d =
|
||||
(ux - x) * (ux - x) + (uy - y) * (uy - y) + (uz - z) * (uz - z);
|
||||
if (d < best1) {
|
||||
best3 = best2;
|
||||
besti3 = besti2;
|
||||
best2 = best1;
|
||||
besti2 = besti1;
|
||||
best1 = d;
|
||||
besti1 = k;
|
||||
} else if (d < best2) {
|
||||
best3 = best2;
|
||||
besti3 = besti2;
|
||||
best2 = d;
|
||||
besti2 = k;
|
||||
} else if (d < best3) {
|
||||
best3 = d;
|
||||
besti3 = k;
|
||||
}
|
||||
}
|
||||
dist2[j * 3 + 0] = best1;
|
||||
dist2[j * 3 + 1] = best2;
|
||||
dist2[j * 3 + 2] = best3;
|
||||
|
||||
idx[j * 3 + 0] = besti1;
|
||||
idx[j * 3 + 1] = besti2;
|
||||
idx[j * 3 + 2] = besti3;
|
||||
}
|
||||
}
|
||||
|
||||
void three_nn_kernel_wrapper(int b, int n, int m, const float *unknown,
|
||||
const float *known, float *dist2, int *idx,
|
||||
cudaStream_t stream) {
|
||||
|
||||
cudaError_t err;
|
||||
three_nn_kernel<<<b, opt_n_threads(n), 0, stream>>>(b, n, m, unknown, known,
|
||||
dist2, idx);
|
||||
|
||||
err = cudaGetLastError();
|
||||
if (cudaSuccess != err) {
|
||||
fprintf(stderr, "CUDA kernel "
|
||||
"failed : %s\n",
|
||||
cudaGetErrorString(err));
|
||||
exit(-1);
|
||||
}
|
||||
}
|
||||
|
||||
// input: points(b, m, c), idx(b, n, 3), weight(b, n, 3)
|
||||
// output: out(b, n, c)
|
||||
__global__ void three_interpolate_kernel(int b, int m, int c, int n,
|
||||
const float *__restrict__ points,
|
||||
const int *__restrict__ idx,
|
||||
const float *__restrict__ weight,
|
||||
float *__restrict__ out) {
|
||||
int batch_index = blockIdx.x;
|
||||
points += batch_index * m * c;
|
||||
|
||||
idx += batch_index * n * 3;
|
||||
weight += batch_index * n * 3;
|
||||
|
||||
out += batch_index * n * c;
|
||||
|
||||
int index = threadIdx.x;
|
||||
int stride = blockDim.x;
|
||||
for (int j = index; j < n; j += stride) {
|
||||
float w1 = weight[j * 3 + 0];
|
||||
float w2 = weight[j * 3 + 1];
|
||||
float w3 = weight[j * 3 + 2];
|
||||
|
||||
int i1 = idx[j * 3 + 0];
|
||||
int i2 = idx[j * 3 + 1];
|
||||
int i3 = idx[j * 3 + 2];
|
||||
|
||||
for (int l = 0; l < c; ++l) {
|
||||
out[j * c + l] = points[i1 * c + l] * w1 + points[i2 * c + l] * w2 +
|
||||
points[i3 * c + l] * w3;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
void three_interpolate_kernel_wrapper(int b, int m, int c, int n,
|
||||
const float *points, const int *idx,
|
||||
const float *weight, float *out,
|
||||
cudaStream_t stream) {
|
||||
|
||||
cudaError_t err;
|
||||
three_interpolate_kernel<<<b, opt_n_threads(n) / 4, 0, stream>>>(
|
||||
b, m, c, n, points, idx, weight, out);
|
||||
|
||||
err = cudaGetLastError();
|
||||
if (cudaSuccess != err) {
|
||||
fprintf(stderr, "CUDA kernel "
|
||||
"failed : %s\n",
|
||||
cudaGetErrorString(err));
|
||||
exit(-1);
|
||||
}
|
||||
}
|
||||
|
||||
// input: grad_out(b, n, c), idx(b, n, 3), weight(b, n, 3)
|
||||
// output: grad_points(b, m, c)
|
||||
|
||||
__global__ void three_interpolate_grad_kernel(
|
||||
int b, int n, int c, int m, const float *__restrict__ grad_out,
|
||||
const int *__restrict__ idx, const float *__restrict__ weight,
|
||||
float *__restrict__ grad_points) {
|
||||
int batch_index = blockIdx.x;
|
||||
grad_out += batch_index * n * c;
|
||||
idx += batch_index * n * 3;
|
||||
weight += batch_index * n * 3;
|
||||
grad_points += batch_index * m * c;
|
||||
|
||||
int index = threadIdx.x;
|
||||
int stride = blockDim.x;
|
||||
for (int j = index; j < n; j += stride) {
|
||||
float w1 = weight[j * 3 + 0];
|
||||
float w2 = weight[j * 3 + 1];
|
||||
float w3 = weight[j * 3 + 2];
|
||||
|
||||
int i1 = idx[j * 3 + 0];
|
||||
int i2 = idx[j * 3 + 1];
|
||||
int i3 = idx[j * 3 + 2];
|
||||
|
||||
for (int l = 0; l < c; ++l) {
|
||||
atomicAdd(grad_points + i1 * c + l, grad_out[j * c + l] * w1);
|
||||
atomicAdd(grad_points + i2 * c + l, grad_out[j * c + l] * w2);
|
||||
atomicAdd(grad_points + i3 * c + l, grad_out[j * c + l] * w3);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
void three_interpolate_grad_kernel_wrapper(int b, int n, int c, int m,
|
||||
const float *grad_out,
|
||||
const int *idx, const float *weight,
|
||||
float *grad_points,
|
||||
cudaStream_t stream) {
|
||||
|
||||
cudaError_t err;
|
||||
three_interpolate_grad_kernel<<<b, opt_n_threads(n) / 4, 0, stream>>>(
|
||||
b, n, c, m, grad_out, idx, weight, grad_points);
|
||||
|
||||
err = cudaGetLastError();
|
||||
if (cudaSuccess != err) {
|
||||
fprintf(stderr, "CUDA kernel "
|
||||
"failed : %s\n",
|
||||
cudaGetErrorString(err));
|
||||
exit(-1);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,157 @@
|
||||
#include <stdio.h>
|
||||
#include <stdlib.h>
|
||||
|
||||
#include "cuda_utils.h"
|
||||
#include "roi_mask_points_gpu.h"
|
||||
|
||||
// roi format: [w, d, h, theta, cx, cy, cz]
|
||||
__device__ bool is_in_roi(const float *__restrict__ xyz,
|
||||
const float *__restrict__ roi) {
|
||||
const float w = roi[0], d = roi[1], h = roi[2], theta = roi[3], cx = roi[4],
|
||||
cy = roi[5], cz = roi[6];
|
||||
const float x = xyz[0], y = xyz[1], z = xyz[2];
|
||||
|
||||
const float sinval = sin(theta);
|
||||
const float cosval = cos(theta);
|
||||
|
||||
const float bx_x = w * cosval;
|
||||
const float bx_y = d * -sinval;
|
||||
|
||||
const float by_x = w * sinval;
|
||||
const float by_y = d * cosval;
|
||||
|
||||
const float dx = fabs(x - cx), dy = fabs(y - cy), dz = fabs(z - cz);
|
||||
|
||||
return dx <= fabs(bx_x + by_x) && dy <= fabs(bx_y + by_y) && dz <= h;
|
||||
}
|
||||
|
||||
// Input rois (n_roi, 7), batch_indices (n_roi), data_xyz (b, n, 3)
|
||||
// Ouput mask (n_roi, n)
|
||||
__global__ void roi_mask_kernel(int n_roi, int b, int n,
|
||||
const float *__restrict__ rois,
|
||||
const long *__restrict__ batch_indices,
|
||||
const float *__restrict__ data_xyz,
|
||||
unsigned char *__restrict__ mask) {
|
||||
|
||||
const int block_idx = blockIdx.x;
|
||||
const float *__restrict__ roi = rois + block_idx * 7;
|
||||
mask += block_idx * n;
|
||||
|
||||
const long batch_idx = batch_indices[block_idx];
|
||||
data_xyz += batch_idx * n * 3;
|
||||
|
||||
const int thread_idx = threadIdx.x;
|
||||
const int thread_stride = blockDim.x;
|
||||
for (int j = thread_idx; j < n; j += thread_stride) {
|
||||
const float *__restrict__ xyz = data_xyz + j * 3;
|
||||
mask[j] = is_in_roi(xyz, roi) ? 1 : 0;
|
||||
}
|
||||
}
|
||||
|
||||
void roi_mask_kernel_wrapper(int n_roi, int b, int n, const float *rois,
|
||||
const long *batch_indices, const float *data_xyz,
|
||||
unsigned char *mask, cudaStream_t stream) {
|
||||
|
||||
cudaError_t err;
|
||||
unsigned int n_threads = opt_n_threads(n);
|
||||
|
||||
roi_mask_kernel<<<n_roi, n_threads, 0, stream>>>(
|
||||
n_roi, b, n, rois, batch_indices, data_xyz, mask);
|
||||
|
||||
err = cudaGetLastError();
|
||||
if (cudaSuccess != err) {
|
||||
fprintf(stderr, "CUDA kernel failed : %s\n", cudaGetErrorString(err));
|
||||
exit(-1);
|
||||
}
|
||||
}
|
||||
|
||||
// Input mask(n_roi, n) batch_indices (n_roi), points (b, n, d)
|
||||
// Ouput count (n_roi,) descriptors (n_roi, d)
|
||||
__global__ void roi_avg_pool_kernel_forward(
|
||||
int n_roi, int b, int n, int d, const unsigned char *__restrict__ mask,
|
||||
const long *__restrict__ batch_indices, const float *__restrict__ points,
|
||||
float *__restrict__ descriptors) {
|
||||
|
||||
const int block_idx = blockIdx.x;
|
||||
mask += block_idx * n;
|
||||
descriptors += block_idx * d;
|
||||
|
||||
const long batch_idx = batch_indices[block_idx];
|
||||
points += batch_idx * n * d;
|
||||
|
||||
const int thread_idx = threadIdx.x;
|
||||
const int thread_stride = blockDim.x;
|
||||
|
||||
for (int j = thread_idx; j < n; j += thread_stride) {
|
||||
if (mask[j] == 1) {
|
||||
for (int c = 0; c < d; ++c) {
|
||||
atomicAdd(descriptors + c, points[j * d + c]);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
void roi_avg_pool_kernel_forward_wrapper(int n_roi, int b, int n, int d,
|
||||
const unsigned char *mask,
|
||||
const long *batch_indices,
|
||||
const float *points,
|
||||
float *descriptors,
|
||||
cudaStream_t stream) {
|
||||
|
||||
cudaError_t err;
|
||||
unsigned int n_threads = opt_n_threads(n);
|
||||
|
||||
roi_avg_pool_kernel_forward<<<n_roi, n_threads, 0, stream>>>(
|
||||
n_roi, b, n, d, mask, batch_indices, points, descriptors);
|
||||
|
||||
err = cudaGetLastError();
|
||||
if (cudaSuccess != err) {
|
||||
fprintf(stderr, "CUDA kernel failed : %s\n", cudaGetErrorString(err));
|
||||
exit(-1);
|
||||
}
|
||||
}
|
||||
|
||||
__global__ void
|
||||
roi_avg_pool_kernel_backward(int n_roi, int b, int n, int d,
|
||||
const unsigned char *__restrict__ mask,
|
||||
const long *__restrict__ batch_indices,
|
||||
const float *__restrict__ grad_descriptors,
|
||||
float *__restrict__ grad_points) {
|
||||
|
||||
const int block_idx = blockIdx.x;
|
||||
mask += block_idx * n;
|
||||
grad_descriptors += block_idx * d;
|
||||
|
||||
const long batch_idx = batch_indices[block_idx];
|
||||
grad_points += batch_idx * n * d;
|
||||
|
||||
const int thread_idx = threadIdx.x;
|
||||
const int thread_stride = blockDim.x;
|
||||
for (int j = thread_idx; j < n; j += thread_stride) {
|
||||
if (mask[j] == 1) {
|
||||
for (int c = 0; c < d; ++c) {
|
||||
atomicAdd(grad_points + j * d + c, grad_descriptors[c]);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
void roi_avg_pool_kernel_backward_wrapper(int n_roi, int b, int n, int d,
|
||||
const unsigned char *mask,
|
||||
const long *batch_indices,
|
||||
const float *grad_descriptors,
|
||||
float *grad_points,
|
||||
cudaStream_t stream) {
|
||||
|
||||
cudaError_t err;
|
||||
unsigned int n_threads = opt_n_threads(n);
|
||||
|
||||
roi_avg_pool_kernel_backward<<<n_roi, n_threads, 0, stream>>>(
|
||||
n_roi, b, n, d, mask, batch_indices, grad_descriptors, grad_points);
|
||||
|
||||
err = cudaGetLastError();
|
||||
if (cudaSuccess != err) {
|
||||
fprintf(stderr, "CUDA kernel failed : %s\n", cudaGetErrorString(err));
|
||||
exit(-1);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,63 @@
|
||||
#include <THC/THC.h>
|
||||
|
||||
#include "roi_mask_points_gpu.h"
|
||||
|
||||
extern THCState *state;
|
||||
|
||||
int roi_mask_wrapper(int n_roi, int b, int n, THCudaTensor *rois_tensor,
|
||||
THCudaLongTensor *batch_indices_tensor,
|
||||
THCudaTensor *data_xyz_tensor,
|
||||
THCudaByteTensor *mask_tensor) {
|
||||
|
||||
const float *rois = THCudaTensor_data(state, rois_tensor);
|
||||
const long *batch_indices =
|
||||
THCudaLongTensor_data(state, batch_indices_tensor);
|
||||
const float *data_xyz = THCudaTensor_data(state, data_xyz_tensor);
|
||||
unsigned char *mask = THCudaByteTensor_data(state, mask_tensor);
|
||||
|
||||
cudaStream_t stream = THCState_getCurrentStream(state);
|
||||
|
||||
roi_mask_kernel_wrapper(n_roi, b, n, rois, batch_indices, data_xyz,
|
||||
mask, stream);
|
||||
return 1;
|
||||
}
|
||||
|
||||
int roi_avg_pool_forward_wrapper(int n_roi, int b, int n, int d,
|
||||
THCudaByteTensor *mask_tensor,
|
||||
THCudaLongTensor *batch_indices_tensor,
|
||||
THCudaTensor *points_tensor,
|
||||
THCudaTensor *descriptors_tensor) {
|
||||
|
||||
const long *batch_indices =
|
||||
THCudaLongTensor_data(state, batch_indices_tensor);
|
||||
const unsigned char *mask = THCudaByteTensor_data(state, mask_tensor);
|
||||
const float *points = THCudaTensor_data(state, points_tensor);
|
||||
float *descriptors = THCudaTensor_data(state, descriptors_tensor);
|
||||
|
||||
cudaStream_t stream = THCState_getCurrentStream(state);
|
||||
roi_avg_pool_kernel_forward_wrapper(n_roi, b, n, d, mask, batch_indices,
|
||||
points, descriptors, stream);
|
||||
|
||||
return 1;
|
||||
}
|
||||
|
||||
int roi_avg_pool_backward_wrapper(int n_roi, int b, int n, int d,
|
||||
THCudaByteTensor *mask_tensor,
|
||||
THCudaLongTensor *batch_indices_tensor,
|
||||
THCudaTensor *grad_descriptors_tensor,
|
||||
THCudaTensor *grad_points_tensor) {
|
||||
|
||||
const long *batch_indices =
|
||||
THCudaLongTensor_data(state, batch_indices_tensor);
|
||||
const unsigned char *mask = THCudaByteTensor_data(state, mask_tensor);
|
||||
const float *grad_descriptors =
|
||||
THCudaTensor_data(state, grad_descriptors_tensor);
|
||||
float *grad_points = THCudaTensor_data(state, grad_points_tensor);
|
||||
|
||||
cudaStream_t stream = THCState_getCurrentStream(state);
|
||||
roi_avg_pool_kernel_backward_wrapper(n_roi, b, n, d, mask,
|
||||
batch_indices, grad_descriptors,
|
||||
grad_points, stream);
|
||||
|
||||
return 1;
|
||||
}
|
||||
@@ -0,0 +1,37 @@
|
||||
#include <THC/THC.h>
|
||||
|
||||
#include "sampling_gpu.h"
|
||||
|
||||
extern THCState *state;
|
||||
|
||||
int gather_points_wrapper(int b, int n, int c, int npoints,
|
||||
THCudaTensor *points_tensor,
|
||||
THCudaIntTensor *idx_tensor,
|
||||
THCudaTensor *out_tensor) {
|
||||
|
||||
const float *points = THCudaTensor_data(state, points_tensor);
|
||||
const int *idx = THCudaIntTensor_data(state, idx_tensor);
|
||||
float *out = THCudaTensor_data(state, out_tensor);
|
||||
|
||||
cudaStream_t stream = THCState_getCurrentStream(state);
|
||||
|
||||
gather_points_kernel_wrapper(b, n, c, npoints, points, idx, out,
|
||||
stream);
|
||||
return 1;
|
||||
}
|
||||
|
||||
int furthest_point_sampling_wrapper(int b, int n, int m,
|
||||
THCudaTensor *points_tensor,
|
||||
THCudaTensor *temp_tensor,
|
||||
THCudaIntTensor *idx_tensor) {
|
||||
|
||||
const float *points = THCudaTensor_data(state, points_tensor);
|
||||
float *temp = THCudaTensor_data(state, temp_tensor);
|
||||
int *idx = THCudaIntTensor_data(state, idx_tensor);
|
||||
|
||||
cudaStream_t stream = THCState_getCurrentStream(state);
|
||||
|
||||
furthest_point_sampling_kernel_wrapper(b, n, m, points, temp, idx,
|
||||
stream);
|
||||
return 1;
|
||||
}
|
||||
@@ -0,0 +1,216 @@
|
||||
#include <stdio.h>
|
||||
#include <stdlib.h>
|
||||
|
||||
#include "cuda_utils.h"
|
||||
#include "sampling_gpu.h"
|
||||
|
||||
// input: points(b, n, c) idx(b, m)
|
||||
// output: out(b, m, c)
|
||||
__global__ void gather_points_kernel(int b, int n, int c, int m,
|
||||
const float *__restrict__ points,
|
||||
const int *__restrict__ idx,
|
||||
float *__restrict__ out) {
|
||||
for (int i = blockIdx.x; i < b; i += gridDim.x) {
|
||||
for (int j = blockIdx.y * blockDim.x + threadIdx.x; j < m;
|
||||
j += blockDim.x * gridDim.y) {
|
||||
int a = idx[i * m + j];
|
||||
memcpy(out + (i * m + j) * c, points + (i * n + a) * c,
|
||||
sizeof(float) * c);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
void gather_points_kernel_wrapper(int b, int n, int c, int npoints,
|
||||
const float *points, const int *idx,
|
||||
float *out, cudaStream_t stream) {
|
||||
|
||||
cudaError_t err;
|
||||
gather_points_kernel<<<dim3(2, 8, 1), opt_n_threads(npoints) / 4, 0,
|
||||
stream>>>(b, n, c, npoints, points, idx, out);
|
||||
|
||||
err = cudaGetLastError();
|
||||
if (cudaSuccess != err) {
|
||||
fprintf(stderr, "CUDA kernel failed : %s\n", cudaGetErrorString(err));
|
||||
exit(-1);
|
||||
}
|
||||
}
|
||||
|
||||
__device__ void __update(float *__restrict__ dists, int *__restrict__ dists_i,
|
||||
int idx1, int idx2) {
|
||||
const float v1 = dists[idx1], v2 = dists[idx2];
|
||||
const int i1 = dists_i[idx1], i2 = dists_i[idx2];
|
||||
dists[idx1] = max(v1, v2);
|
||||
dists_i[idx1] = v2 > v1 ? i2 : i1;
|
||||
}
|
||||
|
||||
// Input dataset: (b, n, 3), tmp: (b, n)
|
||||
// Ouput idxs (b, m)
|
||||
template <unsigned int block_size>
|
||||
__global__ void furthest_point_sampling_kernel(
|
||||
int b, int n, int m, const float *__restrict__ dataset,
|
||||
float *__restrict__ temp, int *__restrict__ idxs) {
|
||||
if (m <= 0)
|
||||
return;
|
||||
__shared__ float dists[block_size];
|
||||
__shared__ int dists_i[block_size];
|
||||
|
||||
int batch_index = blockIdx.x;
|
||||
dataset += batch_index * n * 3;
|
||||
temp += batch_index * n;
|
||||
idxs += batch_index * m;
|
||||
|
||||
int tid = threadIdx.x;
|
||||
const int stride = block_size;
|
||||
|
||||
int old = 0;
|
||||
if (threadIdx.x == 0)
|
||||
idxs[0] = old;
|
||||
|
||||
__syncthreads();
|
||||
for (int j = 1; j < m; j++) {
|
||||
int besti = 0;
|
||||
float best = -1;
|
||||
float x1 = dataset[old * 3 + 0];
|
||||
float y1 = dataset[old * 3 + 1];
|
||||
float z1 = dataset[old * 3 + 2];
|
||||
for (int k = tid; k < n; k += stride) {
|
||||
float x2, y2, z2;
|
||||
x2 = dataset[k * 3 + 0];
|
||||
y2 = dataset[k * 3 + 1];
|
||||
z2 = dataset[k * 3 + 2];
|
||||
float mag = (x2 * x2) + (y2 * y2) + (z2 * z2);
|
||||
if (mag <= 1e-3)
|
||||
continue;
|
||||
|
||||
float d = (x2 - x1) * (x2 - x1) + (y2 - y1) * (y2 - y1) +
|
||||
(z2 - z1) * (z2 - z1);
|
||||
|
||||
float d2 = min(d, temp[k]);
|
||||
temp[k] = d2;
|
||||
besti = d2 > best ? k : besti;
|
||||
best = d2 > best ? d2 : best;
|
||||
}
|
||||
dists[tid] = best;
|
||||
dists_i[tid] = besti;
|
||||
__syncthreads();
|
||||
|
||||
if (block_size >= 512) {
|
||||
if (tid < 256) {
|
||||
__update(dists, dists_i, tid, tid + 256);
|
||||
}
|
||||
__syncthreads();
|
||||
}
|
||||
if (block_size >= 256) {
|
||||
if (tid < 128) {
|
||||
__update(dists, dists_i, tid, tid + 128);
|
||||
}
|
||||
__syncthreads();
|
||||
}
|
||||
if (block_size >= 128) {
|
||||
if (tid < 64) {
|
||||
__update(dists, dists_i, tid, tid + 64);
|
||||
}
|
||||
__syncthreads();
|
||||
}
|
||||
if (block_size >= 64) {
|
||||
if (tid < 32) {
|
||||
__update(dists, dists_i, tid, tid + 32);
|
||||
}
|
||||
__syncthreads();
|
||||
}
|
||||
if (block_size >= 32) {
|
||||
if (tid < 16) {
|
||||
__update(dists, dists_i, tid, tid + 16);
|
||||
}
|
||||
__syncthreads();
|
||||
}
|
||||
if (block_size >= 16) {
|
||||
if (tid < 8) {
|
||||
__update(dists, dists_i, tid, tid + 8);
|
||||
}
|
||||
__syncthreads();
|
||||
}
|
||||
if (block_size >= 8) {
|
||||
if (tid < 4) {
|
||||
__update(dists, dists_i, tid, tid + 4);
|
||||
}
|
||||
__syncthreads();
|
||||
}
|
||||
if (block_size >= 4) {
|
||||
if (tid < 2) {
|
||||
__update(dists, dists_i, tid, tid + 2);
|
||||
}
|
||||
__syncthreads();
|
||||
}
|
||||
if (block_size >= 2) {
|
||||
if (tid < 1) {
|
||||
__update(dists, dists_i, tid, tid + 1);
|
||||
}
|
||||
__syncthreads();
|
||||
}
|
||||
|
||||
old = dists_i[0];
|
||||
if (tid == 0)
|
||||
idxs[j] = old;
|
||||
}
|
||||
}
|
||||
|
||||
void furthest_point_sampling_kernel_wrapper(int b, int n, int m,
|
||||
const float *dataset, float *temp,
|
||||
int *idxs, cudaStream_t stream) {
|
||||
|
||||
cudaError_t err;
|
||||
unsigned int n_threads = opt_n_threads(n);
|
||||
|
||||
switch (n_threads) {
|
||||
case 512:
|
||||
furthest_point_sampling_kernel<512><<<b, n_threads, 0, stream>>>(
|
||||
b, n, m, dataset, temp, idxs);
|
||||
break;
|
||||
case 256:
|
||||
furthest_point_sampling_kernel<256><<<b, n_threads, 0, stream>>>(
|
||||
b, n, m, dataset, temp, idxs);
|
||||
break;
|
||||
case 128:
|
||||
furthest_point_sampling_kernel<128><<<b, n_threads, 0, stream>>>(
|
||||
b, n, m, dataset, temp, idxs);
|
||||
break;
|
||||
case 64:
|
||||
furthest_point_sampling_kernel<64><<<b, n_threads, 0, stream>>>(
|
||||
b, n, m, dataset, temp, idxs);
|
||||
break;
|
||||
case 32:
|
||||
furthest_point_sampling_kernel<32><<<b, n_threads, 0, stream>>>(
|
||||
b, n, m, dataset, temp, idxs);
|
||||
break;
|
||||
case 16:
|
||||
furthest_point_sampling_kernel<16><<<b, n_threads, 0, stream>>>(
|
||||
b, n, m, dataset, temp, idxs);
|
||||
break;
|
||||
case 8:
|
||||
furthest_point_sampling_kernel<8><<<b, n_threads, 0, stream>>>(
|
||||
b, n, m, dataset, temp, idxs);
|
||||
break;
|
||||
case 4:
|
||||
furthest_point_sampling_kernel<4><<<b, n_threads, 0, stream>>>(
|
||||
b, n, m, dataset, temp, idxs);
|
||||
break;
|
||||
case 2:
|
||||
furthest_point_sampling_kernel<2><<<b, n_threads, 0, stream>>>(
|
||||
b, n, m, dataset, temp, idxs);
|
||||
break;
|
||||
case 1:
|
||||
furthest_point_sampling_kernel<1><<<b, n_threads, 0, stream>>>(
|
||||
b, n, m, dataset, temp, idxs);
|
||||
break;
|
||||
default:
|
||||
furthest_point_sampling_kernel<512><<<b, n_threads, 0, stream>>>(
|
||||
b, n, m, dataset, temp, idxs);
|
||||
}
|
||||
|
||||
err = cudaGetLastError();
|
||||
if (cudaSuccess != err) {
|
||||
fprintf(stderr, "CUDA kernel failed : %s\n", cudaGetErrorString(err));
|
||||
exit(-1);
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user