Reconcile variants of ball_query

Author: nicolas-chaulet

cpu/src/bindings.cpp

@@ -9,41 +9,41 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
     m.def("ball_query",
 	      &ball_query,
 	      "compute the radius search of a point cloud using nanoflann"
-	      "-query : a pytorch tensor of size N1 x 3,. used to query the nearest neighbors"
-	      "- support : a pytorch tensor of size N2 x 3. used to build the tree"
-	      "-  radius : float number, size of the ball for the radius search."
+	      "- support : a pytorch tensor of size N1 x 3, points where the neighboors are accessed from"
+	      "- query : a pytorch tensor of size N2 x 3, centre of the balls"
+	      "- radius : float number, size of the ball for the radius search."
 	      "- max_num : int number, indicate the maximum of neaghbors allowed(if -1 then all the possible neighbors will be computed). "
-	      " - mode : int number that indicate which format for the neighborhood"
-	      " mode=0 mean a matrix of neighbors(-1 for shadow neighbors)"
+	      "- mode : int number that indicate which format for the neighborhood"
+	      "mode=0 mean a matrix of neighbors(-1 for shadow neighbors)"
 	      "mode=1 means a matrix of edges of size Num_edge x 2"
-	      "return a tensor of size N1 x M where M is either max_num or the maximum number of neighbors found if mode = 0, if mode=1 return a tensor of size Num_edge x 2 and return a tensor containing the squared distance of the neighbors",
-	      "query"_a, "support"_a, "radius"_a, "max_num"_a=-1, "mode"_a=0);
+	      "return a tensor of size N2 x M where M is either max_num or the maximum number of neighbors found if mode = 0, if mode=1 return a tensor of size Num_edge x 2 and return a tensor containing the squared distance of the neighbors",
+	      "support"_a, "querry"_a, "radius"_a, "max_num"_a=-1, "mode"_a=0);
 
     m.def("batch_ball_query",
 	      &batch_ball_query,
 		"compute the radius search of a point cloud for each batch using nanoflann"
-		"-query : a pytorch tensor (float) of size N1 x 3,. used to query the nearest neighbors"
-		"- support : a pytorch tensor(float) of size N2 x 3. used to build the tree"
-		"- query_batch : a pytorch tensor(long) contains indices of the batch of the query size N1"
-	      "NB : the batch must be sorted"
-		"- support_batch: a pytorch tensor(long) contains indices of the batch of the support size N2"
+		"- support : a pytorch tensor of size N1 x 3, points where the neighboors are accessed from"
+		"- query : a pytorch tensor of size N2 x 3, centre of the balls"
+		"- support_batch: a pytorch tensor(long) contains indices of the batch of the support size N1"
 	      "NB: the batch must be sorted"
+		"- query_batch : a pytorch tensor(long) contains indices of the batch of the query size N2"
+	      "NB : the batch must be sorted"
 		"-radius: float number, size of the ball for the radius search."
 		"- max_num : int number, indicate the maximum of neaghbors allowed(if -1 then all the possible neighbors wrt the radius will be computed)."
 		"- mode : int number that indicate which format for the neighborhood"
-		"mode=0 mean a matrix of neighbors(N2 for shadow neighbors)"
+		"mode=0 mean a matrix of neighbors(N1 for shadow neighbors)"
 		"mode=1 means a matrix of edges of size Num_edge x 2"
-	      "return a tensor of size N1 x M where M is either max_num or the maximum number of neighbors found if mode = 0, if mode=1 return a tensor of size Num_edge x 2 and return a tensor containing the squared distance of the neighbors",
-	      "query"_a, "support"_a, "query_batch"_a, "support_batch"_a, "radius"_a, "max_num"_a=-1, "mode"_a=0);
+	      "return a tensor of size N2 x M where M is either max_num or the maximum number of neighbors found if mode = 0, if mode=1 return a tensor of size Num_edge x 2 and return a tensor containing the squared distance of the neighbors",
+	      "support"_a, "querry"_a, "query_batch"_a, "support_batch"_a, "radius"_a, "max_num"_a=-1, "mode"_a=0);
     m.def("dense_ball_query", &dense_ball_query,
 	  "compute the radius search of a batch of point cloud using nanoflann"
-	      "-query : a pytorch tensor of size B x N1 x 3,. used to query the nearest neighbors"
-	      "- support : a pytorch tensor of size B x N2 x 3. used to build the tree"
-	      "-  radius : float number, size of the ball for the radius search."
+		  "- support : a pytorch tensor of size B x N1 x 3, points where the neighboors are accessed from"
+		  "- query : a pytorch tensor of size B x N2 x 3, centre of the balls"
+	      "- radius : float number, size of the ball for the radius search."
 	      "- max_num : int number, indicate the maximum of neaghbors allowed(if -1 then all the possible neighbors will be computed). "
-	      " - mode : int number that indicate which format for the neighborhood"
-	      " mode=0 mean a matrix of neighbors(-1 for shadow neighbors)"
+	      "- mode : int number that indicate which format for the neighborhood"
+	      "mode=0 mean a matrix of neighbors(-1 for shadow neighbors)"
 	      "mode=1 means a matrix of edges of size Num_edge x 2"
-	      "return a tensor of size N1 x M where M is either max_num or the maximum number of neighbors found if mode = 0, if mode=1 return a tensor of size Num_edge x 2 and return a tensor containing the squared distance of the neighbors",
-	      "query"_a, "support"_a, "radius"_a, "max_num"_a=-1, "mode"_a=0);
+	      "return a tensor of size B x N2 x M where M is either max_num or the maximum number of neighbors found if mode = 0, if mode=1 return a tensor of size Num_edge x 2 and return a tensor containing the squared distance of the neighbors",
+	      "support"_a, "querry"_a, "radius"_a, "max_num"_a=-1, "mode"_a=0);
 }

