Fix KNN query capacity for decoder interpolation

modules/pointops/src/knnquery/knnquery_cuda.cpp

@@ -1,12 +1,20 @@
 #include <vector>
 #include <THC/THC.h>
+#include <torch/extension.h>
 #include <torch/serialize/tensor.h>
 #include <ATen/cuda/CUDAContext.h>
 #include "knnquery_cuda_kernel.h"
 
 
 void knnquery_cuda(int m, int nsample, at::Tensor xyz_tensor, at::Tensor new_xyz_tensor, at::Tensor offset_tensor, at::Tensor new_offset_tensor, at::Tensor idx_tensor, at::Tensor dist2_tensor)
 {
+    TORCH_CHECK(
+        nsample >= 1 && nsample <= KNNQUERY_MAX_NEIGHBORS,
+        "knnquery nsample must be between 1 and ",
+        KNNQUERY_MAX_NEIGHBORS,
+        "; got ",
+        nsample
+    );
     const float *xyz = xyz_tensor.data_ptr<float>();
     const float *new_xyz = new_xyz_tensor.data_ptr<float>();
     const int *offset = offset_tensor.data_ptr<int>();

modules/pointops/src/knnquery/knnquery_cuda_kernel.cu

@@ -83,8 +83,8 @@ __global__ void knnquery_cuda_kernel(int m, int nsample, const float *__restrict
     float new_y = new_xyz[1];
     float new_z = new_xyz[2];
 
-    float best_dist[100];
-    int best_idx[100];
+    float best_dist[KNNQUERY_MAX_NEIGHBORS];
+    int best_idx[KNNQUERY_MAX_NEIGHBORS];
     for(int i = 0; i < nsample; i++){
         best_dist[i] = 1e10;
         best_idx[i] = start;

modules/pointops/src/knnquery/knnquery_cuda_kernel.h

@@ -4,6 +4,8 @@
 #include <torch/serialize/tensor.h>
 #include <ATen/cuda/CUDAContext.h>
 
+constexpr int KNNQUERY_MAX_NEIGHBORS = 256;
+
 void knnquery_cuda(int m, int nsample, at::Tensor xyz_tensor, at::Tensor new_xyz_tensor, at::Tensor offset_tensor, at::Tensor new_offset_tensor, at::Tensor idx_tensor, at::Tensor dist2_tensor);
 
 #ifdef __cplusplus

tests/test_knnquery_capacity.py

@@ -0,0 +1,57 @@
+import unittest
+
+import torch
+
+from modules.pointops.functions import pointops
+
+
+def make_points(count):
+    x = torch.linspace(0.0, 1.0, count, device="cuda", dtype=torch.float32)
+    zeros = torch.zeros_like(x)
+    xyz = torch.stack((x, zeros, zeros), dim=1).contiguous()
+    offset = torch.cuda.IntTensor([count])
+    return xyz, offset
+
+
+@unittest.skipUnless(torch.cuda.is_available(), "CUDA is required")
+class KNNQueryCapacityTest(unittest.TestCase):
+    def test_supports_256_neighbors_without_invalid_indexes(self):
+        xyz, offset = make_points(256)
+        new_xyz = xyz[:1].contiguous()
+        new_offset = torch.cuda.IntTensor([1])
+
+        indexes, distances = pointops.knnquery(
+            256, xyz, new_xyz, offset, new_offset
+        )
+        torch.cuda.synchronize()
+
+        self.assertEqual(tuple(indexes.shape), (1, 256))
+        self.assertEqual(tuple(distances.shape), (1, 256))
+        self.assertGreaterEqual(indexes.min().item(), 0)
+        self.assertLess(indexes.max().item(), xyz.shape[0])
+
+        features = torch.arange(
+            256, device="cuda", dtype=torch.float32
+        ).reshape(256, 1).contiguous()
+        output = pointops.interpolation(
+            xyz, new_xyz, features, offset, new_offset, k=256
+        )
+        torch.cuda.synchronize()
+
+        self.assertEqual(tuple(output.shape), (1, 1))
+        self.assertTrue(torch.isfinite(output).all().item())
+
+    def test_rejects_neighbor_count_above_capacity(self):
+        xyz, offset = make_points(257)
+        new_xyz = xyz[:1].contiguous()
+        new_offset = torch.cuda.IntTensor([1])
+
+        with self.assertRaisesRegex(
+            RuntimeError,
+            r"knnquery nsample must be between 1 and 256; got 257",
+        ):
+            pointops.knnquery(257, xyz, new_xyz, offset, new_offset)
+
+
+if __name__ == "__main__":
+    unittest.main()