Feat: OptiX device 2; hardware boogaloo

Author: RDambrosio016

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

@@ -113,8 +113,7 @@ pub unsafe fn __raygen__renderFrame() {
         RayType::SurfaceRay as u32,
         1,
         RayType::SurfaceRay as u32,
-        &mut p0,
-        &mut p1,
+        [&mut p0, &mut p1],
     );
 
     let fb_index = i.x + i.y * PARAMS.frame.size.x;

crates/optix_device/Cargo.toml

@@ -9,3 +9,4 @@ bitflags = "1.3.2"
 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"

crates/optix_device/src/hit.rs

@@ -0,0 +1,117 @@
+use cuda_std::gpu_only;
+use glam::Vec3;
+
+/// The type of primitive that a ray hit.
+#[repr(u32)]
+#[derive(Debug, Copy, Clone, PartialEq, Eq, Hash)]
+pub enum HitKind {
+    /// A custom primitive, the value is the custom
+    /// 7-bit hit kind set when reporting an intersection.
+    Custom(u8),
+    /// B-spline curve of degree 2 with circular cross-section.
+    RoundQuadraticBSpline,
+    /// B-spline curve of degree 3 with circular cross-section.
+    RoundCubicBSpline,
+    /// Piecewise linear curve with circular cross-section.
+    RoundLinear,
+    /// CatmullRom curve with circular cross-section.
+    RoundCatmullRom,
+    /// ▲
+    Triangle,
+}
+
+#[repr(u32)]
+#[allow(dead_code)]
+enum OptixPrimitiveType {
+    Custom = 0x2500,
+    RoundQuadraticBSpline = 0x2501,
+    RoundCubicBSpline = 0x2502,
+    RoundLinear = 0x2503,
+    RoundCatmullRom = 0x2504,
+    Triangle = 0x2531,
+}
+
+#[gpu_only]
+fn get_primitive_type(val: u8) -> OptixPrimitiveType {
+    let raw: u32;
+    unsafe {
+        asm!("call ({}), _optix_get_primitive_type_from_hit_kind, ({});", out(reg32) raw, in(reg32) val);
+        core::mem::transmute(raw)
+    }
+}
+
+impl HitKind {
+    pub(crate) fn from_raw(val: u8) -> Self {
+        let kind = get_primitive_type(val);
+        match kind {
+            OptixPrimitiveType::Custom => HitKind::Custom(val & 0b0111_1111),
+            OptixPrimitiveType::RoundQuadraticBSpline => HitKind::RoundQuadraticBSpline,
+            OptixPrimitiveType::RoundCubicBSpline => HitKind::RoundCubicBSpline,
+            OptixPrimitiveType::RoundLinear => HitKind::RoundLinear,
+            OptixPrimitiveType::RoundCatmullRom => HitKind::RoundCatmullRom,
+            OptixPrimitiveType::Triangle => HitKind::Triangle,
+        }
+    }
+}
+
+#[gpu_only]
+fn get_hit_kind() -> u8 {
+    let x: u8;
+    unsafe {
+        asm!("call ({}), _optix_get_hit_kind, ();", out(reg32) x);
+    }
+    x
+}
+
+/// Returns the kind of primitive that was hit by the ray.
+pub fn hit_kind() -> HitKind {
+    HitKind::from_raw(get_hit_kind())
+}
+
+#[gpu_only]
+pub fn is_back_face_hit() -> bool {
+    let hit_kind = get_hit_kind();
+    let x: u32;
+    unsafe {
+        asm!("call ({}), _optix_get_backface_from_hit_kind, ({});", out(reg32) x, in(reg32) hit_kind);
+    }
+    x == 1
+}
+
+/// Whether the ray hit a front face.
+pub fn is_front_face_hit() -> bool {
+    !is_back_face_hit()
+}
+
+/// Whether the ray hit a triangle.
+pub fn is_triangle_hit() -> bool {
+    is_triangle_front_face_hit() || is_triangle_back_face_hit()
+}
+
+const OPTIX_HIT_KIND_TRIANGLE_FRONT_FACE: u8 = 0xFE;
+const OPTIX_HIT_KIND_TRIANGLE_BACK_FACE: u8 = 0xFF;
+
+/// Whether the ray hit the front face of a triangle.
+pub fn is_triangle_front_face_hit() -> bool {
+    get_hit_kind() == OPTIX_HIT_KIND_TRIANGLE_FRONT_FACE
+}
+
+/// Whether the ray hit the back face of a triangle.
+pub fn is_triangle_back_face_hit() -> bool {
+    get_hit_kind() == OPTIX_HIT_KIND_TRIANGLE_BACK_FACE
+}
+
+/// Returns the barycentric coordinates of the hit point on the hit triangle.
+#[gpu_only]
+pub fn triangle_barycentrics() -> Vec3 {
+    let x: f32;
+    let y: f32;
+    unsafe {
+        asm!(
+            "call ({}, {}), _optix_get_triangle_barycentrics, ();",
+            out(reg32) x,
+            out(reg32) y
+        );
+    }
+    Vec3::new(x, y, 1.0 - x - y)
+}

crates/optix_device/src/intersect.rs

