feat: add multithreading to batch raycasting

embreex/rtcore_scene.pyx

@@ -11,8 +11,20 @@ from . cimport rtcore_ray as rtcr
 from . cimport rtcore_geometry as rtcg
 
 
+cdef extern from "tbb_parallel.h" nogil:
+    ctypedef void (*embreex_range_fn)(void* ctx, size_t begin, size_t end) noexcept nogil
+
+    void embreex_parallel_for(size_t n, size_t grain, int nthreads,
+                              embreex_range_fn fn, void* ctx) except +
+
+
 log = logging.getLogger('embreex')
 
+# TBB range divisibility threshold, in rays. Value taken from Open3D; not
+# profiled here.
+cdef size_t _TBB_GRAIN_SIZE = 1024
+
+
 cdef void error_printer(void* userPtr, const rtc.RTCError code, const char *_str) noexcept:
     """
     error_printer function for Embree 4.x
@@ -22,6 +34,98 @@ cdef void error_printer(void* userPtr, const rtc.RTCError code, const char *_str
     log.error("ERROR MESSAGE: %s" % _str)
 
 
+# Raw buffers and byte strides let TBB workers traverse strided inputs without
+# accessing Python objects; outputs are contiguous.
+cdef struct RayJob:
+    RTCScene scene
+    char* org
+    Py_ssize_t org_s0
+    Py_ssize_t org_s1
+    char* dir
+    Py_ssize_t dir_s0
+    Py_ssize_t dir_s1
+    int direction_row_step
+    char* tfar
+    Py_ssize_t tfar_s0
+    int* intersect_ids
+    float* u_arr
+    float* v_arr
+    float* Ng_arr
+    int* primID_arr
+    int* geomID_arr
+    int query_type
+    bint use_output
+
+
+cdef void _cast_range(void* ctx, size_t begin, size_t end) noexcept nogil:
+    """Trace rays in the half-open range [begin, end)."""
+    cdef RayJob* j = <RayJob*>ctx
+    cdef rtcr.RTCRayHit rayhit
+    cdef unsigned int INVALID_GEOMETRY_ID = 0xFFFFFFFF
+    cdef size_t i
+    cdef Py_ssize_t vd_i
+    cdef char* origin_ptr
+    cdef char* direction_ptr
+    cdef char* distance_ptr
+
+    for i in range(begin, end):
+        origin_ptr = j.org + <Py_ssize_t>i * j.org_s0
+        # Broadcast a single direction row across all origins.
+        vd_i = <Py_ssize_t>i * j.direction_row_step
+        direction_ptr = j.dir + vd_i * j.dir_s0
+        distance_ptr = j.tfar + <Py_ssize_t>i * j.tfar_s0
+
+        rayhit.ray.org_x = (<float*>origin_ptr)[0]
+        rayhit.ray.org_y = (<float*>(origin_ptr + j.org_s1))[0]
+        rayhit.ray.org_z = (<float*>(origin_ptr + 2 * j.org_s1))[0]
+        rayhit.ray.dir_x = (<float*>direction_ptr)[0]
+        rayhit.ray.dir_y = (<float*>(direction_ptr + j.dir_s1))[0]
+        rayhit.ray.dir_z = (<float*>(direction_ptr + 2 * j.dir_s1))[0]
+        rayhit.ray.tnear = 0.0
+        rayhit.ray.tfar = (<float*>distance_ptr)[0]
+        rayhit.hit.geomID = INVALID_GEOMETRY_ID
+        rayhit.hit.primID = INVALID_GEOMETRY_ID
+        rayhit.hit.instID[0] = INVALID_GEOMETRY_ID
+        rayhit.ray.mask = 0xFFFFFFFF
+        rayhit.ray.time = 0.0
+        rayhit.ray.flags = 0
+
+        if j.query_type == <int>intersect or j.query_type == <int>distance:
+            rtcIntersect1(j.scene, &rayhit, NULL)
+            if not j.use_output:
+                if j.query_type == <int>intersect:
+                    j.intersect_ids[i] = -1 if rayhit.hit.primID == INVALID_GEOMETRY_ID else <int>rayhit.hit.primID
+                else:
+                    (<float*>distance_ptr)[0] = rayhit.ray.tfar
+            else:
+                j.primID_arr[i] = -1 if rayhit.hit.primID == INVALID_GEOMETRY_ID else <int>rayhit.hit.primID
+                j.geomID_arr[i] = -1 if rayhit.hit.geomID == INVALID_GEOMETRY_ID else <int>rayhit.hit.geomID
+                j.u_arr[i] = rayhit.hit.u
+                j.v_arr[i] = rayhit.hit.v
+                (<float*>distance_ptr)[0] = rayhit.ray.tfar
+                j.Ng_arr[3 * i + 0] = rayhit.hit.Ng_x
+                j.Ng_arr[3 * i + 1] = rayhit.hit.Ng_y
+                j.Ng_arr[3 * i + 2] = rayhit.hit.Ng_z
+        else:
+            rtcOccluded1(j.scene, &rayhit.ray, NULL)
+            # In Embree 4, occlusion is signaled by setting ray.tfar to -inf
+            j.intersect_ids[i] = 0 if rayhit.ray.tfar < 0 else -1
+
+
+cdef bint _ray_buffers_are_thread_safe(int nv,
+                                       Py_ssize_t org_s0, Py_ssize_t dir_s0,
+                                       Py_ssize_t tfar_s0, bint user_dists) nogil:
+    if nv <= 1:
+        return True
+    if org_s0 < <Py_ssize_t>sizeof(float):
+        return False
+    if dir_s0 < <Py_ssize_t>sizeof(float):
+        return False
+    if user_dists and tfar_s0 < <Py_ssize_t>sizeof(float):
+        return False
+    return True
+
+
 cdef class EmbreeScene:
     def __init__(self, rtc.EmbreeDevice device=None, robust=True):
         if device is None:
@@ -39,28 +143,29 @@ cdef class EmbreeScene:
 
     def run(self, np.ndarray[np.float32_t, ndim=2] vec_origins,
                   np.ndarray[np.float32_t, ndim=2] vec_directions,
-                  dists=None,query='INTERSECT',output=None):
-
-        if self.is_committed == 0:
-            rtcCommitScene(self.scene_i)
-            self.is_committed = 1
-
+                  dists=None, query='INTERSECT', output=None, threads=0):
         cdef int nv = vec_origins.shape[0]
-        cdef int i, vd_i, vd_step
         cdef np.ndarray[np.int32_t, ndim=1] intersect_ids
         cdef np.ndarray[np.float32_t, ndim=1] tfars
         cdef np.ndarray[np.float32_t, ndim=1] u_arr, v_arr
         cdef np.ndarray[np.float32_t, ndim=2] Ng_arr
         cdef np.ndarray[np.int32_t, ndim=1] primID_arr, geomID_arr
         cdef rayQueryType query_type
+        cdef int nthreads
+        cdef np.ndarray[np.float32_t, ndim=1] work_tfars
+        cdef bint user_dists = False
+        cdef bint copy_dists_back = False
+
+        if not isinstance(threads, numbers.Integral):
+            raise TypeError("`threads` must be an integer, got %r" % (threads,))
+        nthreads = int(threads)
 
         if query == 'INTERSECT':
             query_type = intersect
         elif query == 'OCCLUDED':
             query_type = occluded
         elif query == 'DISTANCE':
             query_type = distance
-
         else:
             raise ValueError("Embree ray query type %s not recognized."
                 "\nAccepted types are (INTERSECT,OCCLUDED,DISTANCE)" % (query))
@@ -72,6 +177,7 @@ cdef class EmbreeScene:
             tfars = np.empty(nv, 'float32')
             tfars.fill(dists)
         else:
+            user_dists = True
             tfars = dists
 
         if output:
@@ -85,59 +191,74 @@ cdef class EmbreeScene:
         if not output or query_type == occluded:
             intersect_ids = np.empty(nv, dtype="int32")
 
-        cdef rtcr.RTCRayHit rayhit
-        cdef unsigned int INVALID_GEOMETRY_ID = 0xFFFFFFFF
-        cdef bint use_output = bool(output)
-        vd_i = 0
-        vd_step = 1
-        # If vec_directions is 1 long, we won't be updating it.
-        if vec_directions.shape[0] == 1: vd_step = 0
-
-        with nogil:
-            for i in range(nv):
-                rayhit.ray.org_x = vec_origins[i, 0]
-                rayhit.ray.org_y = vec_origins[i, 1]
-                rayhit.ray.org_z = vec_origins[i, 2]
-                rayhit.ray.dir_x = vec_directions[vd_i, 0]
-                rayhit.ray.dir_y = vec_directions[vd_i, 1]
-                rayhit.ray.dir_z = vec_directions[vd_i, 2]
-                rayhit.ray.tnear = 0.0
-                rayhit.ray.tfar = tfars[i]
-                rayhit.hit.geomID = INVALID_GEOMETRY_ID
-                rayhit.hit.primID = INVALID_GEOMETRY_ID
-                rayhit.hit.instID[0] = INVALID_GEOMETRY_ID
-                rayhit.ray.mask = 0xFFFFFFFF
-                rayhit.ray.time = 0.0
-                rayhit.ray.flags = 0
-                vd_i += vd_step
-
-                if query_type == intersect or query_type == distance:
-                    rtcIntersect1(self.scene_i, &rayhit, NULL)
-                    if not use_output:
-                        if query_type == intersect:
-                            intersect_ids[i] = -1 if rayhit.hit.primID == INVALID_GEOMETRY_ID else <int>rayhit.hit.primID
-                        else:
-                            tfars[i] = rayhit.ray.tfar
-                    else:
-                        primID_arr[i] = -1 if rayhit.hit.primID == INVALID_GEOMETRY_ID else <int>rayhit.hit.primID
-                        geomID_arr[i] = -1 if rayhit.hit.geomID == INVALID_GEOMETRY_ID else <int>rayhit.hit.geomID
-                        u_arr[i] = rayhit.hit.u
-                        v_arr[i] = rayhit.hit.v
-                        tfars[i] = rayhit.ray.tfar
-                        Ng_arr[i, 0] = rayhit.hit.Ng_x
-                        Ng_arr[i, 1] = rayhit.hit.Ng_y
-                        Ng_arr[i, 2] = rayhit.hit.Ng_z
-                else:
-                    rtcOccluded1(self.scene_i, &rayhit.ray, NULL)
-                    # In Embree 4, occlusion is signaled by setting ray.tfar to -inf
-                    intersect_ids[i] = 0 if rayhit.ray.tfar < 0 else -1
+        cdef RayJob job
+        job.org_s0 = np.PyArray_STRIDES(vec_origins)[0]
+        job.dir_s0 = np.PyArray_STRIDES(vec_directions)[0]
+        job.tfar_s0 = np.PyArray_STRIDES(tfars)[0]
+
+        if user_dists and nv > 1:
+            if tfars.shape[0] != nv:
+                raise ValueError(
+                    "dists must have one entry per ray for threaded queries"
+                )
+
+        if not _ray_buffers_are_thread_safe(
+            nv, job.org_s0, job.dir_s0, job.tfar_s0, user_dists
+        ):
+            raise ValueError(
+                "threaded queries require distinct per-ray rows in origins, "
+                "directions, and dists"
+            )
+
+        if user_dists and nv > 1 and not np.PyArray_ISCONTIGUOUS(tfars):
+            work_tfars = np.ascontiguousarray(tfars)
+            copy_dists_back = True
+        else:
+            work_tfars = tfars
+
+        # Commit while holding the GIL; Embree forbids commit/traversal overlap.
+        if self.is_committed == 0:
+            rtcCommitScene(self.scene_i)
+            self.is_committed = 1
+
+        job.scene = self.scene_i
+        job.org = <char*>np.PyArray_DATA(vec_origins)
+        job.org_s1 = np.PyArray_STRIDES(vec_origins)[1]
+        job.dir = <char*>np.PyArray_DATA(vec_directions)
+        job.dir_s1 = np.PyArray_STRIDES(vec_directions)[1]
+        job.direction_row_step = 0 if vec_directions.shape[0] == 1 else 1
+        job.tfar = <char*>np.PyArray_DATA(work_tfars)
+        job.tfar_s0 = np.PyArray_STRIDES(work_tfars)[0]
+        job.query_type = <int>query_type
+        job.use_output = bool(output)
+        job.intersect_ids = <int*>np.PyArray_DATA(intersect_ids) if (not output or query_type == occluded) else NULL
+        if output:
+            job.u_arr = <float*>np.PyArray_DATA(u_arr)
+            job.v_arr = <float*>np.PyArray_DATA(v_arr)
+            job.Ng_arr = <float*>np.PyArray_DATA(Ng_arr)
+            job.primID_arr = <int*>np.PyArray_DATA(primID_arr)
+            job.geomID_arr = <int*>np.PyArray_DATA(geomID_arr)
+        else:
+            job.u_arr = NULL
+            job.v_arr = NULL
+            job.Ng_arr = NULL
+            job.primID_arr = NULL
+            job.geomID_arr = NULL
+
+        if nv > 0:
+            with nogil:
+                embreex_parallel_for(<size_t>nv, _TBB_GRAIN_SIZE, nthreads,
+                                     _cast_range, <void*>&job)
+
+        if copy_dists_back:
+            dists[:] = work_tfars
 
         if output:
-            return {'u': u_arr, 'v': v_arr, 'Ng': Ng_arr, 'tfar': tfars,
+            return {'u': u_arr, 'v': v_arr, 'Ng': Ng_arr, 'tfar': work_tfars,
                     'primID': primID_arr, 'geomID': geomID_arr}
         else:
             if query_type == distance:
-                return tfars
+                return dists if user_dists else work_tfars
             else:
                 return intersect_ids
 

embreex/tbb_parallel.h

@@ -0,0 +1,33 @@
+// Adapter from Cython's range callback to oneTBB's templated parallel_for.
+
+#pragma once
+
+#include <cstddef>
+
+#include <oneapi/tbb/blocked_range.h>
+#include <oneapi/tbb/parallel_for.h>
+#include <oneapi/tbb/task_arena.h>
+
+// Body of one chunk of rays: [begin, end).
+typedef void (*embreex_range_fn)(void *ctx, size_t begin, size_t end);
+
+// Apply fn to subranges of [0, n); grain controls divisibility. nthreads > 0
+// caps concurrency via task_arena; nthreads <= 0 uses the TBB default pool
+// (Open3D RaycastingScene semantics: any nthreads <= 0 for automatic).
+inline void embreex_parallel_for(size_t n,
+                                 size_t grain,
+                                 int nthreads,
+                                 embreex_range_fn fn,
+                                 void *ctx) {
+    auto body = [&](const tbb::blocked_range<size_t> &r) {
+        fn(ctx, r.begin(), r.end());
+    };
+    const tbb::blocked_range<size_t> range(0, n, grain);
+
+    if (nthreads > 0) {
+        tbb::task_arena arena(nthreads);
+        arena.execute([&]() { tbb::parallel_for(range, body); });
+    } else {
+        tbb::parallel_for(range, body);
+    }
+}

package/embree.json

@@ -44,5 +44,20 @@
     "sha256": "d951e5e6bd295c54cdd66be9cdb44a4e8c42fb38a99f94f79305e48765fc3454",
     "target": "../embree4",
     "url": "https://github.com/RenderKit/embree/releases/download/v4.4.0/embree-4.4.0.x64.windows.zip"
+  },
+  {
+    "comment": "Headers only; pin no newer than the oldest bundled TBB runtime (2021.11 for embree 4.4.0). Revalidate when updating Embree.",
+    "extract_skip": [
+      ".*", ".github/*", "*.md", "*.bazel", "*.bazelrc", "*.bazelversion",
+      "CMakeLists.txt", "third-party-programs.txt",
+      "cmake/*", "doc/*", "examples/*", "integration/*", "python/*",
+      "src/*", "test/*"
+    ],
+    "name": "tbb-headers",
+    "platform": "any",
+    "sha256": "782ce0cab62df9ea125cdea253a50534862b563f1d85d4cda7ad4e77550ac363",
+    "strip_components": 1,
+    "target": "../tbb",
+    "url": "https://github.com/uxlfoundation/oneTBB/archive/refs/tags/v2021.11.0.tar.gz"
   }
 ]

package/fetch-embree.py

@@ -198,9 +198,10 @@ def is_current_platform(platform: str, architecture: Optional[str]) -> bool:
     Parameters
     ----------
     platform
-      Checked against `platform.system`
+      Checked against `platform.system`, or `"any"` for a
+      platform-independent resource such as a source archive.
     architecture
-      Checked against `platform.uname.machine`
+      Checked against `platform.uname.machine`, or `None` for any.
 
     Returns
     -------
@@ -210,6 +211,9 @@ def is_current_platform(platform: str, architecture: Optional[str]) -> bool:
     """
     # 'linux', 'darwin', 'windows'
 
