Some faster/better kernels. Tensors with points are now kept in (b, c, ...) format as this is easier for pytorch

This commit is contained in:
erikwijmans
2018-02-10 20:32:52 -05:00
parent 8bce353da4
commit 65a127f3d2
19 changed files with 372 additions and 296 deletions
+18 -2
View File
@@ -4,7 +4,7 @@
extern THCState *state;
int gather_points_wrapper(int b, int n, int c, int npoints,
int gather_points_wrapper(int b, int c, int n, int npoints,
THCudaTensor *points_tensor,
THCudaIntTensor *idx_tensor,
THCudaTensor *out_tensor) {
@@ -15,7 +15,23 @@ int gather_points_wrapper(int b, int n, int c, int npoints,
cudaStream_t stream = THCState_getCurrentStream(state);
gather_points_kernel_wrapper(b, n, c, npoints, points, idx, out, stream);
gather_points_kernel_wrapper(b, c, n, npoints, points, idx, out, stream);
return 1;
}
int gather_points_grad_wrapper(int b, int c, int n, int npoints,
THCudaTensor *grad_out_tensor,
THCudaIntTensor *idx_tensor,
THCudaTensor *grad_points_tensor) {
const float *grad_out = THCudaTensor_data(state, grad_out_tensor);
const int *idx = THCudaIntTensor_data(state, idx_tensor);
float *grad_points = THCudaTensor_data(state, grad_points_tensor);
cudaStream_t stream = THCState_getCurrentStream(state);
gather_points_grad_kernel_wrapper(b, c, n, npoints, grad_out, idx,
grad_points, stream);
return 1;
}