@@ -0,0 +1,130 @@
+use cuda_std::gpu_only;
+use paste::paste;
+use seq_macro::seq;
+
+pub trait IntersectionPayload {
+    fn report_intersection(hit_t: f32, hit_kind: u8, payload: Self) -> bool;
+}
+
+macro_rules! impl_intersection_payload {
+    ($num:tt) => {
+        paste! {
+            impl IntersectionPayload for [u32; $num] {
+                #[gpu_only]
+                fn report_intersection(
+                    hit_t: f32,
+                    hit_kind: u8,
+                    payload: Self,
+                ) -> bool {
+                    seq!(
+                        P in 0..$num {{
+                            let [#(p~P,)*] = payload;
+                            let out: u32;
+                            unsafe {
+                                asm!(
+                                    concat!("call ({}), _optix_report_intersection_", stringify!($num)),
+                                    concat!(", ({}, {}, ", #(concat!("{", stringify!(p~P), "},"),)* ");"),
+                                    out(reg32) out,
+                                    in(reg32) hit_t,
+                                    in(reg32) hit_kind,
+                                    #(
+                                        p~P = in(reg32) p~P,
+                                    )*
+                                );
+                            }
+                            out != 0
+                        }}
+                    )
+                }
+            }
+        }
+    };
+    () => {
+        seq! {
+            N in 0..=7 {
+                impl_intersection_payload! { N }
+            }
+        }
+    }
+}
+
+impl_intersection_payload! {}
+
+/// Reports an intersection and passes custom attributes to further programs.
+///
+/// If `tmin <= hit_t <= tmax` then the anyhit program associated with this intersection will be invoked,
+/// then the program will do one of three things:
+/// - Ignore the intersection; no hit is recorded and this function returns `false`.
+/// - Terminate the ray; a hit is recorded and this function does not return. No further traversal occurs and the associated
+/// closesthit program is invoked.
+/// - Neither; A hit is recorded and this function returns `true`.
+///
+/// **Only the lower 7 bits of the `hit_kind` should be written, the top 127 values are reserved for hardware primitives.**
+pub fn report_intersection<P: IntersectionPayload>(hit_t: f32, hit_kind: u8, payload: P) -> bool {
+    P::report_intersection(hit_t, hit_kind, payload)
+}
+
+/// Records the hit, stops traversal, then proceeds to the closesthit program.
+#[gpu_only]
+pub fn terminate_ray() {
+    unsafe {
+        asm!("call _optix_terminate_ray, ();");
+    }
+}
+
+/// Discards the hit and returns control to the calling intersection program or the built-in intersection hardware.
+#[gpu_only]
+pub fn ignore_intersection() {
+    unsafe {
+        asm!("call _optix_ignore_intersection, ();");
+    }
+}
+
+macro_rules! get_attribute_fns {
+    ($num:tt) => {
+        paste! {
+            #[gpu_only]
+            #[allow(clippy::missing_safety_doc)]
+            unsafe fn [<get_attribute_ $num>]() -> u32 {
+                let out: u32;
+                asm!(
+                    concat!("call ({}), _optix_get_attribute_", stringify!($num), ", ();"),
+                    out(reg32) out
+                );
+                out
+            }
+        }
+    };
+    () => {
+        seq! {
+            N in 0..=7 {
+                get_attribute_fns! { N }
+            }
+        }
+    };
+}
+
+get_attribute_fns! {}
+
+/// Retrieves an attribute set by the intersection program when reporting an intersection.
+///
+/// # Safety
+///
+/// The attribute must have been set by the intersection program.
+///
+/// # Panics
+///
+/// Panics if the idx is over `7`.
+pub unsafe fn get_attribute(idx: u8) -> u32 {
+    match idx {
+        0 => get_attribute_0(),
+        1 => get_attribute_1(),
+        2 => get_attribute_2(),
+        3 => get_attribute_3(),
+        4 => get_attribute_4(),
+        5 => get_attribute_5(),
+        6 => get_attribute_6(),
+        7 => get_attribute_7(),
+        _ => panic!("Invalid attribute index"),
+    }
+}

crates/optix_device/src/lib.rs

@@ -7,8 +7,11 @@
 
 extern crate alloc;
 
+mod hit;
+mod intersect;
 pub mod misc;
 pub mod payload;
+mod ray;
 pub mod sys;
 pub mod trace;
 pub mod util;
@@ -53,6 +56,48 @@ pub fn get_launch_dimensions() -> UVec3 {
     UVec3::new(x, y, z)
 }
 
+/// Functions/items only available in raygen programs (`__raygen__`).
 pub mod raygen {
+    #[doc(inline)]
     pub use crate::trace::*;
 }
+
+/// Functions/items only available in miss programs (`__miss__`).
+pub mod intersection {
+    #[doc(inline)]
+    pub use crate::intersect::{get_attribute, report_intersection};
+    #[doc(inline)]
+    pub use crate::ray::*;
+}
+
+/// Functions/items only available in anyhit programs (`__anyhit__`).
+pub mod anyhit {
+    #[doc(inline)]
+    pub use crate::hit::*;
+    #[doc(inline)]
+    pub use crate::intersect::{get_attribute, ignore_intersection, terminate_ray};
+    #[doc(inline)]
+    pub use crate::ray::*;
+}
+
+/// Functions/items only available in closesthit programs (`__closesthit__`).
+pub mod closesthit {
+    #[doc(inline)]
+    pub use crate::hit::*;
+    #[doc(inline)]
+    pub use crate::intersect::get_attribute;
+    #[doc(inline)]
+    pub use crate::ray::{
+        ray_flags, ray_time, ray_tmax, ray_tmin, ray_visibility_mask, ray_world_direction,
+        ray_world_origin,
+    };
+}
+
+/// Functions/items only available in miss programs (`__miss__`).
+pub mod miss {
+    #[doc(inline)]
+    pub use crate::ray::{
+        ray_flags, ray_time, ray_tmax, ray_tmin, ray_visibility_mask, ray_world_direction,
+        ray_world_origin,
+    };
+}