+    if platform == "any":
+        return True
+
     if architecture is not None:
         # Check for cibuildwheel target architecture first
         target_arch = os.environ.get("ARCHFLAGS", "")
@@ -266,5 +270,6 @@ def is_current_platform(platform: str, architecture: Optional[str]) -> bool:
             subset = option.copy()
             subset.pop("name")
             subset.pop("platform")
-            subset.pop("architecture")
+            subset.pop("architecture", None)
+            subset.pop("comment", None)
             handle_fetch(**subset)

pyproject.toml

@@ -34,10 +34,10 @@ skip =  "*i686 *-win32 *musllinux*"
 manylinux-x86_64-image = "manylinux_2_28"
 before-test = "pip install pytest"
 test-command = "pytest -v {project}/tests"
-before-build = "python {project}/package/fetch-embree.py --install embree4"
+before-build = "python {project}/package/fetch-embree.py --install embree4,tbb-headers"
 
 [tool.cibuildwheel.windows]
-before-build = "pip install delvewheel && python {project}\\package\\fetch-embree.py --install embree4"
+before-build = "pip install delvewheel && python {project}\\package\\fetch-embree.py --install embree4,tbb-headers"
 repair-wheel-command = "delvewheel repair --add-path embree4\\bin --no-mangle tbb12.dll;embree4.dll -w {dest_dir} {wheel}"
 
 [tool.cibuildwheel.macos]