Feat: add OptiX (hardware rt) support to toy path tracer

Author: RDambrosio016

crates/cuda_builder/src/lib.rs

@@ -138,6 +138,8 @@ pub struct CudaBuilder {
     pub override_libm: bool,
     /// Whether to generate any debug info and what level of info to generate.
     pub debug: DebugInfo,
+    /// Additional arguments passed to cargo during `cargo build`.
+    pub build_args: Vec<String>,
 }
 
 impl CudaBuilder {
@@ -158,9 +160,17 @@ impl CudaBuilder {
             optix: false,
             override_libm: true,
             debug: DebugInfo::None,
+            build_args: vec![],
         }
     }
 
+    /// Additional arguments passed to cargo during `cargo build`.
+    pub fn build_args(mut self, args: &[impl AsRef<str>]) -> Self {
+        self.build_args
+            .extend(args.iter().map(|s| s.as_ref().to_owned()));
+        self
+    }
+
     /// Whether to generate any debug info and what level of info to generate.
     pub fn debug(mut self, debug: DebugInfo) -> Self {
         self.debug = debug;
@@ -433,6 +443,8 @@ fn invoke_rustc(builder: &CudaBuilder) -> Result<PathBuf, CudaBuilderError> {
         target,
     ]);
 
+    cargo.args(&builder.build_args);
+
     cargo.env(dylib_path_envvar(), new_path);
 
     if builder.release {

crates/cust/src/memory/device/device_slice.rs

@@ -464,6 +464,15 @@ fn slice_end_index_overflow_fail() -> ! {
     panic!("attempted to index slice up to maximum usize");
 }
 
+impl<T: DeviceCopy> DeviceSliceIndex<T> for usize {
+    unsafe fn get_unchecked(self, slice: &DeviceSlice<T>) -> DeviceSlice<T> {
+        (self..self + 1).get_unchecked(slice)
+    }
+    fn index(self, slice: &DeviceSlice<T>) -> DeviceSlice<T> {
+        slice.index(self..self + 1)
+    }
+}
+
 impl<T: DeviceCopy> DeviceSliceIndex<T> for Range<usize> {
     unsafe fn get_unchecked(self, slice: &DeviceSlice<T>) -> DeviceSlice<T> {
         DeviceSlice::from_raw_parts(slice.as_device_ptr().add(self.start), self.end - self.start)

crates/optix/examples/rust/ex04_mesh_gpu/src/lib.rs

@@ -8,7 +8,7 @@
 
 use cuda_std::kernel;
 use optix_device::{
-    get_launch_index,
+    closesthit, get_launch_index,
     glam::*,
     misc::*,
     payload,
@@ -65,7 +65,7 @@ fn random_color(i: u32) -> Vec4 {
 
 #[kernel]
 pub unsafe fn __closesthit__radiance() {
-    let prim_id = primitive_index();
+    let prim_id = closesthit::primitive_index();
     let buf = get_color_buf();
     *buf = random_color(prim_id);
 }

crates/optix/src/acceleration.rs

@@ -877,7 +877,7 @@ pub enum AccelEmitDesc {
 /// Used to communicate bounds info to and from OptiX for bounding custom primitives
 /// and instances
 #[repr(C)]
-#[derive(DeviceCopy, Copy, Clone)]
+#[derive(Debug, DeviceCopy, Copy, Clone)]
 pub struct Aabb {
     min: Vector3<f32>,
     max: Vector3<f32>,

crates/optix_device/Cargo.toml

@@ -10,3 +10,4 @@ cuda_std = { version = "0.2", path = "../cuda_std" }
 glam = { version = "0.20", features=["cuda", "libm"], default-features=false }
 paste = "1.0.6"
 seq-macro = "0.3.0"
+cust_core = { version = "0.1", path = "../cust_core" }

crates/optix_device/src/intersect.rs

@@ -2,6 +2,15 @@ use cuda_std::gpu_only;
 use paste::paste;
 use seq_macro::seq;
 
+#[gpu_only]
+pub fn primitive_index() -> u32 {
+    let mut idx: u32;
+    unsafe {
+        asm!("call ({}), _optix_read_primitive_idx, ();", out(reg32) idx);
+    }
+    idx
+}
+
 pub trait IntersectionPayload {
     fn report_intersection(hit_t: f32, hit_kind: u8, payload: Self) -> bool;
 }
@@ -22,11 +31,15 @@ macro_rules! impl_intersection_payload {
                             let out: u32;
                             unsafe {
                                 asm!(
+                                    "{{",
+                                    ".reg .f32 %f0;",
+                                    "mov.f32 %f0, {hit_t};",
                                     concat!("call ({}), _optix_report_intersection_", stringify!($num)),
-                                    concat!(", ({}, {}, ", #(concat!("{", stringify!(p~P), "},"),)* ");"),
+                                    concat!(", (%f0, {}", #(concat!(", {", stringify!(p~P), "}"),)* ");"),
+                                    "}}",
                                     out(reg32) out,
-                                    in(reg32) hit_t,
                                     in(reg32) hit_kind,
+                                    hit_t = in(reg32) hit_t,
                                     #(
                                         p~P = in(reg32) p~P,
                                     )*

crates/optix_device/src/lib.rs

@@ -14,12 +14,15 @@ pub mod payload;
 mod ray;
 pub mod sys;
 pub mod trace;
+pub mod transform;
 pub mod util;
 
 use cuda_std::*;
 pub use glam;
 use glam::UVec3;
 
+pub use misc::*;
+
 extern "C" {
     pub fn vprintf(format: *const u8, valist: *const core::ffi::c_void) -> i32;
 }
@@ -65,7 +68,7 @@ pub mod raygen {
 /// Functions/items only available in miss programs (`__miss__`).
 pub mod intersection {
     #[doc(inline)]
-    pub use crate::intersect::{get_attribute, report_intersection};
+    pub use crate::intersect::{get_attribute, primitive_index, report_intersection};
     #[doc(inline)]
     pub use crate::ray::*;
 }
@@ -75,7 +78,9 @@ pub mod anyhit {
     #[doc(inline)]
     pub use crate::hit::*;
     #[doc(inline)]
-    pub use crate::intersect::{get_attribute, ignore_intersection, terminate_ray};
+    pub use crate::intersect::{
+        get_attribute, ignore_intersection, primitive_index, terminate_ray,
+    };
     #[doc(inline)]
     pub use crate::ray::*;
 }
@@ -85,7 +90,7 @@ pub mod closesthit {
     #[doc(inline)]
     pub use crate::hit::*;
     #[doc(inline)]
-    pub use crate::intersect::get_attribute;
+    pub use crate::intersect::{get_attribute, primitive_index};
     #[doc(inline)]
     pub use crate::ray::{
         ray_flags, ray_time, ray_tmax, ray_tmin, ray_visibility_mask, ray_world_direction,

crates/optix_device/src/misc.rs

@@ -1,10 +1,13 @@
 use cuda_std::gpu_only;
 
+/// Retrieves the data past the SBT header for this particular program
+///
+/// # Safety
+///
+/// The type requested must match with what is stored in the SBT.
 #[gpu_only]
-pub fn primitive_index() -> u32 {
-    let mut idx: u32;
-    unsafe {
-        asm!("call ({}), _optix_read_primitive_idx, ();", out(reg32) idx);
-    }
-    idx
+pub unsafe fn sbt_data<T>() -> &'static T {
+    let ptr: *const T;
+    asm!("call ({}), _optix_get_sbt_data_ptr_64, ();", out(reg64) ptr);
+    &*ptr
 }

crates/optix_device/src/ray.rs

@@ -46,7 +46,14 @@ pub fn ray_tmin() -> f32 {
     let x: f32;
 
     unsafe {
-        asm!("call ({}), _optix_get_ray_tmin, ();", out(reg32) x);
+        asm!(
+            "{{",
+            ".reg .f32 %f<1>;",
+            "call (%f0), _optix_get_ray_tmin, ();",
+            "mov.f32 {}, %f0;",
+            "}}",
+            out(reg32) x
+        );
     }
 
     x
@@ -58,7 +65,14 @@ pub fn ray_tmax() -> f32 {
     let x: f32;
 
     unsafe {
-        asm!("call ({}), _optix_get_ray_tmax, ();", out(reg32) x);
+        asm!(
+            "{{",
+            ".reg .f32 %f<1>;",
+            "call (%f0), _optix_get_ray_tmax, ();",
+            "mov.f32 {}, %f0;",
+            "}}",
+            out(reg32) x
+        );
     }
 
     x

crates/optix_device/src/trace.rs

@@ -1,11 +1,12 @@
 use crate::sys::*;
+use cust_core::DeviceCopy;
 use glam::Vec3;
 use paste::paste;
 use seq_macro::seq;
 
 /// An opaque handle to a traversable BVH.
 #[repr(transparent)]
-#[derive(Clone, Copy, Debug, PartialEq, Eq)]
+#[derive(Clone, Copy, Debug, PartialEq, Eq, DeviceCopy)]
 pub struct TraversableHandle(pub(crate) u64);
 
 impl TraversableHandle {