Work around the prange non-resolution

prange is not resolved when used in a cypclass.
Co-authored-by: Xavier Thompson <xavier.thompson@nexedi.com>

kdtree.pyx

@@ -515,6 +515,24 @@ cdef cypclass Node activable:
             counter=counter,
         )
 
+# TODO (jjerphan): integrate this where needed once the
+#       compilation problem with Cython has been resolved
+cdef void prange_workaround(
+    KDTree tree,
+    NeighborsHeaps heaps,
+    I_t n_query,
+    I_t num_threads,
+    I_t n_features,
+    D_t * _query_points_ptr,
+) nogil:
+    cdef:
+        D_t rdist_lower_bound = 0
+        I_t i
+
+    for i in prange(n_query, schedule='static', num_threads=num_threads):
+        rdist_lower_bound = tree.min_rdist(0, _query_points_ptr + i * n_features)
+        tree._query_single_depthfirst(0, _query_points_ptr, i, heaps, rdist_lower_bound)
+
 
 cdef cypclass KDTree:
     """A KDTree based on asynchronous and parallel computations.
@@ -699,9 +717,7 @@ cdef cypclass KDTree:
                 n_neighbors
             )
 
-        for i in prange(n_query, schedule='static', num_threads=num_threads):
-            rdist_lower_bound = self.min_rdist(0, _query_points_ptr + i * n_features)
-            self._query_single_depthfirst(0, _query_points_ptr, i, heaps, rdist_lower_bound)
+        prange_workaround(self, heaps, n_query, num_threads, n_features, _query_points_ptr)
 
         heaps.sort()
 

[Julien Jerphanion]

This compiles but blocks at runtime. Might this problem be due to improper interactions between OpenMP threads and Cython+ runtime threads?

A similar unresolved issue exists for OpenMP: https://github.com/xianyi/OpenBLAS/issues/3187

---

Script to reproduce the problem:

import kdtree
import numpy as np

if __name__ == "__main__":
    X = np.load("X.npy")
    tree = kdtree.KDTree(X, leaf_size=256)

    n_query = X.shape[0] // 10
    k = 1
    Y = np.random.rand(n_query, X.shape[1])

    knn_indices = np.zeros((n_query, k), dtype=np.int32)
    knn_distances = np.zeros((n_query, k), dtype=np.float64)
    tree.query(Y, knn_indices, knn_distances)