Replace CUDA C kernels with CUDA Oxide Rust

.github/workflows/ci.yml

@@ -20,7 +20,9 @@ jobs:
       - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1
       - run: rustup component add rustfmt clippy
       - run: python3 tools/test_corpus_build.py
-      - run: cargo fmt --all -- --check
+      - run: |
+          cargo fmt --all -- --check
+          rustfmt --edition 2024 --check crates/jxr-cuda/src/kernels/main.rs
       - run: cargo clippy --workspace --all-targets --all-features -- -D warnings
       - run: cargo test --workspace --all-features
       - run: cargo test -p jxr --no-default-features

.github/workflows/cuda-hardware.yml

@@ -15,30 +15,36 @@ jobs:
     if: ${{ github.event_name == 'workflow_dispatch' && github.ref == 'refs/heads/main' && github.actor == 'jcwal1516' && github.triggering_actor == 'jcwal1516' }}
     runs-on: [self-hosted, Linux, X64, cuda]
     timeout-minutes: 180
+    env:
+      JXR_REQUIRE_CUDA_OXIDE: "1"
     steps:
       - name: Check out repository
         uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1
       - name: Record CUDA environment
         run: |
           mkdir -p target/t834-conformance/reports
           nvidia-smi | tee target/t834-conformance/reports/nvidia-smi.txt
-          nvcc --version | tee target/t834-conformance/reports/nvcc-version.txt
+          cargo oxide --version | tee target/t834-conformance/reports/cuda-oxide-version.txt
       - name: Install Rust checks
         run: rustup component add rustfmt clippy
       - name: Prepare T.834 vectors
         run: tools/t834-conformance/build.sh
       - name: Build T.835 oracle
         run: tools/t835-oracle/build.sh
       - name: Check formatting
-        run: cargo fmt --all -- --check
+        run: |
+          cargo fmt --all -- --check
+          rustfmt --edition 2024 --check crates/jxr-cuda/src/kernels/main.rs
       - name: Lint all targets and features
         run: cargo clippy --workspace --all-targets --all-features -- -D warnings
       - name: Compile CUDA backend without static toolkit linkage
         run: cargo check -p jxr-cuda --all-targets
       - name: Test full workspace
         run: cargo test --workspace --all-features
       - name: Run ignored CUDA lifecycle and differential tests
-        run: cargo test -p jxr --features cuda --test cuda_contracts -- --ignored --test-threads=1
+        run: |
+          cargo test -p jxr-cuda --lib -- --ignored --test-threads=1
+          cargo test -p jxr --features cuda --test cuda_contracts -- --ignored --test-threads=1
       - name: Differentially test the full CUDA corpus and ROI boundaries against CPU
         run: cargo test -p jxr-test-support --features cuda --test cuda_hardware -- --ignored --test-threads=1
       - name: Run complete CUDA T.834/T.835 conformance
@@ -65,7 +71,7 @@ jobs:
           name: jpeg-xr-cuda-hardware-reports
           path: |
             target/t834-conformance/reports/nvidia-smi.txt
-            target/t834-conformance/reports/nvcc-version.txt
+            target/t834-conformance/reports/cuda-oxide-version.txt
             target/t834-conformance/reports/ci-cuda.tsv
             target/t834-conformance/reports/cuda-pathology-bench.tsv
           if-no-files-found: warn

README.md

@@ -194,7 +194,7 @@ buffers, caller-owned destinations, and batches.
 | --- | --- | --- |
 | CPU | Default; no features needed | No GPU runtime |
 | Metal | `cargo add jxr --features metal` | Apple-silicon macOS (M1 or newer) |
-| CUDA | `cargo add jxr --features cuda` | Compatible NVIDIA GPU, driver, and NVRTC; hardware validation currently covers Linux |
+| CUDA | `cargo add jxr --features cuda` | Compatible NVIDIA GPU and driver; CUDA Oxide kernels built on Linux (see guide) |
 
 For single-image GPU decoding, attach a reusable session to the decoder;
 enabling a feature alone does not attach one. See the

THIRD_PARTY_NOTICES.md

@@ -9,8 +9,13 @@ must retain the applicable BSD-3-Clause copyright and disclaimer here.
 `fearless_simd` 1.0.0 is used through its safe capability-token API and is
 available under MIT or Apache-2.0 licensing.
 
-`cudarc` 0.19.10 provides dynamically loaded CUDA Driver API and NVRTC bindings
+`cudarc` 0.19.10 provides dynamically loaded CUDA Driver API bindings
 for the optional `jxr-cuda` crate and is available under MIT or Apache-2.0
 licensing. No NVIDIA toolkit code or headers are distributed by this repository.
 Its dynamic loader dependency, `libloading` 0.9.0, is available under the ISC
 license; `cfg-if` and the Windows loader shim are MIT OR Apache-2.0.