cpu/src/neighbors.cpp

@@ -23,7 +23,7 @@ int nanoflann_neighbors(vector<scalar_t>& queries,
 	int i0 = 0;
 
 	// Counting vector
-	int max_count = 1;
+	size_t max_count = 1;
 
 
 	// Nanoflann related variables
@@ -63,36 +63,30 @@ int nanoflann_neighbors(vector<scalar_t>& queries,
 		list_matches[i0].reserve(max_count);
 		std::vector<std::pair<size_t, scalar_t> >   ret_matches;
 
-
 		const size_t nMatches = index->radiusSearch(&query_pt[0], search_radius, ret_matches, search_params);
-		list_matches[i0] = ret_matches;
-		if((size_t)max_count < nMatches) max_count = nMatches;
+		if (nMatches == 0)
+			list_matches[i0] = {std::make_pair(0,-1)};
+		else
+			list_matches[i0] = ret_matches;
+		max_count = max(max_count,nMatches);
 		i0++;
-
-
 	}
 	// Reserve the memory
 	if(max_num > 0) {
 		max_count = max_num;
 	}
 	if(mode == 0){
-
-		neighbors_indices.resize(list_matches.size() * max_count);
-		dists.resize(list_matches.size() * max_count);
-
+		neighbors_indices.resize(list_matches.size() * max_count, 0);
+		dists.resize(list_matches.size() * max_count, -1);
 		i0 = 0;
-
 		int token = 0;
 		for (auto& inds : list_matches){
 			token = inds[0].first;
-			for (int j = 0; j < max_count; j++){
-				if ((unsigned int)j < inds.size()){
+			for (size_t j = 0; j < max_count; j++){
+				if (j < inds.size()){
 					neighbors_indices[i0 * max_count + j] = inds[j].first;
 					dists[i0 * max_count + j] = (float) inds[j].second;
-
-
 				}
-
 				else {
 					neighbors_indices[i0 * max_count + j] = token;
 					dists[i0 * max_count + j] = -1;
@@ -103,9 +97,9 @@ int nanoflann_neighbors(vector<scalar_t>& queries,
 
 	}
 	else if(mode == 1){
-		int size = 0; // total number of edges
+		size_t size = 0; // total number of edges
 		for (auto& inds : list_matches){
-			if((int)inds.size() <= max_count)
+			if(inds.size() <= max_count)
 				size += inds.size();
 			else
 				size += max_count;
@@ -115,8 +109,8 @@ int nanoflann_neighbors(vector<scalar_t>& queries,
 		int i0 = 0; // index of the query points
 		int u = 0; // curent index of the neighbors_indices
 		for (auto& inds : list_matches){
-			for (int j = 0; j < max_count; j++){
-				if((unsigned int)j < inds.size()){
+			for (size_t j = 0; j < max_count; j++){
+				if(j < inds.size()){
 					neighbors_indices[u] = inds[j].first;
 					neighbors_indices[u + 1] = i0;
 					dists[u/2] = (float) inds[j].second;
@@ -129,11 +123,8 @@ int nanoflann_neighbors(vector<scalar_t>& queries,
 
 	}
 	return max_count;
-
-
-
-
 }
+
 template<typename scalar_t>
 int batch_nanoflann_neighbors (vector<scalar_t>& queries,
                                vector<scalar_t>& supports,

cpu/src/torch_nearest_neighbors.cpp

@@ -7,39 +7,35 @@
 #include <iostream>
 
 
-std::pair<at::Tensor, at::Tensor> ball_query(at::Tensor query,
-					     at::Tensor support,
+std::pair<at::Tensor, at::Tensor> ball_query(at::Tensor support,
+					     at::Tensor query,
 					     float radius, int max_num, int mode){
 
 	at::Tensor out;
 	at::Tensor out_dists;
-	std::vector<long> neighbors_indices;
+	std::vector<long> neighbors_indices(query.size(0),0);
+	std::vector<float> neighbors_dists(query.size(0), -1);
 
 	auto options = torch::TensorOptions().dtype(torch::kLong).device(torch::kCPU);
 	auto options_dist = torch::TensorOptions().dtype(torch::kFloat32).device(torch::kCPU);
 	int max_count = 0;
 
-	std::vector<float> neighbors_dists;
-
 	AT_DISPATCH_ALL_TYPES(query.scalar_type(), "radius_search", [&] {
-
-
-	auto data_q = query.DATA_PTR<scalar_t>();
-	auto data_s = support.DATA_PTR<scalar_t>();
-	std::vector<scalar_t> queries_stl = std::vector<scalar_t>(data_q,
-								   data_q + query.size(0)*query.size(1));
-	std::vector<scalar_t> supports_stl = std::vector<scalar_t>(data_s,
-								   data_s + support.size(0)*support.size(1));
-
-	max_count = nanoflann_neighbors<scalar_t>(queries_stl,
-						  supports_stl,
-						  neighbors_indices,
-						  neighbors_dists,
-						  radius,
-						  max_num,
-						  mode);
+		auto data_q = query.DATA_PTR<scalar_t>();
+		auto data_s = support.DATA_PTR<scalar_t>();
+		std::vector<scalar_t> queries_stl = std::vector<scalar_t>(data_q,
+									data_q + query.size(0)*query.size(1));
+		std::vector<scalar_t> supports_stl = std::vector<scalar_t>(data_s,
+									data_s + support.size(0)*support.size(1));
+
+		max_count = nanoflann_neighbors<scalar_t>(queries_stl,
+							supports_stl,
+							neighbors_indices,
+							neighbors_dists,
+							radius,
+							max_num,
+							mode);
 	});
-
 	auto neighbors_dists_ptr = neighbors_dists.data();
 	long* neighbors_indices_ptr = neighbors_indices.data();
 	if(mode == 0){
@@ -65,10 +61,10 @@ at::Tensor degree(at::Tensor row, int64_t num_nodes) {
 	return zero.scatter_add_(0, row, one);
 }
 
-std::pair<at::Tensor, at::Tensor> batch_ball_query(at::Tensor query,
-						   at::Tensor support,
-						   at::Tensor query_batch,
+std::pair<at::Tensor, at::Tensor> batch_ball_query(at::Tensor support,
+						   at::Tensor query,
 						   at::Tensor support_batch,
+						   at::Tensor query_batch,
 						   float radius, int max_num, int mode) {
 	at::Tensor idx;
 
@@ -92,8 +88,7 @@ std::pair<at::Tensor, at::Tensor> batch_ball_query(at::Tensor query,
 	AT_DISPATCH_ALL_TYPES(query.scalar_type(), "batch_radius_search", [&] {
 
         std::vector<scalar_t> queries_stl(query.DATA_PTR<scalar_t>(), query.DATA_PTR<scalar_t>() + query.numel());
-	std::vector<scalar_t> supports_stl(support.DATA_PTR<scalar_t>(), support.DATA_PTR<scalar_t>() + support.numel());
-
+		std::vector<scalar_t> supports_stl(support.DATA_PTR<scalar_t>(), support.DATA_PTR<scalar_t>() + support.numel());
 
         max_count = batch_nanoflann_neighbors<scalar_t>(queries_stl,
 							supports_stl,
@@ -114,7 +109,6 @@ std::pair<at::Tensor, at::Tensor> batch_ball_query(at::Tensor query,
 		dist = torch::from_blob(neighbors_dists_ptr,
 					{query.size(0), max_count},
 					options=options_dist);
-
 	}
 	else if(mode ==1){
 		idx = torch::from_blob(neighbors_indices_ptr, {(int)neighbors_indices.size()/2, 2}, options=options);
@@ -123,12 +117,11 @@ std::pair<at::Tensor, at::Tensor> batch_ball_query(at::Tensor query,
 					options=options_dist);
 	}
 	return std::make_pair(idx.clone(), dist.clone());
-
 }
 
 
-std::pair<at::Tensor, at::Tensor> dense_ball_query(at::Tensor query,
-						   at::Tensor support,
+std::pair<at::Tensor, at::Tensor> dense_ball_query(at::Tensor support,
+						   at::Tensor query,
 						   float radius, int max_num, int mode){
 
 	int b = query.size(0);

cuda/include/compat.h

@@ -0,0 +1,5 @@
+#ifdef VERSION_GE_1_3
+#define DATA_PTR data_ptr
+#else
+#define DATA_PTR data
+#endif

cuda/src/ball_query.cpp

@@ -1,5 +1,6 @@
 #include "ball_query.h"
 #include "utils.h"
+#include "compat.h"
 
 void query_ball_point_kernel_dense_wrapper(int b, int n, int m, float radius,
 					   int nsample, const float *new_xyz,
@@ -33,8 +34,8 @@ at::Tensor ball_query_dense(at::Tensor new_xyz, at::Tensor xyz, const float radi
 
   if (new_xyz.type().is_cuda()) {
     query_ball_point_kernel_dense_wrapper(xyz.size(0), xyz.size(1), new_xyz.size(1),
-					  radius, nsample, new_xyz.data_ptr<float>(),
-					  xyz.data_ptr<float>(), idx.data_ptr<int>());
+					  radius, nsample, new_xyz.DATA_PTR<float>(),
+					  xyz.DATA_PTR<float>(), idx.DATA_PTR<int>());
   } else {
     TORCH_CHECK(false, "CPU not supported");
   }
@@ -74,7 +75,7 @@ std::pair<at::Tensor, at::Tensor> ball_query_partial_dense(at::Tensor x,
 
 	cudaSetDevice(x.get_device());
 	auto batch_sizes = (int64_t *)malloc(sizeof(int64_t));
-	cudaMemcpy(batch_sizes, batch_x[-1].data_ptr<int64_t>(), sizeof(int64_t),
+	cudaMemcpy(batch_sizes, batch_x[-1].DATA_PTR<int64_t>(), sizeof(int64_t),
 				cudaMemcpyDeviceToHost);
 	auto batch_size = batch_sizes[0] + 1;
 
@@ -88,12 +89,12 @@ std::pair<at::Tensor, at::Tensor> ball_query_partial_dense(at::Tensor x,
 							x.size(0),
 							y.size(0),
 							radius, nsample,
-							x.data_ptr<float>(),
-							y.data_ptr<float>(),
-							batch_x.data_ptr<long>(),
-							batch_y.data_ptr<long>(),
-							idx.data_ptr<long>(),
-							dist.data_ptr<float>());
+							x.DATA_PTR<float>(),
+							y.DATA_PTR<float>(),
+							batch_x.DATA_PTR<long>(),
+							batch_y.DATA_PTR<long>(),
+							idx.DATA_PTR<long>(),
+							dist.DATA_PTR<float>());
 	} else {
 	  TORCH_CHECK(false, "CPU not supported");
 	}

cuda/src/group_points.cpp

@@ -1,5 +1,6 @@
 #include "group_points.h"
 #include "utils.h"
+#include "compat.h"
 
 void group_points_kernel_wrapper(int b, int c, int n, int npoints, int nsample,
                                  const float *points, const int *idx,
@@ -25,8 +26,8 @@ at::Tensor group_points(at::Tensor points, at::Tensor idx) {
 
   if (points.type().is_cuda()) {
     group_points_kernel_wrapper(points.size(0), points.size(1), points.size(2),
-                                idx.size(1), idx.size(2), points.data_ptr<float>(),
-                                idx.data_ptr<int>(), output.data_ptr<float>());
+                                idx.size(1), idx.size(2), points.DATA_PTR<float>(),
+                                idx.DATA_PTR<int>(), output.DATA_PTR<float>());
   } else {
     TORCH_CHECK(false, "CPU not supported");
   }
@@ -51,7 +52,7 @@ at::Tensor group_points_grad(at::Tensor grad_out, at::Tensor idx, const int n) {
   if (grad_out.type().is_cuda()) {
     group_points_grad_kernel_wrapper(
         grad_out.size(0), grad_out.size(1), n, idx.size(1), idx.size(2),
-        grad_out.data_ptr<float>(), idx.data_ptr<int>(), output.data_ptr<float>());
+        grad_out.DATA_PTR<float>(), idx.DATA_PTR<int>(), output.DATA_PTR<float>());
   } else {
     TORCH_CHECK(false, "CPU not supported");
   }

cuda/src/interpolate.cpp

@@ -1,5 +1,6 @@
 #include "interpolate.h"
 #include "utils.h"
+#include "compat.h"
 
 void three_nn_kernel_wrapper(int b, int n, int m, const float *unknown,
                              const float *known, float *dist2, int *idx);
@@ -30,8 +31,8 @@ std::vector<at::Tensor> three_nn(at::Tensor unknowns, at::Tensor knows) {
 
   if (unknowns.type().is_cuda()) {
     three_nn_kernel_wrapper(unknowns.size(0), unknowns.size(1), knows.size(1),
-                            unknowns.data_ptr<float>(), knows.data_ptr<float>(),
-                            dist2.data_ptr<float>(), idx.data_ptr<int>());
+                            unknowns.DATA_PTR<float>(), knows.DATA_PTR<float>(),
+                            dist2.DATA_PTR<float>(), idx.DATA_PTR<int>());
   } else {
     TORCH_CHECK(false, "CPU not supported");
   }
@@ -60,8 +61,8 @@ at::Tensor three_interpolate(at::Tensor points, at::Tensor idx,
   if (points.type().is_cuda()) {
     three_interpolate_kernel_wrapper(
         points.size(0), points.size(1), points.size(2), idx.size(1),
-        points.data_ptr<float>(), idx.data_ptr<int>(), weight.data_ptr<float>(),
-        output.data_ptr<float>());
+        points.DATA_PTR<float>(), idx.DATA_PTR<int>(), weight.DATA_PTR<float>(),
+        output.DATA_PTR<float>());
   } else {
     TORCH_CHECK(false, "CPU not supported");
   }
@@ -89,8 +90,8 @@ at::Tensor three_interpolate_grad(at::Tensor grad_out, at::Tensor idx,
   if (grad_out.type().is_cuda()) {
     three_interpolate_grad_kernel_wrapper(
         grad_out.size(0), grad_out.size(1), grad_out.size(2), m,
-        grad_out.data_ptr<float>(), idx.data_ptr<int>(), weight.data_ptr<float>(),
-        output.data_ptr<float>());
+        grad_out.DATA_PTR<float>(), idx.DATA_PTR<int>(), weight.DATA_PTR<float>(),
+        output.DATA_PTR<float>());
   } else {
     TORCH_CHECK(false, "CPU not supported");
   }