Feat: rust OptiX hardware raytracing MVP

Author: RDambrosio016

Cargo.toml

@@ -3,6 +3,8 @@ members = [
   "crates/*",
   "crates/optix/examples/ex*",
   "crates/optix/examples/ex*/device",
+  "crates/optix/examples/rust/ex*",
+
   "xtask",
 
   "examples/optix/*",

crates/optix/examples/ex04_mesh/Cargo.toml

@@ -16,3 +16,4 @@ glam = { version = "0.20", features=["cuda"] }
 
 [build-dependencies]
 find_cuda_helper = { version = "0.2", path = "../../../find_cuda_helper" }
+cuda_builder = { version = "0.2", path = "../../../cuda_builder" }

crates/optix/examples/ex04_mesh/build.rs

@@ -1,42 +1,8 @@
-use find_cuda_helper::find_optix_root;
+use cuda_builder::CudaBuilder;
 
 fn main() {
-    let manifest_dir = std::env::var("CARGO_MANIFEST_DIR").unwrap();
-
-    let mut optix_include = find_optix_root().expect(
-        "Unable to find the OptiX SDK, make sure you installed it and
-    that OPTIX_ROOT or OPTIX_ROOT_DIR are set",
-    );
-    optix_include = optix_include.join("include");
-
-    let args = vec![
-        format!("-I{}", optix_include.display()),
-        format!("-I{}/../common/gdt", manifest_dir),
-    ];
-
-    compile_to_ptx("src/ex04_mesh.cu", &args);
-}
-
-fn compile_to_ptx(cu_path: &str, args: &[String]) {
-    println!("cargo:rerun-if-changed={}", cu_path);
-
-    let full_path =
-        std::path::PathBuf::from(std::env::var("CARGO_MANIFEST_DIR").unwrap()).join(cu_path);
-
-    let mut ptx_path = std::path::PathBuf::from(std::env::var("OUT_DIR").unwrap()).join(cu_path);
-    ptx_path.set_extension("ptx");
-    std::fs::create_dir_all(ptx_path.parent().unwrap()).unwrap();
-
-    let output = std::process::Command::new("nvcc")
-        .arg("-ptx")
-        .arg(&full_path)
-        .arg("-o")
-        .arg(&ptx_path)
-        .args(args)
-        .output()
-        .expect("failed to fun nvcc");
-
-    if !output.status.success() {
-        panic!("{}", unsafe { String::from_utf8_unchecked(output.stderr) });
-    }
+    CudaBuilder::new("../rust/ex04_mesh_gpu")
+        .copy_to("../resources/ex04_mesh.ptx")
+        .build()
+        .unwrap();
 }

crates/optix/examples/ex04_mesh/src/ex04_mesh.cu

@@ -63,7 +63,7 @@ impl Renderer {
             .traversable_graph_flags(TraversableGraphFlags::ALLOW_SINGLE_GAS)
             .exception_flags(ExceptionFlags::NONE);
 
-        let ptx = include_str!(concat!(env!("OUT_DIR"), "/src/ex04_mesh.ptx"));
+        let ptx = include_str!("../../resources/ex04_mesh.ptx");
 
         let (module, _log) = Module::new(
             &mut ctx,

crates/optix/examples/ex04_mesh/src/renderer.rs

@@ -0,0 +1 @@
+*.ptx

crates/optix/examples/resources/.gitignore

@@ -0,0 +1,11 @@
+[package]
+name = "ex04_mesh_gpu"
+version = "0.1.0"
+edition = "2021"
+
+[dependencies]
+cuda_std = { version = "0.2.0", path = "../../../../cuda_std" }
+optix_device = { version = "0.1.0", path = "../../../../optix_device" }
+
+[lib]
+crate-type = ["cdylib", "rlib"]

crates/optix/examples/rust/ex04_mesh_gpu/Cargo.toml

@@ -0,0 +1,123 @@
+#![cfg_attr(
+    target_os = "cuda",
+    no_std,
+    feature(register_attr),
+    register_attr(nvvm_internal)
+)]
+#![allow(non_snake_case, clippy::missing_safety_doc)]
+
+use cuda_std::kernel;
+use optix_device::{
+    get_launch_index,
+    glam::*,
+    misc::*,
+    payload,
+    trace::TraversableHandle,
+    trace::{trace, RayFlags},
+    util::*,
+};
+
+#[derive(Debug, Clone, Copy)]
+#[repr(C)]
+pub struct LaunchParams {
+    pub frame: Frame,
+    pub camera: Camera,
+    pub traversable: TraversableHandle,
+}
+
+#[derive(Debug, Clone, Copy)]
+#[repr(C)]
+pub struct Frame {
+    pub color_buf: *mut Vec4,
+    pub size: UVec2,
+}
+
+#[derive(Debug, Clone, Copy)]
+#[repr(C)]
+pub struct Camera {
+    pub position: Vec3,
+    pub direction: Vec3,
+    pub horizontal: Vec3,
+    pub vertical: Vec3,
+}
+
+#[repr(u32)]
+#[derive(Debug, Clone, Copy)]
+pub enum RayType {
+    SurfaceRay = 0,
+}
+
+fn get_color_buf() -> *mut Vec4 {
+    unsafe { unpack_pointer(payload::get_payload(0), payload::get_payload(1)) }
+}
+
+fn random_color(i: u32) -> Vec4 {
+    let r = i * 13 * 17 + 0x234235;
+    let g = i * 7 * 3 * 5 + 0x773477;
+    let b = i * 11 * 19 + 0x223766;
+    Vec4::new(
+        (r & 255) as f32 / 255.0,
+        (g & 255) as f32 / 255.0,
+        (b & 255) as f32 / 255.0,
+        1.0,
+    )
+}
+
+#[kernel]
+pub unsafe fn __closesthit__radiance() {
+    let prim_id = primitive_index();
+    let buf = get_color_buf();
+    *buf = random_color(prim_id);
+}
+
+#[kernel]
+pub unsafe fn __anyhit__radiance() {}
+
+#[kernel]
+pub unsafe fn __miss__radiance() {
+    let buf = get_color_buf();
+    // pure white background
+    *buf = Vec4::ONE;
+}
+
+extern "C" {
+    #[cfg_attr(target_os = "cuda", nvvm_internal(addrspace(4)))]
+    static PARAMS: LaunchParams;
+}
+
+#[kernel]
+pub unsafe fn __raygen__renderFrame() {
+    let i = get_launch_index();
+    let i = UVec2::new(i.x, i.y);
+
+    let camera = PARAMS.camera;
+
+    let px_color = Vec3::ZERO;
+    let (mut p0, mut p1) = pack_pointer(&px_color as *const _ as *mut Vec3);
+
+    let screen = (i.as_vec2() + Vec2::splat(0.5)) / PARAMS.frame.size.as_vec2();
+    let ray_dir = (camera.direction
+        + (screen.x - 0.5) * camera.horizontal
+        + (screen.y - 0.5) * camera.vertical)
+        .normalize();
+
+    trace(
+        PARAMS.traversable,
+        camera.position,
+        ray_dir,
+        0.0,
+        1e20,
+        0.0,
+        255,
+        RayFlags::DISABLE_ANYHIT,
+        RayType::SurfaceRay as u32,
+        1,
+        RayType::SurfaceRay as u32,
+        &mut p0,
+        &mut p1,
+    );
+
+    let fb_index = i.x + i.y * PARAMS.frame.size.x;
+    *PARAMS.frame.color_buf.add(fb_index as usize) =
+        Vec4::new(px_color.x, px_color.y, px_color.z, 1.0);
+}

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

@@ -5,5 +5,7 @@ edition = "2021"
 authors = ["Anders Langlands <[email protected]>", "Riccardo D'Ambrosio <[email protected]>"]
 
 [dependencies]
+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"

crates/optix_device/Cargo.toml

@@ -1,13 +1,20 @@
 #![cfg_attr(
     target_arch = "nvptx64",
     no_std,
-    feature(register_attr, asm),
+    feature(register_attr, asm, asm_experimental_arch),
     register_attr(nvvm_internal)
 )]
 
 extern crate alloc;
 
+pub mod misc;
+pub mod payload;
+pub mod sys;
+pub mod trace;
+pub mod util;
+
 use cuda_std::*;
+pub use glam;
 use glam::UVec3;
 
 extern "C" {
@@ -45,3 +52,7 @@ pub fn get_launch_dimensions() -> UVec3 {
 
     UVec3::new(x, y, z)
 }
+
+pub mod raygen {
+    pub use crate::trace::*;
+}