+
+CUDA Oxide supplies the Rust device compiler and device/runtime crates used
+only when building the optional CUDA kernels. The pinned upstream revision is
+Apache-2.0 licensed. JXR uses the upstream APIs; no compiler or toolkit sources
+are vendored in this repository.

crates/jxr-cuda/build.rs

@@ -0,0 +1,96 @@
+// SPDX-License-Identifier: MIT OR Apache-2.0
+
+use std::{env, fs, io, path::PathBuf, process::Command};
+
+fn main() -> io::Result<()> {
+    println!("cargo:rustc-check-cfg=cfg(jxr_cuda_oxide)");
+    println!("cargo:rerun-if-env-changed=JXR_CUDA_OXIDE_ARCH");
+    println!("cargo:rerun-if-env-changed=JXR_CUDA_OXIDE_PTX");
+    println!("cargo:rerun-if-env-changed=JXR_REQUIRE_CUDA_OXIDE");
+    println!("cargo:rerun-if-changed=src/kernels");
+    let math =
+        PathBuf::from(env::var_os("DEP_JXR_MATH_MANIFEST_DIR").ok_or_else(|| {
+            io::Error::other("Cargo did not provide the jxr-math source directory")
+        })?);
+    println!("cargo:rerun-if-changed={}", math.display());
+    let manifest = PathBuf::from(env::var_os("CARGO_MANIFEST_DIR").unwrap());
+    let out = PathBuf::from(env::var_os("OUT_DIR").unwrap());
+    let ptx = out.join("jxr_reconstruction.ptx");
+    if let Some(prebuilt) = env::var_os("JXR_CUDA_OXIDE_PTX") {
+        println!(
+            "cargo:rerun-if-changed={}",
+            PathBuf::from(&prebuilt).display()
+        );
+        fs::copy(prebuilt, &ptx)?;
+    } else {
+        let available = env::var("HOST").unwrap().contains("linux")
+            && Command::new("cargo-oxide")
+                .arg("--help")
+                .output()
+                .is_ok_and(|r| r.status.success());
+        if !available {
+            if env::var_os("JXR_REQUIRE_CUDA_OXIDE").is_some() {
+                return Err(io::Error::other(
+                    "CUDA Oxide requires a Linux build host with cargo-oxide installed, or JXR_CUDA_OXIDE_PTX",
+                ));
+            }
+            println!(
+                "cargo:warning=CUDA Oxide kernels unavailable; install cargo-oxide on Linux or set JXR_CUDA_OXIDE_PTX to enable CUDA execution"
+            );
+            fs::write(ptx, "")?;
+            return Ok(());
+        }
+        let project = out.join("cuda-oxide");
+        fs::create_dir_all(project.join("simt/src"))?;
+        fs::create_dir_all(project.join("src"))?;
+        fs::write(project.join("src/main.rs"), "fn main() {}\n")?;
+        for name in ["Cargo.toml", "rust-toolchain.toml"] {
+            fs::copy(
+                manifest.join("src/kernels").join(if name == "Cargo.toml" {
+                    "Cargo.toml.in"
+                } else {
+                    name
+                }),
+                project.join(name),
+            )?;
+        }
+        let template = fs::read_to_string(manifest.join("src/kernels/simt.toml"))?;
+
+        fs::write(
+            project.join("simt/Cargo.toml"),
+            template.replace("__JXR_MATH__", math.to_str().unwrap()),
+        )?;
+        for name in ["main.rs", "reconstruction.rs", "output.rs"] {
+            fs::copy(
+                manifest.join("src/kernels").join(name),
+                project.join("simt/src").join(name),
+            )?;
+        }
+        let arch = env::var("JXR_CUDA_OXIDE_ARCH").unwrap_or_else(|_| "sm_50".into());
+        let status = Command::new("cargo")
+            .args(["oxide", "build", "--arch", &arch, "--no-fmad"])
+            .current_dir(&project)
+            .env_remove("RUSTUP_TOOLCHAIN")
+            .env_remove("RUSTC")
+            .env_remove("RUSTC_WRAPPER")
+            .env_remove("RUSTC_WORKSPACE_WRAPPER")
+            .env_remove("RUSTDOC")
+            .env_remove("RUSTFLAGS")
+            .env_remove("CARGO_ENCODED_RUSTFLAGS")
+            .env("CARGO_TARGET_DIR", project.join("target"))
+            .env("JXR_MATH_DEVICE", "1")
+            .status()?;
+        if !status.success() {
+            return Err(io::Error::other("CUDA Oxide kernel compilation failed"));
+        }
+        fs::copy(project.join("ptx/jxr_reconstruction.ptx"), &ptx)?;
+    }
+    let source = fs::read_to_string(ptx)?;
+    if !source.contains(".entry jxr_dequantize_first_transform") {
+        return Err(io::Error::other(
+            "CUDA Oxide artifact omits reconstruction kernels",
+        ));
+    }
+    println!("cargo:rustc-cfg=jxr_cuda_oxide");
+    Ok(())
+}

