diff --git a/embreex/rtcore_scene.pyx b/embreex/rtcore_scene.pyx index ee14b46..ddc02c5 100644 --- a/embreex/rtcore_scene.pyx +++ b/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 = 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 + i * j.org_s0 + # Broadcast a single direction row across all origins. + vd_i = i * j.direction_row_step + direction_ptr = j.dir + vd_i * j.dir_s0 + distance_ptr = j.tfar + i * j.tfar_s0 + + rayhit.ray.org_x = (origin_ptr)[0] + rayhit.ray.org_y = ((origin_ptr + j.org_s1))[0] + rayhit.ray.org_z = ((origin_ptr + 2 * j.org_s1))[0] + rayhit.ray.dir_x = (direction_ptr)[0] + rayhit.ray.dir_y = ((direction_ptr + j.dir_s1))[0] + rayhit.ray.dir_z = ((direction_ptr + 2 * j.dir_s1))[0] + rayhit.ray.tnear = 0.0 + rayhit.ray.tfar = (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 == intersect or j.query_type == distance: + rtcIntersect1(j.scene, &rayhit, NULL) + if not j.use_output: + if j.query_type == intersect: + j.intersect_ids[i] = -1 if rayhit.hit.primID == INVALID_GEOMETRY_ID else rayhit.hit.primID + else: + (distance_ptr)[0] = rayhit.ray.tfar + else: + j.primID_arr[i] = -1 if rayhit.hit.primID == INVALID_GEOMETRY_ID else rayhit.hit.primID + j.geomID_arr[i] = -1 if rayhit.hit.geomID == INVALID_GEOMETRY_ID else rayhit.hit.geomID + j.u_arr[i] = rayhit.hit.u + j.v_arr[i] = rayhit.hit.v + (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 < sizeof(float): + return False + if dir_s0 < sizeof(float): + return False + if user_dists and tfar_s0 < sizeof(float): + return False + return True + + cdef class EmbreeScene: def __init__(self, rtc.EmbreeDevice device=None, robust=True): if device is None: @@ -39,20 +143,22 @@ 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 @@ -60,7 +166,6 @@ cdef class EmbreeScene: 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 rayhit.hit.primID - else: - tfars[i] = rayhit.ray.tfar - else: - primID_arr[i] = -1 if rayhit.hit.primID == INVALID_GEOMETRY_ID else rayhit.hit.primID - geomID_arr[i] = -1 if rayhit.hit.geomID == INVALID_GEOMETRY_ID else 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 = np.PyArray_DATA(vec_origins) + job.org_s1 = np.PyArray_STRIDES(vec_origins)[1] + job.dir = 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 = np.PyArray_DATA(work_tfars) + job.tfar_s0 = np.PyArray_STRIDES(work_tfars)[0] + job.query_type = query_type + job.use_output = bool(output) + job.intersect_ids = np.PyArray_DATA(intersect_ids) if (not output or query_type == occluded) else NULL + if output: + job.u_arr = np.PyArray_DATA(u_arr) + job.v_arr = np.PyArray_DATA(v_arr) + job.Ng_arr = np.PyArray_DATA(Ng_arr) + job.primID_arr = np.PyArray_DATA(primID_arr) + job.geomID_arr = 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(nv, _TBB_GRAIN_SIZE, nthreads, + _cast_range, &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 diff --git a/embreex/tbb_parallel.h b/embreex/tbb_parallel.h new file mode 100644 index 0000000..125d84c --- /dev/null +++ b/embreex/tbb_parallel.h @@ -0,0 +1,33 @@ +// Adapter from Cython's range callback to oneTBB's templated parallel_for. + +#pragma once + +#include + +#include +#include +#include + +// 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 &r) { + fn(ctx, r.begin(), r.end()); + }; + const tbb::blocked_range 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); + } +} diff --git a/package/embree.json b/package/embree.json index 1637bc9..aac9115 100644 --- a/package/embree.json +++ b/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" } ] diff --git a/package/fetch-embree.py b/package/fetch-embree.py index 1fb38b5..6707119 100755 --- a/package/fetch-embree.py +++ b/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) diff --git a/pyproject.toml b/pyproject.toml index 8ea5e66..b5d3a94 100644 --- a/pyproject.toml +++ b/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] diff --git a/setup.py b/setup.py index 50b9917..cad7276 100755 --- a/setup.py +++ b/setup.py @@ -13,12 +13,17 @@ def ext_modules(): """Generate a list of extension modules for embreex.""" + # Fetch oneTBB headers separately; Embree's bundle contains only the runtime. + tbb_include = os.path.join(_cwd, "tbb", "include") + cxx_std = ["-std=c++11"] if os.name != "nt" else ["/std:c++14"] + if os.name == "nt": # embree search locations on windows includes = [ get_include(), "c:/Program Files/Intel/Embree4/include", os.path.join(_cwd, "embree4", "include"), + tbb_include, ] libraries = [ "c:/Program Files/Intel/Embree4/lib", @@ -30,6 +35,7 @@ def ext_modules(): get_include(), "/opt/local/include", os.path.join(_cwd, "embree4", "include"), + tbb_include, ] libraries = ["/opt/local/lib", os.path.join(_cwd, "embree4", "lib")] @@ -37,16 +43,21 @@ def ext_modules(): for ext in ext_modules: ext.include_dirs = includes ext.library_dirs = libraries + ext.extra_compile_args = getattr(ext, "extra_compile_args", []) + cxx_std # on macOS with Embree 4.x, link against the versioned library directly if sys.platform == "darwin": - ext.libraries = ["embree4.4"] + # `libtbb.dylib` is a symlink created by `package/embree.json` + ext.libraries = ["embree4.4", "tbb"] # Add rpath to find libembree4 during build and set loader_path for runtime ext.extra_link_args = [ "-Wl,-rpath,@loader_path", "-Wl,-rpath," + os.path.join(_cwd, "embree4", "lib"), ] + elif os.name == "nt": + # the import library in the embree bundle is `tbb12.lib` + ext.libraries = ["embree4", "tbb12"] else: - ext.libraries = ["embree4"] + ext.libraries = ["embree4", "tbb"] return ext_modules diff --git a/tests/test_intersection.py b/tests/test_intersection.py index a10fc8c..4f4356c 100644 --- a/tests/test_intersection.py +++ b/tests/test_intersection.py @@ -218,6 +218,154 @@ def test_occluded(self): self.assertEqual(res[3], -1) +class TestThreadedQueries(TestCase): + """Parity tests across Open3D-style thread counts (0 = automatic).""" + + def setUp(self): + rng = np.random.default_rng(0) + scene = rtcs.EmbreeScene() + TriangleMesh(scene, np.array(xplane(7.0), "float32")) + self.scene = scene + n = 20_000 + self.origins = np.zeros((n, 3), dtype="float32") + self.origins[:, 0] = 0.1 + # Include both hits and misses across multiple chunks. + self.origins[:, 1] = rng.uniform(-3.0, 3.0, n).astype("float32") + self.origins[:, 2] = rng.uniform(-3.0, 3.0, n).astype("float32") + self.dirs = np.zeros((n, 3), dtype="float32") + self.dirs[:, 0] = 1.0 + + def test_threads_match_serial(self): + for query in ("INTERSECT", "OCCLUDED", "DISTANCE"): + ref = self.scene.run(self.origins, self.dirs, query=query, threads=1) + for threads in (0, -1, -2, 2, 3, 8): + got = self.scene.run( + self.origins, self.dirs, query=query, threads=threads + ) + np.testing.assert_array_equal(ref, got, err_msg=f"{query} t={threads}") + + def test_threads_match_serial_output_dict(self): + ref = self.scene.run(self.origins, self.dirs, output=True, threads=1) + hit = ref["primID"] != -1 + # Guard against a fixture that exercises only one output branch. + self.assertTrue(hit.any() and not hit.all()) + for threads in (0, -1, -2, 2, 3, 8): + got = self.scene.run( + self.origins, self.dirs, output=True, threads=threads + ) + for key in ("primID", "geomID", "tfar", "u", "v", "Ng"): + np.testing.assert_array_equal( + ref[key], got[key], err_msg=f"{key} t={threads}" + ) + + def test_threads_output_dict_is_repeatable(self): + a = self.scene.run(self.origins, self.dirs, output=True, threads=8) + b = self.scene.run(self.origins, self.dirs, output=True, threads=8) + for key in ("primID", "geomID", "tfar", "u", "v", "Ng"): + np.testing.assert_array_equal(a[key], b[key], err_msg=key) + + def test_threads_broadcast_direction(self): + repeated_dirs = np.tile(self.dirs[:1], (len(self.origins), 1)) + ref = self.scene.run(self.origins, repeated_dirs, threads=1) + broadcast = self.scene.run(self.origins, self.dirs[:1], threads=1) + threaded = self.scene.run(self.origins, self.dirs[:1], threads=8) + np.testing.assert_array_equal(ref, broadcast) + np.testing.assert_array_equal(ref, threaded) + + def test_threads_non_contiguous_input(self): + origins = self.origins[::3] + directions = self.dirs[::3] + ref = self.scene.run( + np.ascontiguousarray(origins), + np.ascontiguousarray(directions), + threads=1, + ) + np.testing.assert_array_equal( + ref, self.scene.run(origins, directions, threads=1) + ) + np.testing.assert_array_equal( + ref, self.scene.run(origins, directions, threads=8) + ) + + def test_threads_degenerate_sizes(self): + for n in (0, 1, 2, 1023, 1024, 1025): + origins, directions = self.origins[:n], self.dirs[:n] + np.testing.assert_array_equal( + self.scene.run(origins, directions, threads=1), + self.scene.run(origins, directions, threads=8), + err_msg=str(n), + ) + + def test_threads_dists_written_in_place(self): + ref = np.full(len(self.origins), 20.0, dtype="float32") + got = ref.copy() + self.scene.run(self.origins, self.dirs, dists=ref, query="DISTANCE", threads=1) + self.scene.run(self.origins, self.dirs, dists=got, query="DISTANCE", threads=8) + np.testing.assert_array_equal(ref, got) + + def test_threads_non_contiguous_dists(self): + dists = np.full(len(self.origins) * 2, 20.0, dtype="float32")[::2] + expected = dists.copy() + self.scene.run( + self.origins, self.dirs, dists=expected, query="DISTANCE", threads=1 + ) + work = dists.copy() + self.scene.run(self.origins, self.dirs, dists=work, query="DISTANCE", threads=8) + np.testing.assert_array_equal(expected, work) + + def test_threads_non_integer_raises(self): + for value in (-0.5, 1.9, "3"): + with self.assertRaises(TypeError): + self.scene.run(self.origins, self.dirs, threads=value) + + def test_invalid_threads_leaves_scene_usable(self): + scene = rtcs.EmbreeScene() + empty = np.empty((0, 3), dtype="float32") + with self.assertRaises(TypeError): + scene.run(empty, empty, threads=1.5) + TriangleMesh(scene, np.array(xplane(7.0), "float32")) + origins = np.array([[0.1, 0.0, 0.0]], dtype="float32") + directions = np.array([[1.0, 0.0, 0.0]], dtype="float32") + np.testing.assert_array_equal(scene.run(origins, directions), [0]) + + def test_threads_overlapping_dists_rejected(self): + dists = np.broadcast_to(np.array([20.0], dtype="float32"), (2,)) + with self.assertRaises(ValueError): + self.scene.run( + self.origins[:2], + self.dirs[:2], + dists=dists, + query="DISTANCE", + threads=2, + ) + + def test_threads_mismatched_dists_length_rejected(self): + dists = np.array([20.0], dtype="float32") + with self.assertRaises(ValueError): + self.scene.run( + self.origins[:2], + self.dirs[:2], + dists=dists, + query="DISTANCE", + threads=2, + ) + + def test_threads_from_python_threads(self): + """Concurrent `run()` calls on one scene must not interfere.""" + from concurrent.futures import ThreadPoolExecutor + + ref = self.scene.run(self.origins, self.dirs) + with ThreadPoolExecutor(4) as pool: + got = list( + pool.map( + lambda _: self.scene.run(self.origins, self.dirs, threads=4), + range(8), + ) + ) + for i, g in enumerate(got): + np.testing.assert_array_equal(ref, g, err_msg=str(i)) + + if __name__ == "__main__": from unittest import main