feat(lib): add BVH acceleration

differt-core/Cargo.toml

@@ -2,7 +2,12 @@
 harness = false
 name = "bench_main"
 
+[build-dependencies]
+cxx-build = {version = ">=1.0,<1.0.178"}
+pyo3-build-config = {version = "0.25"}
+
 [dependencies]
+cxx = {version = ">=1.0,<1.0.178"}
 indexmap = {version = "2.5.0", features = ["serde"]}
 log = "0.4"
 nalgebra = "0.32.3"

differt-core/build.rs

@@ -0,0 +1,60 @@
+/// Build script for differt-core.
+///
+/// 1. Queries JAX for XLA FFI header locations
+/// 2. Compiles the C++ FFI shim via cxx-build
+use std::process::exit;
+
+fn main() {
+    // Only build FFI when the feature is enabled
+    use std::{env, path::PathBuf, str::from_utf8};
+
+    // Find the Python interpreter:
+    // 1. PYTHON env var (covers VIRTUAL_ENV and explicit overrides)
+    // 2. pyo3_build_config: the interpreter pyo3 itself was built against
+    // 3. Fall back to "python3"
+    let python = env::var("PYTHON")
+        .ok()
+        .or_else(|| pyo3_build_config::get().executable.clone())
+        .unwrap_or_else(|| "python3".to_owned());
+
+    // Query JAX for its XLA FFI include directory
+    let output = std::process::Command::new(&python)
+        .args([
+            "-c",
+            "from jax.ffi import include_dir; print(include_dir())",
+        ])
+        .output();
+
+    let include_path = match output {
+        Ok(ref out) if out.status.success() => {
+            let path = from_utf8(&out.stdout)
+                .expect("Invalid UTF-8 from JAX include_dir()")
+                .trim()
+                .to_string();
+            if path.is_empty() {
+                None
+            } else {
+                Some(PathBuf::from(path))
+            }
+        },
+        _ => None,
+    };
+
+    if let Some(include_path) = include_path {
+        println!("cargo:rerun-if-changed=src/ffi.cc");
+        println!("cargo:rerun-if-changed=include/ffi.h");
+
+        cxx_build::bridge("src/accel/ffi.rs")
+            .file("src/ffi.cc")
+            .std("c++17")
+            .include(&include_path)
+            .include("include")
+            .compile("differt-ffi");
+    } else {
+        println!(
+            "cargo:error=JAX not found or missing jax.ffi.include_dir(). Python interpreter used: \
+             {python}"
+        );
+        exit(1);
+    }
+}

differt-core/include/ffi.h

@@ -0,0 +1,17 @@
+#pragma once
+
+#include "xla/ffi/api/ffi.h"
+
+// BVH nearest-hit: for each ray, find the closest active triangle.
+// Inputs: ray_origins [num_rays, 3], ray_directions [num_rays, 3],
+//         active_mask [num_triangles] (PRED, or [0] for no mask)
+// Attrs: bvh_id (u64)
+// Outputs: hit_indices [num_rays] (i32), hit_t [num_rays] (f32)
+extern "C" XLA_FFI_Error *BvhNearestHit(XLA_FFI_CallFrame *call_frame);
+
+// BVH get-candidates: for each ray, find candidate triangles with expanded AABBs.
+// Inputs: ray_origins [num_rays, 3], ray_directions [num_rays, 3]
+// Attrs: bvh_id (u64), expansion (f32), max_candidates (i32)
+// Outputs: candidate_indices [num_rays, max_candidates] (i32),
+//          candidate_counts [num_rays] (i32)
+extern "C" XLA_FFI_Error *BvhGetCandidates(XLA_FFI_CallFrame *call_frame);

differt-core/pyproject.toml

@@ -1,6 +1,6 @@
 [build-system]
 build-backend = "maturin"
-requires = ["maturin>=1.9.6,<2"]
+requires = ["maturin>=1.9.6,<2", "jax>=0.8.1"]
 
 [project]
 authors = [
@@ -31,6 +31,8 @@ bindings = "pyo3"
 features = ["pyo3/extension-module"]
 include = [
   {path = "src/**/*", format = "sdist"},
+  {path = "include/**/*", format = "sdist"},
+  {path = "build.rs", format = "sdist"},
   {path = "LICENSE.md", format = "sdist"},
   {path = "README.md", format = "sdist"},
 ]

differt-core/python/differt_core/_differt_core/__init__.pyi

@@ -3,6 +3,7 @@ from types import ModuleType
 __version__: str
 __version_info__: tuple[int, int, int]
 
+accel: ModuleType
 geometry: ModuleType
 rt: ModuleType
 scene: ModuleType

differt-core/python/differt_core/_differt_core/accel/bvh.pyi

@@ -0,0 +1,32 @@
+import numpy as np
+from jaxtyping import Float, Int
+
+class TriangleBvh:
+    def __init__(
+        self, triangle_vertices: Float[np.ndarray, "num_triangles 9"]
+    ) -> None: ...
+    @property
+    def num_triangles(self) -> int: ...
+    @property
+    def num_nodes(self) -> int: ...
+    def register(self) -> int: ...
+    def unregister(self) -> None: ...
+    def nearest_hit(
+        self,
+        ray_origins: Float[np.ndarray, "num_rays 3"],
+        ray_directions: Float[np.ndarray, "num_rays 3"],
+        active_mask: np.ndarray | None = None,
+    ) -> tuple[Int[np.ndarray, " num_rays"], Float[np.ndarray, " num_rays"]]: ...
+    def get_candidates(
+        self,
+        ray_origins: Float[np.ndarray, "num_rays 3"],
+        ray_directions: Float[np.ndarray, "num_rays 3"],
+        expansion: float = 0.0,
+        max_candidates: int = 256,
+    ) -> tuple[
+        Int[np.ndarray, "num_rays max_candidates"],
+        Int[np.ndarray, " num_rays"],
+    ]: ...
+
+def bvh_nearest_hit_capsule() -> object: ...
+def bvh_get_candidates_capsule() -> object: ...

differt-core/python/differt_core/accel/__init__.py

@@ -0,0 +1 @@
+"""Acceleration structures for ray tracing."""

differt-core/python/differt_core/accel/bvh.py

@@ -0,0 +1,7 @@
+__all__ = ("TriangleBvh",)
+
+from differt_core import _differt_core
+
+TriangleBvh = _differt_core.accel.bvh.TriangleBvh
+bvh_nearest_hit_capsule = _differt_core.accel.bvh.bvh_nearest_hit_capsule
+bvh_get_candidates_capsule = _differt_core.accel.bvh.bvh_get_candidates_capsule