crates/jxr-cuda/src/abi.rs

@@ -56,7 +56,7 @@ pub(crate) struct JxrPlaneAbi {
 }
 
 // SAFETY: Twelve consecutive u32 fields have the asserted layout and every
-// bit pattern is valid for the matching CUDA C structure.
+// bit pattern is valid for the matching CUDA Oxide structure.
 unsafe impl DeviceRepr for JxrPlaneAbi {}
 // SAFETY: Every field is a u32, so an all-zero value is valid.
 unsafe impl ValidAsZeroBits for JxrPlaneAbi {}

crates/jxr-cuda/src/encode.rs

@@ -246,7 +246,7 @@ fn launch_second_transforms(
             launch.launch(LaunchConfig {
                 grid_dim: (abi.macroblock_count, 1, 1),
                 block_dim: (16, 1, 1),
-                shared_mem_bytes: 256 * 4,
+                shared_mem_bytes: 0,
             })?;
         }
     }

crates/jxr-cuda/src/error.rs

@@ -6,7 +6,7 @@ use jxr_core::BackendRequest;
 #[derive(Debug, thiserror::Error)]
 #[non_exhaustive]
 pub enum CudaError {
-    /// The CUDA driver or NVRTC runtime is unavailable on this host.
+    /// The CUDA driver or compiled CUDA Oxide kernels are unavailable.
     #[error("CUDA is unavailable on this host")]
     Unavailable,
     /// The selected NVIDIA device does not meet the backend baseline.
@@ -42,10 +42,10 @@ pub enum CudaError {
         /// Stable description of the rejected representation.
         reason: &'static str,
     },
-    /// Runtime kernel compilation or module creation failed.
+    /// Device module creation failed.
     #[error("CUDA pipeline initialization failed: {message}")]
     RuntimeInitialization {
-        /// NVRTC or driver diagnostic.
+        /// CUDA driver diagnostic.
         message: String,
     },
     /// Checked device arithmetic exceeded the scalar contract.

crates/jxr-cuda/src/kernels.rs

@@ -20,7 +20,7 @@ pub enum KernelStage {
 pub struct CudaKernelManifest {
     /// Reconstruction phase.
     pub stage: KernelStage,
-    /// Packaged CUDA C entry point.
+    /// Packaged CUDA Oxide entry point.
     pub entrypoint: &'static str,
 }
 
@@ -48,50 +48,4 @@ pub const RECONSTRUCTION_KERNELS: [CudaKernelManifest; 5] = [
     },
 ];
 
-pub(crate) fn source() -> String {
-    [
-        jxr_math::tables::CUDA_RECONSTRUCTION_CONSTANTS,
-        include_str!("kernels/common.cu"),
-        include_str!("kernels/first_transform.cu"),
-        include_str!("kernels/overlap.cu"),
-        include_str!("kernels/highpass_transform.cu"),
-        include_str!("kernels/output.cu"),
-    ]
-    .join("\n")
-}
-
-#[cfg(test)]
-mod tests {
-    use super::source;
-
-    #[test]
-    fn packages_every_runtime_entrypoint_without_metal_syntax() {
-        let source = source();
-        for entrypoint in [
-            "jxr_dequantize_first_transform",
-            "jxr_first_overlap",
-            "jxr_highpass_second_transform",
-            "jxr_second_overlap",
-            "jxr_output_bits",
-            "jxr_output_u8",
-            "jxr_output_u16",
-            "jxr_output_i16",
-            "jxr_output_i32",
-            "jxr_output_f16",
-            "jxr_output_f32",
-            "jxr_output_packed16",
-            "jxr_output_packed32",
-        ] {
-            assert!(source.contains(entrypoint), "missing {entrypoint}");
-        }
-        assert!(!source.contains("[[buffer("));
-        assert!(!source.contains("metal_stdlib"));
-    }
-
-    #[test]
-    fn kernel_entry_status_is_broadcast_once_per_thread_block() {
-        let source = source();
-        assert!(source.contains("bool jxr_block_failed(uint *status)"));
-        assert!(source.contains("__shared__ uint jxr_block_status"));
-    }
-}
+pub(crate) const PTX: &str = include_str!(concat!(env!("OUT_DIR"), "/jxr_reconstruction.ptx"));

crates/jxr-cuda/src/kernels/Cargo.toml.in

@@ -0,0 +1,12 @@
+[package]
+name = "jxr-cuda-oxide-host"
+version = "0.1.0"
+edition = "2024"
+
+[[package.metadata.cuda-oxide.device-crates]]
+manifest-path = "simt/Cargo.toml"
+ptx-dir = "ptx"
+artifact-name = "jxr_reconstruction"
+
+[workspace]
+resolver = "2"