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

utils/csrc/group_points.c

@@ -4,7 +4,7 @@
 
 extern THCState *state;
 
-int group_points_wrapper(int b, int n, int c, int npoints, int nsample,
+int group_points_wrapper(int b, int c, int n, int npoints, int nsample,
 			 THCudaTensor *points_tensor,
 			 THCudaIntTensor *idx_tensor,
 			 THCudaTensor *out_tensor) {
@@ -15,12 +15,12 @@ int group_points_wrapper(int b, int n, int c, int npoints, int nsample,
 
     cudaStream_t stream = THCState_getCurrentStream(state);
 
-    group_points_kernel_wrapper(b, n, c, npoints, nsample, points, idx, out,
+    group_points_kernel_wrapper(b, c, n, npoints, nsample, points, idx, out,
 				stream);
     return 1;
 }
 
-int group_points_grad_wrapper(int b, int n, int c, int npoints, int nsample,
+int group_points_grad_wrapper(int b, int c, int n, int npoints, int nsample,
 			      THCudaTensor *grad_out_tensor,
 			      THCudaIntTensor *idx_tensor,
 			      THCudaTensor *grad_points_tensor) {
@@ -31,7 +31,7 @@ int group_points_grad_wrapper(int b, int n, int c, int npoints, int nsample,
 
     cudaStream_t stream = THCState_getCurrentStream(state);
 
-    group_points_grad_kernel_wrapper(b, n, c, npoints, nsample, grad_out, idx,
+    group_points_grad_kernel_wrapper(b, c, n, npoints, nsample, grad_out, idx,
 				     grad_points, stream);
     return 1;
 }

utils/csrc/group_points_gpu.cu

@@ -4,9 +4,9 @@
 #include "cuda_utils.h"
 #include "group_points_gpu.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,
+// input: points(b, c, n) idx(b, npoints, nsample)
+// output: out(b, c, npoints, nsample)
+__global__ void group_points_kernel(int b, int c, int n, int npoints,
 				    int nsample,
 				    const float *__restrict__ points,
 				    const int *__restrict__ idx,
@@ -16,25 +16,25 @@ __global__ void group_points_kernel(int b, int n, int c, int npoints,
     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) {
+    const int index = threadIdx.y * blockDim.x + threadIdx.x;
+    const int stride = blockDim.y * blockDim.x;
+    for (int i = index; i < c * npoints; i += stride) {
+	const int l = i / npoints;
+	const int j = i % npoints;
 	for (int k = 0; k < nsample; ++k) {
 	    int ii = idx[j * nsample + k];
-	    for (int l = 0; l < c; ++l) {
-		out[j * nsample * c + k * c + l] = points[ii * c + l];
-	    }
+	    out[(l * npoints + j) * nsample + k] = points[l * n + ii];
 	}
     }
 }
 
-void group_points_kernel_wrapper(int b, int n, int c, int npoints, int nsample,
+void group_points_kernel_wrapper(int b, int c, int n, 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);
+    group_points_kernel<<<b, opt_block_config(npoints, c), 0, stream>>>(
+	b, c, n, npoints, nsample, points, idx, out);
 
     err = cudaGetLastError();
     if (cudaSuccess != err) {
@@ -43,38 +43,38 @@ void group_points_kernel_wrapper(int b, int n, int c, int npoints, int nsample,
     }
 }
 
-// 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,
+// input: grad_out(b, c, npoints, nsample), idx(b, npoints, nsample)
+// output: grad_points(b, c, n)
+__global__ void group_points_grad_kernel(int b, int c, int n, 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;
+    idx += batch_index * npoints * nsample;
+    grad_points += batch_index * n * c;
 
-    int index = threadIdx.x;
-    int stride = blockDim.x;
-    for (int j = index; j < npoints; j += stride) {
+    const int index = threadIdx.y * blockDim.x + threadIdx.x;
+    const int stride = blockDim.y * blockDim.x;
+    for (int i = index; i < c * npoints; i += stride) {
+	const int l = i / npoints;
+	const int j = i % npoints;
 	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]);
-	    }
+	    atomicAdd(grad_points + l * n + ii,
+		      grad_out[(l * npoints + j) * nsample + k]);
 	}
     }
 }
 
-void group_points_grad_kernel_wrapper(int b, int n, int c, int npoints,
+void group_points_grad_kernel_wrapper(int b, int c, int n, 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);
+    group_points_grad_kernel<<<b, opt_block_config(npoints, c), 0, stream>>>(
+	b, c, n, npoints, nsample, grad_out, idx, grad_points);
 
     err = cudaGetLastError();
     if (cudaSuccess != err) {

utils/csrc/interpolate.c

@@ -19,7 +19,7 @@ void three_nn_wrapper(int b, int n, int m, THCudaTensor *unknown_tensor,
     three_nn_kernel_wrapper(b, n, m, unknown, known, dist2, idx, stream);
 }
 
-void three_interpolate_wrapper(int b, int m, int c, int n,
+void three_interpolate_wrapper(int b, int c, int m, int n,
 			       THCudaTensor *points_tensor,
 			       THCudaIntTensor *idx_tensor,
 			       THCudaTensor *weight_tensor,
@@ -31,11 +31,11 @@ void three_interpolate_wrapper(int b, int m, int c, int n,
     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,
+    three_interpolate_kernel_wrapper(b, c, m, n, points, idx, weight, out,
 				     stream);
 }
 
-void three_interpolate_grad_wrapper(int b, int n, int c, int m,
+void three_interpolate_grad_wrapper(int b, int c, int n, int m,
 				    THCudaTensor *grad_out_tensor,
 				    THCudaIntTensor *idx_tensor,
 				    THCudaTensor *weight_tensor,
@@ -47,6 +47,6 @@ void three_interpolate_grad_wrapper(int b, int n, int c, int m,
     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,
+    three_interpolate_grad_kernel_wrapper(b, c, n, m, grad_out, idx, weight,
 					  grad_points, stream);
 }

utils/csrc/interpolate_gpu.cu

@@ -77,9 +77,9 @@ void three_nn_kernel_wrapper(int b, int n, int m, const float *unknown,
     }
 }
 
-// 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,
+// input: points(b, c, m), idx(b, n, 3), weight(b, n, 3)
+// output: out(b, c, n)
+__global__ void three_interpolate_kernel(int b, int c, int m, int n,
 					 const float *__restrict__ points,
 					 const int *__restrict__ idx,
 					 const float *__restrict__ weight,
@@ -92,9 +92,11 @@ __global__ void three_interpolate_kernel(int b, int m, int c, int n,
 
     out += batch_index * n * c;
 
-    int index = threadIdx.x;
-    int stride = blockDim.x;
-    for (int j = index; j < n; j += stride) {
+    const int index = threadIdx.y * blockDim.x + threadIdx.x;
+    const int stride = blockDim.y * blockDim.x;
+    for (int i = index; i < c * n; i += stride) {
+	const int l = i / n;
+	const int j = i % n;
 	float w1 = weight[j * 3 + 0];
 	float w2 = weight[j * 3 + 1];
 	float w3 = weight[j * 3 + 2];
@@ -103,21 +105,19 @@ __global__ void three_interpolate_kernel(int b, int m, int c, int n,
 	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;
-	}
+	out[i] = points[l * m + i1] * w1 + points[l * m + i2] * w2 +
+		 points[l * m + i3] * w3;
     }
 }
 
-void three_interpolate_kernel_wrapper(int b, int m, int c, int n,
+void three_interpolate_kernel_wrapper(int b, int c, int m, 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);
+    three_interpolate_kernel<<<b, opt_block_config(n, c), 0, stream>>>(
+	b, c, m, n, points, idx, weight, out);
 
     err = cudaGetLastError();
     if (cudaSuccess != err) {
@@ -128,11 +128,11 @@ void three_interpolate_kernel_wrapper(int b, int m, int c, int n,
     }
 }
 
-// input: grad_out(b, n, c), idx(b, n, 3), weight(b, n, 3)
-// output: grad_points(b, m, c)
+// input: grad_out(b, c, n), idx(b, n, 3), weight(b, n, 3)
+// output: grad_points(b, c, m)
 
 __global__ void three_interpolate_grad_kernel(
-    int b, int n, int c, int m, const float *__restrict__ grad_out,
+    int b, int c, int n, int m, const float *__restrict__ grad_out,
     const int *__restrict__ idx, const float *__restrict__ weight,
     float *__restrict__ grad_points) {
     int batch_index = blockIdx.x;
@@ -141,9 +141,11 @@ __global__ void three_interpolate_grad_kernel(
     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) {
+    const int index = threadIdx.y * blockDim.x + threadIdx.x;
+    const int stride = blockDim.y * blockDim.x;
+    for (int i = index; i < c * n; i += stride) {
+	const int l = i / n;
+	const int j = i % n;
 	float w1 = weight[j * 3 + 0];
 	float w2 = weight[j * 3 + 1];
 	float w3 = weight[j * 3 + 2];
@@ -152,11 +154,9 @@ __global__ void three_interpolate_grad_kernel(
 	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);
-	}
+	atomicAdd(grad_points + l * m + i1, grad_out[i] * w1);
+	atomicAdd(grad_points + l * m + i2, grad_out[i] * w2);
+	atomicAdd(grad_points + l * m + i3, grad_out[i] * w3);
     }
 }
 
@@ -167,7 +167,7 @@ void three_interpolate_grad_kernel_wrapper(int b, int n, int c, int m,
 					   cudaStream_t stream) {
 
     cudaError_t err;
-    three_interpolate_grad_kernel<<<b, opt_n_threads(n) / 4, 0, stream>>>(
+    three_interpolate_grad_kernel<<<b, opt_block_config(n, c), 0, stream>>>(
 	b, n, c, m, grad_out, idx, weight, grad_points);
 
     err = cudaGetLastError();

utils/csrc/sampling.c

@@ -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;
 }
 

utils/csrc/sampling_gpu.cu

@@ -4,30 +4,63 @@
 #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,
+// input: points(b, c, n) idx(b, m)
+// output: out(b, c, m)
+__global__ void gather_points_kernel(int b, int c, int n, 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) {
-	    const int jj = idx[i * m + j];
-	    for (int l = 0; l < c; ++l) {
-		out[(i * m + j) * c + l] = points[(i * n + jj) * c + l];
+	for (int l = blockIdx.y; l < c; l += gridDim.y) {
+	    for (int j = threadIdx.x; j < m; j += blockDim.x) {
+		int a = idx[i * m + j];
+		out[(i * c + l) * m + j] = points[(i * c + l) * n + a];
 	    }
 	}
     }
 }
 
-void gather_points_kernel_wrapper(int b, int n, int c, int npoints,
+void gather_points_kernel_wrapper(int b, int c, int n, int npoints,
 				  const float *points, const int *idx,
 				  float *out, cudaStream_t stream) {
 
     cudaError_t err;
-    gather_points_kernel<<<dim3(b, 8, 1), opt_n_threads(npoints), 0,
-			   stream>>>(b, n, c, npoints, points, idx, out);
+    gather_points_kernel<<<dim3(b, c, 1), opt_n_threads(npoints), 0, stream>>>(
+	b, c, n, npoints, points, idx, out);
+
+    err = cudaGetLastError();
+    if (cudaSuccess != err) {
+	fprintf(stderr, "CUDA kernel failed : %s\n", cudaGetErrorString(err));
+	exit(-1);
+    }
+}
+
+// input: grad_out(b, c, m) idx(b, m)
+// output: grad_points(b, c, n)
+__global__ void gather_points_grad_kernel(int b, int c, int n, int m,
+					  const float *__restrict__ grad_out,
+					  const int *__restrict__ idx,
+					  float *__restrict__ grad_points) {
+    for (int i = blockIdx.x; i < b; i += gridDim.x) {
+	for (int l = blockIdx.y; l < c; l += gridDim.y) {
+	    for (int j = threadIdx.x; j < m; j += blockDim.x) {
+		int a = idx[i * m + j];
+		atomicAdd(grad_points + (i * c + l) * n + a,
+			  grad_out[(i * c + l) * m + j]);
+	    }
+	}
+    }
+}
+
+void gather_points_grad_kernel_wrapper(int b, int c, int n, int npoints,
+				       const float *grad_out, const int *idx,
+				       float *grad_points,
+				       cudaStream_t stream) {
+
+    cudaError_t err;
+    gather_points_grad_kernel<<<dim3(b, c, 1), opt_n_threads(npoints), 0,
+				stream>>>(b, c, n, npoints, grad_out, idx,
+					  grad_points);
 
     err = cudaGetLastError();
     if (cudaSuccess != err) {