diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 602cabb..254e91a 100644 --- a/.github/workflows/ci.yml +++ b/.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 diff --git a/.github/workflows/cuda-hardware.yml b/.github/workflows/cuda-hardware.yml index 497e586..35c4ef7 100644 --- a/.github/workflows/cuda-hardware.yml +++ b/.github/workflows/cuda-hardware.yml @@ -15,6 +15,8 @@ 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 @@ -22,7 +24,7 @@ jobs: 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 @@ -30,7 +32,9 @@ jobs: - 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 @@ -38,7 +42,9 @@ jobs: - 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 diff --git a/README.md b/README.md index 569aa83..a76c3d5 100644 --- a/README.md +++ b/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 diff --git a/THIRD_PARTY_NOTICES.md b/THIRD_PARTY_NOTICES.md index dc2ef39..f6c7053 100644 --- a/THIRD_PARTY_NOTICES.md +++ b/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. diff --git a/crates/jxr-cuda/build.rs b/crates/jxr-cuda/build.rs new file mode 100644 index 0000000..726ca56 --- /dev/null +++ b/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(()) +} diff --git a/crates/jxr-cuda/src/abi.rs b/crates/jxr-cuda/src/abi.rs index 3162156..ffd31c4 100644 --- a/crates/jxr-cuda/src/abi.rs +++ b/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 {} diff --git a/crates/jxr-cuda/src/encode.rs b/crates/jxr-cuda/src/encode.rs index ebaa743..23daa01 100644 --- a/crates/jxr-cuda/src/encode.rs +++ b/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, })?; } } diff --git a/crates/jxr-cuda/src/error.rs b/crates/jxr-cuda/src/error.rs index 9294f9b..3839172 100644 --- a/crates/jxr-cuda/src/error.rs +++ b/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. diff --git a/crates/jxr-cuda/src/kernels.rs b/crates/jxr-cuda/src/kernels.rs index 6bc9b81..7faa720 100644 --- a/crates/jxr-cuda/src/kernels.rs +++ b/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")); diff --git a/crates/jxr-cuda/src/kernels/Cargo.toml.in b/crates/jxr-cuda/src/kernels/Cargo.toml.in new file mode 100644 index 0000000..2480063 --- /dev/null +++ b/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" diff --git a/crates/jxr-cuda/src/kernels/common.cu b/crates/jxr-cuda/src/kernels/common.cu deleted file mode 100644 index 3c6e2f3..0000000 --- a/crates/jxr-cuda/src/kernels/common.cu +++ /dev/null @@ -1,170 +0,0 @@ -// SPDX-License-Identifier: MIT OR Apache-2.0 - -typedef unsigned int uint; -typedef unsigned short ushort; -typedef unsigned char uchar; -typedef unsigned long long ulong; -typedef long long jxr_long; - -#define JXR_INT_MIN (-2147483647 - 1) -#define JXR_INT_MAX 2147483647 -#define JXR_UINT_MAX 0xffffffffu - -struct JxrUint4 { uint x; uint y; uint z; uint w; }; - -__device__ __forceinline__ JxrUint4 jxr_uint4(uint x, uint y, uint z, uint w) { - JxrUint4 result = {x, y, z, w}; - return result; -} - -__device__ __forceinline__ int jxr_min_i32(int a, int b) { return a < b ? a : b; } -__device__ __forceinline__ int jxr_max_i32(int a, int b) { return a > b ? a : b; } -__device__ __forceinline__ uint jxr_min_u32(uint a, uint b) { return a < b ? a : b; } -__device__ __forceinline__ uint jxr_max_u32(uint a, uint b) { return a > b ? a : b; } -__device__ __forceinline__ int jxr_clamp_i32(int value, int low, int high) { - return jxr_min_i32(jxr_max_i32(value, low), high); -} -__device__ __forceinline__ float jxr_clamp_f32(float value, float low, float high) { - return value < low ? low : (value > high ? high : value); -} - -__device__ __forceinline__ bool jxr_failed(uint *status) { - return atomicAdd(status, 0u) != 0u; -} - -__device__ __forceinline__ bool jxr_block_failed(uint *status) { - __shared__ uint jxr_block_status; - const uint linear_thread = threadIdx.x - + blockDim.x * (threadIdx.y + blockDim.y * threadIdx.z); - if (linear_thread == 0u) jxr_block_status = (uint)jxr_failed(status); - __syncthreads(); - return jxr_block_status != 0u; -} - -__device__ __forceinline__ void jxr_fail(uint *status, uint code) { - atomicCAS(status, 0u, code); -} - -__device__ __forceinline__ bool jxr_narrow( - jxr_long value, int &result, uint *status, uint code) { - if (value < (jxr_long)JXR_INT_MIN || value > (jxr_long)JXR_INT_MAX) { - jxr_fail(status, code); - return false; - } - result = (int)value; - return true; -} - -__device__ __forceinline__ bool jxr_add( - int a, int b, int &result, uint *status, uint code) { - return jxr_narrow((jxr_long)a + (jxr_long)b, result, status, code); -} - -__device__ __forceinline__ bool jxr_sub( - int a, int b, int &result, uint *status, uint code) { - return jxr_narrow((jxr_long)a - (jxr_long)b, result, status, code); -} - -__device__ __forceinline__ bool jxr_mul( - int a, uint b, int &result, uint *status, uint code) { - return jxr_narrow((jxr_long)a * (jxr_long)b, result, status, code); -} - -__device__ __forceinline__ bool jxr_mul3_round( - int value, int rounding, uint shift, int &result, uint *status, uint code) { - int product; - int rounded; - return jxr_narrow((jxr_long)value * 3ll, product, status, code) - && jxr_add(product, rounding, rounded, status, code) - && (result = rounded >> shift, true); -} - -__device__ __forceinline__ bool jxr_t2x2h( - int v[4], int rounding, uint *status, uint code) { - int difference; - int midpoint; - int old2 = v[2]; - return jxr_add(v[0], v[3], v[0], status, code) - && jxr_sub(v[1], v[2], v[1], status, code) - && jxr_sub(v[0], v[1], difference, status, code) - && jxr_add(difference, rounding, midpoint, status, code) - && (midpoint >>= 1, true) - && jxr_sub(midpoint, v[3], v[2], status, code) - && jxr_sub(midpoint, old2, v[3], status, code) - && jxr_sub(v[0], v[3], v[0], status, code) - && jxr_add(v[1], v[2], v[1], status, code); -} - -__device__ __forceinline__ bool jxr_todd(int v[4], uint *status, uint code) { - int temporary; - return jxr_add(v[1], v[3], v[1], status, code) - && jxr_sub(v[0], v[2], v[0], status, code) - && jxr_sub(v[3], v[1] >> 1, v[3], status, code) - && jxr_add(v[0], 1, temporary, status, code) - && jxr_add(v[2], temporary >> 1, v[2], status, code) - && jxr_mul3_round(v[1], 4, 3, temporary, status, code) - && jxr_sub(v[0], temporary, v[0], status, code) - && jxr_mul3_round(v[0], 4, 3, temporary, status, code) - && jxr_add(v[1], temporary, v[1], status, code) - && jxr_mul3_round(v[3], 4, 3, temporary, status, code) - && jxr_sub(v[2], temporary, v[2], status, code) - && jxr_mul3_round(v[2], 4, 3, temporary, status, code) - && jxr_add(v[3], temporary, v[3], status, code) - && jxr_add(v[1], 1, temporary, status, code) - && jxr_sub(v[2], temporary >> 1, v[2], status, code) - && jxr_add(v[0], 1, temporary, status, code) - && jxr_sub(temporary >> 1, v[3], v[3], status, code) - && jxr_add(v[1], v[2], v[1], status, code) - && jxr_sub(v[0], v[3], v[0], status, code); -} - -__device__ __forceinline__ bool jxr_todd_odd(int v[4], uint *status, uint code) { - int temporary; - int first; - int second; - return jxr_add(v[3], v[0], v[3], status, code) - && jxr_sub(v[2], v[1], v[2], status, code) - && (first = v[3] >> 1, second = v[2] >> 1, true) - && jxr_sub(v[0], first, v[0], status, code) - && jxr_add(v[1], second, v[1], status, code) - && jxr_mul3_round(v[1], 3, 3, temporary, status, code) - && jxr_sub(v[0], temporary, v[0], status, code) - && jxr_mul3_round(v[0], 3, 2, temporary, status, code) - && jxr_add(v[1], temporary, v[1], status, code) - && jxr_mul3_round(v[1], 4, 3, temporary, status, code) - && jxr_sub(v[0], temporary, v[0], status, code) - && jxr_sub(v[1], second, v[1], status, code) - && jxr_add(v[0], first, v[0], status, code) - && jxr_add(v[2], v[1], v[2], status, code) - && jxr_sub(v[3], v[0], v[3], status, code) - && jxr_narrow(-(jxr_long)v[1], v[1], status, code) - && jxr_narrow(-(jxr_long)v[2], v[2], status, code); -} - -__device__ __forceinline__ bool jxr_transform_group( - int c[16], JxrUint4 indices, uint kind, int rounding, uint *status, uint code) { - int v[4] = {c[indices.x], c[indices.y], c[indices.z], c[indices.w]}; - bool ok = kind == 0u ? jxr_t2x2h(v, rounding, status, code) - : (kind == 1u ? jxr_todd(v, status, code) : jxr_todd_odd(v, status, code)); - if (!ok) return false; - c[indices.x] = v[0]; - c[indices.y] = v[1]; - c[indices.z] = v[2]; - c[indices.w] = v[3]; - return true; -} - -__device__ __forceinline__ bool jxr_inverse_transform( - int c[16], uint *status, uint code) { - int input[16]; - for (uint i = 0; i < 16; ++i) input[i] = c[i]; - for (uint i = 0; i < 16; ++i) c[JXR_INVERSE_PERMUTATION[i]] = input[i]; - return jxr_transform_group(c, jxr_uint4(0, 1, 4, 5), 0, 1, status, code) - && jxr_transform_group(c, jxr_uint4(2, 3, 6, 7), 1, 0, status, code) - && jxr_transform_group(c, jxr_uint4(8, 12, 9, 13), 1, 0, status, code) - && jxr_transform_group(c, jxr_uint4(10, 11, 14, 15), 2, 0, status, code) - && jxr_transform_group(c, jxr_uint4(0, 3, 12, 15), 0, 0, status, code) - && jxr_transform_group(c, jxr_uint4(5, 6, 9, 10), 0, 0, status, code) - && jxr_transform_group(c, jxr_uint4(1, 2, 13, 14), 0, 0, status, code) - && jxr_transform_group(c, jxr_uint4(4, 7, 8, 11), 0, 0, status, code); -} diff --git a/crates/jxr-cuda/src/kernels/first_transform.cu b/crates/jxr-cuda/src/kernels/first_transform.cu deleted file mode 100644 index 4225f48..0000000 --- a/crates/jxr-cuda/src/kernels/first_transform.cu +++ /dev/null @@ -1,94 +0,0 @@ -// SPDX-License-Identifier: MIT OR Apache-2.0 - -__device__ __forceinline__ void jxr_dequantize_first_transform_one( - const int *packed, - const JxrMacroblockAbi *macroblocks, - int *low_plane, - uint *status, - JxrPlaneAbi plane, - uint gid) { - const bool failed = jxr_block_failed(status); - if (gid >= plane.macroblock_count || failed) return; - const uint metadata_index = plane.macroblock_offset + gid; - const JxrMacroblockAbi metadata = macroblocks[metadata_index]; - const uint block_count = plane.block_columns * plane.block_rows; - int low[16]; - for (uint i = 0; i < 16; ++i) low[i] = 0; - if (!jxr_mul( - packed[metadata.coefficient_offset], metadata.quantizer_dc, low[0], status, 1u)) { - return; - } - if (metadata.bands != 0u) { - for (uint block = 1; block < block_count; ++block) { - const uint source = metadata.bands >= 2u - ? metadata.coefficient_offset + block * 16u - : metadata.coefficient_offset + block; - if (!jxr_mul( - packed[source], metadata.quantizer_low_pass, low[block], status, 1u)) { - return; - } - } - } - if (plane.block_columns == 4u) { - if (!jxr_inverse_transform(low, status, 1u)) return; - } else if (plane.block_rows == 2u) { - int values[4] = {low[0], low[1], low[2], low[3]}; - if (!jxr_t2x2h(values, 0, status, 1u)) return; - low[0] = values[0]; - low[1] = values[2]; - low[2] = values[1]; - low[3] = values[3]; - } else { - int pair0; - int pair1; - int temporary; - if (!jxr_add(low[4], 1, temporary, status, 1u) - || !jxr_sub(low[0], temporary >> 1, pair0, status, 1u) - || !jxr_add(low[4], pair0, pair1, status, 1u)) { - return; - } - low[0] = pair0; - low[4] = pair1; - int first[4] = {low[0], low[1], low[2], low[3]}; - int second[4] = {low[4], low[6], low[5], low[7]}; - if (!jxr_t2x2h(first, 0, status, 1u) - || !jxr_t2x2h(second, 0, status, 1u)) { - return; - } - low[0] = first[0]; - low[1] = first[2]; - low[2] = first[1]; - low[3] = first[3]; - low[4] = second[0]; - low[5] = second[1]; - low[6] = second[2]; - low[7] = second[3]; - } - if (plane.scale_after_first_transform != 0u) { - for (uint block = 0; block < block_count; ++block) { - if (!jxr_narrow((jxr_long)low[block] * 2ll, low[block], status, 1u)) return; - } - } - const uint local_x = metadata.coded_x - plane.macroblock_origin_x; - const uint local_y = metadata.coded_y - plane.macroblock_origin_y; - const uint base_x = local_x * plane.block_columns; - const uint base_y = local_y * plane.block_rows; - for (uint row = 0; row < plane.block_rows; ++row) { - for (uint column = 0; column < plane.block_columns; ++column) { - low_plane[ - plane.low_offset + (base_y + row) * plane.low_width + base_x + column] = - low[row * plane.block_columns + column]; - } - } -} - -extern "C" __global__ void jxr_dequantize_first_transform( - const int *packed, - const JxrMacroblockAbi *macroblocks, - int *low_plane, - uint *status, - JxrPlaneAbi plane) { - const uint gid = blockIdx.x * blockDim.x + threadIdx.x; - jxr_dequantize_first_transform_one( - packed, macroblocks, low_plane, status, plane, gid); -} diff --git a/crates/jxr-cuda/src/kernels/highpass_transform.cu b/crates/jxr-cuda/src/kernels/highpass_transform.cu deleted file mode 100644 index 3564247..0000000 --- a/crates/jxr-cuda/src/kernels/highpass_transform.cu +++ /dev/null @@ -1,87 +0,0 @@ -// SPDX-License-Identifier: MIT OR Apache-2.0 - -extern "C" __global__ void jxr_highpass_second_transform( - const int *packed, - const JxrMacroblockAbi *macroblocks, - const int *low_plane, - int *samples, - uint *status, - JxrPlaneAbi plane) { - const uint group_id = blockIdx.x; - const uint tid = threadIdx.x; - const uint threads = blockDim.x; - const bool failed = jxr_block_failed(status); - if (group_id >= plane.macroblock_count || failed) return; - const uint metadata_index = plane.macroblock_offset + group_id; - const JxrMacroblockAbi metadata = macroblocks[metadata_index]; - const uint block_count = plane.block_columns * plane.block_rows; - const uint coefficient_count = block_count * 16u; - extern __shared__ int high[]; - for (uint index = tid; index < coefficient_count; index += threads) { - if (metadata.bands < 2u || (index & 15u) == 0u) { - high[index] = 0; - } else { - high[index] = packed[metadata.coefficient_offset + index]; - } - } - __syncthreads(); - if (tid == 0u && metadata.bands >= 2u) { - if (metadata.hp_prediction == 1u) { - for (uint block = 1; block < block_count; ++block) { - if ((block % plane.block_columns) != 0u) { - for (uint coefficient = 4; coefficient <= 12; coefficient += 4) { - const uint destination = block * 16u + coefficient; - int predicted; - if (!jxr_add( - high[destination], high[destination - 16u], - predicted, status, 3u)) { - break; - } - high[destination] = predicted; - } - } - } - } else if (metadata.hp_prediction == 2u) { - for (uint block = plane.block_columns; block < block_count; ++block) { - for (uint coefficient = 1; coefficient <= 3; ++coefficient) { - const uint destination = block * 16u + coefficient; - const uint source = destination - plane.block_columns * 16u; - int predicted; - if (!jxr_add(high[destination], high[source], predicted, status, 3u)) break; - high[destination] = predicted; - } - } - } - } - __syncthreads(); - if (jxr_block_failed(status)) return; - for (uint block = tid; block < block_count; block += threads) { - int coefficients[16]; - coefficients[0] = 0; - for (uint coefficient = 1; coefficient < 16; ++coefficient) { - if (!jxr_mul( - high[block * 16u + coefficient], metadata.quantizer_high_pass, - coefficients[coefficient], status, 4u)) { - return; - } - } - const uint local_x = metadata.coded_x - plane.macroblock_origin_x; - const uint local_y = metadata.coded_y - plane.macroblock_origin_y; - const uint low_x = local_x * plane.block_columns + block % plane.block_columns; - const uint low_y = local_y * plane.block_rows + block / plane.block_columns; - coefficients[0] = - low_plane[plane.low_offset + low_y * plane.low_width + low_x]; - if (!jxr_inverse_transform(coefficients, status, 4u)) return; - const uint output_x = - local_x * plane.block_columns * 4u + (block % plane.block_columns) * 4u; - const uint output_y = - local_y * plane.block_rows * 4u + (block / plane.block_columns) * 4u; - for (uint row = 0; row < 4; ++row) { - for (uint column = 0; column < 4; ++column) { - samples[ - plane.sample_offset + (output_y + row) * plane.sample_width - + output_x + column] = coefficients[row * 4u + column]; - } - } - } -} diff --git a/crates/jxr-cuda/src/kernels/main.rs b/crates/jxr-cuda/src/kernels/main.rs new file mode 100644 index 0000000..ade8f39 --- /dev/null +++ b/crates/jxr-cuda/src/kernels/main.rs @@ -0,0 +1,319 @@ +// SPDX-License-Identifier: MIT OR Apache-2.0 + +#![allow( + static_mut_refs, + reason = "CUDA shared arrays are scoped to a thread block" +)] + +use cuda_device::{ + SharedArray, + atomic::{AtomicOrdering, DeviceAtomicU32}, + kernel, thread, +}; +use cuda_host::cuda_module; +use jxr_math::tables::reconstruction_abi::*; +use jxr_math::{ + MathError, + arithmetic::{checked_add, checked_i64_to_i32, checked_mul, checked_sub}, +}; + +// CUDA Oxide cannot lower the nested ControlFlow layout produced by `?`. +macro_rules! math_try { + ($value:expr) => { + match $value { + Ok(value) => value, + Err(error) => return Err(error), + } + }; +} + +mod output; +mod reconstruction; + +#[inline(always)] +fn device_load(pointer: *const T, index: u32) -> T { + // SAFETY: Host plans validate all buffer ranges and device descriptor offsets. + unsafe { *pointer.add(index as usize) } +} + +#[inline(always)] +fn store(pointer: *mut T, index: u32, value: T) { + // SAFETY: Host plans validate destination extents; each work item owns its writes. + unsafe { + *pointer.add(index as usize) = value; + } +} + +#[inline(always)] +fn dequantize(value: i32, quantizer: u32) -> Result { + checked_i64_to_i32(value as i64 * quantizer as i64) +} + +fn fail(status: *mut u32, code: u32) { + // SAFETY: Every submission retains one aligned status word until completion. + let atomic = unsafe { DeviceAtomicU32::from_ptr(status) }; + let _ = atomic.compare_exchange(0, code, AtomicOrdering::Relaxed, AtomicOrdering::Relaxed); +} + +fn block_failed(status: *mut u32) -> bool { + static mut STATUS: SharedArray = SharedArray::UNINIT; + let tid = thread::threadIdx_x() + + thread::blockDim_x() + * (thread::threadIdx_y() + thread::blockDim_y() * thread::threadIdx_z()); + // SAFETY: Only thread zero writes; every thread reaches the barrier before reading. + unsafe { + if tid == 0 { + STATUS[0] = DeviceAtomicU32::from_ptr(status).load(AtomicOrdering::Relaxed); + } + thread::sync_threads(); + STATUS[0] != 0 + } +} + +fn output_store( + samples: *const i32, + planes: *const JxrSamplePlaneAbi, + surfaces: *const JxrSurfacePlaneAbi, + output: *mut u8, + status: *mut u32, + params: JxrOutputAbi, + output_base: u32, + kind: u32, +) { + let x = thread::blockIdx_x() * thread::blockDim_x() + thread::threadIdx_x(); + let y = thread::blockIdx_y() * thread::blockDim_y() + thread::threadIdx_y(); + let surface = device_load(surfaces, params.output_plane); + let failed = block_failed(status); + let width = if kind == 0 { + (surface.width + 7) / 8 + } else { + surface.width + }; + if x >= width || y >= surface.height || failed { + return; + } + // SAFETY: Output allocation and the per-image base are validated before launch. + let output = unsafe { output.add(output_base as usize) }; + if output::write(samples, planes, surface, output, params, x, y, kind).is_err() { + fail(status, 16); + } +} + +#[cuda_module] +mod kernels { + use super::*; + + #[kernel] + pub unsafe fn jxr_dequantize_first_transform( + packed: *const i32, + macroblocks: *const JxrMacroblockAbi, + low_plane: *mut i32, + status: *mut u32, + plane: JxrPlaneAbi, + ) { + let gid = thread::blockIdx_x() * thread::blockDim_x() + thread::threadIdx_x(); + let failed = block_failed(status); + if gid >= plane.macroblock_count || failed { + return; + } + if reconstruction::first_transform(packed, macroblocks, low_plane, plane, gid).is_err() { + fail(status, 1); + } + } + + #[kernel] + pub unsafe fn jxr_highpass_second_transform( + packed: *const i32, + macroblocks: *const JxrMacroblockAbi, + low_plane: *const i32, + samples: *mut i32, + status: *mut u32, + plane: JxrPlaneAbi, + ) { + static mut HIGH: SharedArray = SharedArray::UNINIT; + let group = thread::blockIdx_x(); + let tid = thread::threadIdx_x(); + let threads = thread::blockDim_x(); + let failed = block_failed(status); + if group >= plane.macroblock_count || failed { + return; + } + let metadata = device_load(macroblocks, plane.macroblock_offset + group); + let count = plane.block_columns * plane.block_rows; + // SAFETY: Each lane initializes distinct entries, all lanes synchronize before prediction. + let high = unsafe { HIGH.as_mut_ptr() }; + let mut index = tid; + while index < count * 16 { + let value = if metadata.bands < 2 || index & 15 == 0 { + 0 + } else { + device_load(packed, metadata.coefficient_offset + index) + }; + store(high, index, value); + index += threads; + } + thread::sync_threads(); + if tid == 0 && reconstruction::predict_high(high, metadata, plane).is_err() { + fail(status, 3); + } + thread::sync_threads(); + if block_failed(status) { + return; + } + let mut block = tid; + while block < count { + if reconstruction::second_transform(high, low_plane, samples, metadata, plane, block) + .is_err() + { + fail(status, 4); + return; + } + block += threads; + } + } + + #[kernel] + pub unsafe fn jxr_first_overlap( + samples: *mut i32, + work: *const JxrOverlapWorkAbi, + status: *mut u32, + count: u32, + ) { + let gid = thread::blockIdx_x() * thread::blockDim_x() + thread::threadIdx_x(); + let failed = block_failed(status); + if gid < count + && !failed + && reconstruction::overlap(samples, device_load(work, gid)).is_err() + { + fail(status, 2); + } + } + + #[kernel] + pub unsafe fn jxr_second_overlap( + samples: *mut i32, + work: *const JxrOverlapWorkAbi, + status: *mut u32, + count: u32, + ) { + let gid = thread::blockIdx_x() * thread::blockDim_x() + thread::threadIdx_x(); + let failed = block_failed(status); + if gid < count + && !failed + && reconstruction::overlap(samples, device_load(work, gid)).is_err() + { + fail(status, 2); + } + } + + #[kernel] + pub unsafe fn jxr_output_bits( + samples: *const i32, + planes: *const JxrSamplePlaneAbi, + surfaces: *const JxrSurfacePlaneAbi, + output: *mut u8, + status: *mut u32, + params: JxrOutputAbi, + base: u32, + ) { + output_store(samples, planes, surfaces, output, status, params, base, 0); + } + #[kernel] + pub unsafe fn jxr_output_u8( + samples: *const i32, + planes: *const JxrSamplePlaneAbi, + surfaces: *const JxrSurfacePlaneAbi, + output: *mut u8, + status: *mut u32, + params: JxrOutputAbi, + base: u32, + ) { + output_store(samples, planes, surfaces, output, status, params, base, 1); + } + #[kernel] + pub unsafe fn jxr_output_u16( + samples: *const i32, + planes: *const JxrSamplePlaneAbi, + surfaces: *const JxrSurfacePlaneAbi, + output: *mut u8, + status: *mut u32, + params: JxrOutputAbi, + base: u32, + ) { + output_store(samples, planes, surfaces, output, status, params, base, 2); + } + #[kernel] + pub unsafe fn jxr_output_i16( + samples: *const i32, + planes: *const JxrSamplePlaneAbi, + surfaces: *const JxrSurfacePlaneAbi, + output: *mut u8, + status: *mut u32, + params: JxrOutputAbi, + base: u32, + ) { + output_store(samples, planes, surfaces, output, status, params, base, 3); + } + #[kernel] + pub unsafe fn jxr_output_i32( + samples: *const i32, + planes: *const JxrSamplePlaneAbi, + surfaces: *const JxrSurfacePlaneAbi, + output: *mut u8, + status: *mut u32, + params: JxrOutputAbi, + base: u32, + ) { + output_store(samples, planes, surfaces, output, status, params, base, 4); + } + #[kernel] + pub unsafe fn jxr_output_f16( + samples: *const i32, + planes: *const JxrSamplePlaneAbi, + surfaces: *const JxrSurfacePlaneAbi, + output: *mut u8, + status: *mut u32, + params: JxrOutputAbi, + base: u32, + ) { + output_store(samples, planes, surfaces, output, status, params, base, 5); + } + #[kernel] + pub unsafe fn jxr_output_f32( + samples: *const i32, + planes: *const JxrSamplePlaneAbi, + surfaces: *const JxrSurfacePlaneAbi, + output: *mut u8, + status: *mut u32, + params: JxrOutputAbi, + base: u32, + ) { + output_store(samples, planes, surfaces, output, status, params, base, 6); + } + #[kernel] + pub unsafe fn jxr_output_packed16( + samples: *const i32, + planes: *const JxrSamplePlaneAbi, + surfaces: *const JxrSurfacePlaneAbi, + output: *mut u8, + status: *mut u32, + params: JxrOutputAbi, + base: u32, + ) { + output_store(samples, planes, surfaces, output, status, params, base, 7); + } + #[kernel] + pub unsafe fn jxr_output_packed32( + samples: *const i32, + planes: *const JxrSamplePlaneAbi, + surfaces: *const JxrSurfacePlaneAbi, + output: *mut u8, + status: *mut u32, + params: JxrOutputAbi, + base: u32, + ) { + output_store(samples, planes, surfaces, output, status, params, base, 8); + } +} + +fn main() {} diff --git a/crates/jxr-cuda/src/kernels/output.cu b/crates/jxr-cuda/src/kernels/output.cu deleted file mode 100644 index a93a146..0000000 --- a/crates/jxr-cuda/src/kernels/output.cu +++ /dev/null @@ -1,605 +0,0 @@ -// SPDX-License-Identifier: MIT OR Apache-2.0 - -struct JxrInt4 { int x; int y; int z; int w; }; - -__device__ __forceinline__ int jxr_read_plane( - const int *samples, JxrSamplePlaneAbi plane, uint x, uint y) { - return samples[plane.sample_offset + (y - plane.origin_y) * plane.width + x - plane.origin_x]; -} - -__device__ __forceinline__ JxrInt4 jxr_centering(uint code) { - switch (code) { - case 0u: return {4, 4, 0, 8}; - case 1u: return {5, 3, 1, 7}; - case 2u: return {6, 2, 2, 6}; - case 3u: return {7, 1, 3, 5}; - default: return {8, 0, 4, 4}; - } -} - -__device__ __forceinline__ bool jxr_weighted( - int first, int first_weight, int second, int second_weight, - int &result, uint *status) { - jxr_long value = (jxr_long)first * first_weight - + (jxr_long)second * second_weight + 4ll; - return jxr_narrow(value >> 3, result, status, 16u); -} - -__device__ __forceinline__ bool jxr_upsample_pair( - int previous, int current, int next, uint centering, int pair[2], uint *status) { - JxrInt4 weights = jxr_centering(centering); - return jxr_weighted(previous, weights.z, current, weights.w, pair[0], status) - && jxr_weighted(current, weights.x, next, weights.y, pair[1], status); -} - -__device__ __forceinline__ int jxr_clamped_plane( - const int *samples, JxrSamplePlaneAbi plane, int x, int y) { - int local_x = jxr_clamp_i32(x - (int)plane.origin_x, 0, (int)plane.width - 1); - int local_y = jxr_clamp_i32(y - (int)plane.origin_y, 0, (int)plane.height - 1); - return samples[plane.sample_offset + (uint)local_y * plane.width + (uint)local_x]; -} - -__device__ __forceinline__ bool jxr_chroma_sample( - const int *samples, JxrSamplePlaneAbi plane, uint full_x, uint full_y, - const JxrOutputAbi ¶ms, int &result, uint *status) { - if (params.chroma_sampling == 3u) { - result = jxr_read_plane(samples, plane, full_x, full_y); - return true; - } - int chroma_x = (int)(full_x >> 1u); - if (params.chroma_sampling == 2u) { - int pair[2]; - if (!jxr_upsample_pair( - jxr_clamped_plane(samples, plane, chroma_x - 1, (int)full_y), - jxr_clamped_plane(samples, plane, chroma_x, (int)full_y), - jxr_clamped_plane(samples, plane, chroma_x + 1, (int)full_y), - params.chroma_centering_x, pair, status)) { - return false; - } - result = pair[full_x & 1u]; - return true; - } - int chroma_y = (int)(full_y >> 1u); - int vertical[3]; - for (int column = -1; column <= 1; ++column) { - int pair[2]; - if (!jxr_upsample_pair( - jxr_clamped_plane(samples, plane, chroma_x + column, chroma_y - 1), - jxr_clamped_plane(samples, plane, chroma_x + column, chroma_y), - jxr_clamped_plane(samples, plane, chroma_x + column, chroma_y + 1), - params.chroma_centering_y, pair, status)) { - return false; - } - vertical[column + 1] = pair[full_y & 1u]; - } - int horizontal[2]; - if (!jxr_upsample_pair( - vertical[0], vertical[1], vertical[2], params.chroma_centering_x, - horizontal, status)) { - return false; - } - result = horizontal[full_x & 1u]; - return true; -} - -__device__ __forceinline__ bool jxr_primary_values( - const int *samples, const JxrSamplePlaneAbi *planes, uint x, uint y, - const JxrOutputAbi ¶ms, int values[4], uint *status) { - for (uint i = 0; i < 4; ++i) values[i] = 0; - values[0] = jxr_read_plane(samples, planes[0], x, y); - if (params.component_count == 1u) return true; - if (params.internal_color == 1u) { - if (!jxr_chroma_sample(samples, planes[1], x, y, params, values[1], status) - || !jxr_chroma_sample(samples, planes[2], x, y, params, values[2], status)) { - return false; - } - } else { - for (uint i = 1; i < jxr_min_u32(params.component_count, 4u); ++i) { - values[i] = jxr_read_plane(samples, planes[i], x, y); - } - } - return true; -} - -__device__ __forceinline__ bool jxr_converted( - const int *samples, const JxrSamplePlaneAbi *planes, uint x, uint y, - const JxrOutputAbi ¶ms, int values[4], uint *status) { - int input[4]; - if (!jxr_primary_values(samples, planes, x, y, params, input, status)) return false; - for (uint i = 0; i < 4; ++i) values[i] = input[i]; - if (params.internal_color == 0u && params.output_color == 2u) { - values[1] = input[0]; - values[2] = input[0]; - } else if (params.internal_color == 1u - && (params.output_color == 2u || params.output_color == 6u)) { - int temporary; - int green; - int red; - if (!jxr_sub(0, input[1], temporary, status, 16u) - || !jxr_sub(input[0], temporary >> 1, green, status, 16u) - || !jxr_add(temporary, green, red, status, 16u) - || !jxr_sub(red, (input[2] >> 1) + (input[2] & 1), red, status, 16u) - || !jxr_add(input[2], red, values[2], status, 16u)) { - return false; - } - values[0] = red; - values[1] = green; - if (params.bit_depth >= 8u && params.bit_depth <= 10u - && params.red_blue_not_swapped == 0u) { - int swap = values[0]; - values[0] = values[2]; - values[2] = swap; - } - } else if (params.internal_color == 5u && params.output_color == 3u) { - int black; - int magenta; - int cyan; - if (!jxr_add(input[3], input[0] >> 1, black, status, 16u) - || !jxr_sub(black, input[0], magenta, status, 16u) - || !jxr_sub(magenta, input[1] >> 1, magenta, status, 16u) - || !jxr_add(input[1], magenta, cyan, status, 16u) - || !jxr_add(cyan, input[2] >> 1, cyan, status, 16u) - || !jxr_sub(cyan, input[2], values[2], status, 16u)) { - return false; - } - values[0] = cyan; - values[1] = magenta; - values[3] = black; - } else if (params.internal_color == 5u && params.output_color == 4u) { - values[0] = input[1]; - values[1] = input[2]; - values[2] = input[3]; - values[3] = input[0]; - } - return true; -} - -__device__ __forceinline__ int jxr_base_bias(uint depth, uint shift_bits) { - switch (depth) { - case 1u: return 128; - case 2u: - case 9u: return 512; - case 3u: return 32768 >> shift_bits; - case 8u: return 16; - case 10u: return 32; - default: return 0; - } -} - -__device__ __forceinline__ bool jxr_scale( - int sample, uint component, bool alpha, const JxrOutputAbi ¶ms, - int &result, uint *status) { - uint depth = params.bit_depth; - uint scaled = alpha ? params.alpha_scaled : params.scaled; - uint shift_bits = alpha ? params.alpha_shift_bits : params.shift_bits; - int bias = jxr_base_bias(depth, shift_bits); - if (!alpha && params.output_color == 3u) { - bias = component < 3u ? bias >> 1 : -(bias >> 1); - } - if (!alpha && params.output_color == 6u) bias = 0; - uint scale = scaled * 3u; - int shifted_bias; - if (!jxr_narrow((jxr_long)bias * (1ll << scale), shifted_bias, status, 16u)) return false; - int biased; - if (!jxr_add(sample, shifted_bias, biased, status, 16u)) return false; - int rounding = scaled == 0u ? 0 : ((depth == 0u || depth == 3u) ? 4 : 3); - int rounded; - if (!jxr_add(biased, rounding, rounded, status, 16u)) return false; - uint downshift = scale + ((depth == 10u && component != 1u) ? 1u : 0u); - int downscaled = rounded >> downshift; - if (depth == 3u || depth == 4u || depth == 6u) { - return jxr_narrow( - (jxr_long)downscaled * (1ll << shift_bits), result, status, 16u); - } - result = downscaled; - return true; -} - -__device__ __forceinline__ ulong jxr_abs_long(int value) { - return value < 0 ? (ulong)(-(jxr_long)value) : (ulong)value; -} - -__device__ __forceinline__ uint jxr_f16_bits( - int sample, bool alpha, const JxrOutputAbi ¶ms, uint *status) { - int scaled; - if (!jxr_scale(sample, 0u, alpha, params, scaled, status)) return 0u; - uint sign = scaled < 0 ? 0x8000u : 0u; - uint magnitude = (uint)jxr_min_u32((uint)jxr_abs_long(scaled), 32767u); - return sign | magnitude; -} - -__device__ __forceinline__ uint jxr_f32_bits( - int sample, bool alpha, const JxrOutputAbi ¶ms, uint *status) { - int scaled; - if (!jxr_scale(sample, 0u, alpha, params, scaled, status)) return 0u; - uint length = alpha ? params.alpha_mantissa_length : params.mantissa_length; - int exponent_bias = (int)(alpha ? params.alpha_exponent_bias_bits : params.exponent_bias_bits); - uint sign = scaled < 0 ? 0x80000000u : 0u; - ulong magnitude = jxr_abs_long(scaled); - ulong implicit = 1ull << length; - jxr_long exponent = (jxr_long)(magnitude >> length); - ulong mantissa = (magnitude & (implicit - 1ull)) | implicit; - if (exponent == 0ll) { - mantissa ^= implicit; - exponent = 1ll; - } - exponent = exponent - (jxr_long)exponent_bias + 127ll; - while (mantissa < implicit && exponent > 1ll && mantissa > 0ull) { - --exponent; - mantissa <<= 1u; - } - if (mantissa < implicit) { - exponent = 0ll; - } else { - mantissa ^= implicit; - } - mantissa <<= 23u - length; - if (exponent < 0ll || exponent > 255ll || mantissa > 0x7fffffull) { - jxr_fail(status, 16u); - return 0u; - } - return sign | ((uint)exponent << 23u) | (uint)mantissa; -} - -__device__ __forceinline__ uint jxr_unsigned_premultiply( - uint value, uint alpha, uint maximum) { - return (uint)(((ulong)jxr_min_u32(value, maximum) * (ulong)jxr_min_u32(alpha, maximum) - + (ulong)(maximum / 2u)) / (ulong)maximum); -} - -__device__ __forceinline__ int jxr_signed_premultiply( - int value, int alpha, int maximum) { - jxr_long magnitude = (jxr_long)jxr_abs_long(value); - jxr_long result = (magnitude * (jxr_long)jxr_clamp_i32(alpha, 0, maximum) - + (jxr_long)(maximum / 2)) / (jxr_long)maximum; - return (int)(value < 0 ? -result : result); -} - -__device__ __forceinline__ bool jxr_component( - const int *samples, const JxrSamplePlaneAbi *planes, uint channel, uint x, uint y, - const JxrOutputAbi ¶ms, int &sample, bool &alpha, uint *status) { - uint primary_count = params.alpha_plane == JXR_UINT_MAX - ? params.channels : params.channels - 1u; - alpha = params.alpha_plane != JXR_UINT_MAX && channel == primary_count; - if (alpha) { - sample = jxr_read_plane(samples, planes[params.alpha_plane], x, y); - return true; - } - if (params.output_color == 7u) { - sample = jxr_read_plane(samples, planes[channel], x, y); - return true; - } - if ((params.channel_layout == 12u || params.channel_layout == 13u) && channel == 3u) { - sample = 0; - alpha = false; - return true; - } - int values[4]; - if (!jxr_converted(samples, planes, x, y, params, values, status)) return false; - uint source_channel = channel; - if ((params.channel_layout == 6u || params.channel_layout == 7u - || params.channel_layout == 13u) && channel < 3u) { - source_channel = 2u - channel; - } - sample = values[source_channel]; - return true; -} - -__device__ __forceinline__ bool jxr_formatted_integer( - const int *samples, const JxrSamplePlaneAbi *planes, uint channel, uint x, uint y, - const JxrOutputAbi ¶ms, int &value, uint *status) { - if ((params.channel_layout == 12u || params.channel_layout == 13u) && channel == 3u) { - value = 0; - return true; - } - int sample; - bool alpha; - if (!jxr_component(samples, planes, channel, x, y, params, sample, alpha, status) - || !jxr_scale(sample, channel, alpha, params, value, status)) { - return false; - } - if (params.premultiply_alpha != 0u && !alpha && params.alpha_plane != JXR_UINT_MAX) { - int alpha_sample = jxr_read_plane(samples, planes[params.alpha_plane], x, y); - int alpha_value; - if (!jxr_scale(alpha_sample, 0u, true, params, alpha_value, status)) return false; - if (params.bit_depth == 1u) { - value = (int)jxr_unsigned_premultiply( - (uint)jxr_clamp_i32(value, 0, 255), - (uint)jxr_clamp_i32(alpha_value, 0, 255), 255u); - } else if (params.bit_depth == 2u || params.bit_depth == 3u) { - value = (int)jxr_unsigned_premultiply( - (uint)jxr_clamp_i32(value, 0, 65535), - (uint)jxr_clamp_i32(alpha_value, 0, 65535), 65535u); - } else if (params.bit_depth == 4u) { - value = jxr_signed_premultiply(value, alpha_value, 32767); - } else if (params.bit_depth == 6u) { - value = jxr_signed_premultiply(value, alpha_value, JXR_INT_MAX); - } - } - return true; -} - -__device__ __forceinline__ uint jxr_output_index( - JxrSurfacePlaneAbi surface, uint x, uint y, uint channel, uint bytes) { - return surface.byte_offset + y * surface.row_stride_bytes - + (x * surface.channels + channel) * bytes; -} - -__device__ __forceinline__ void jxr_store_u16(uchar *output, uint index, uint value) { - output[index] = (uchar)value; - output[index + 1u] = (uchar)(value >> 8u); -} - -__device__ __forceinline__ void jxr_store_u32(uchar *output, uint index, uint value) { - output[index] = (uchar)value; - output[index + 1u] = (uchar)(value >> 8u); - output[index + 2u] = (uchar)(value >> 16u); - output[index + 3u] = (uchar)(value >> 24u); -} - -extern "C" __global__ void jxr_output_u8( - const int *samples, const JxrSamplePlaneAbi *planes, - const JxrSurfacePlaneAbi *surfaces, uchar *output, uint *status, - JxrOutputAbi params, uint output_base) { - output += output_base; - uint x_out = blockIdx.x * blockDim.x + threadIdx.x; - uint y_out = blockIdx.y * blockDim.y + threadIdx.y; - JxrSurfacePlaneAbi surface = surfaces[params.output_plane]; - const bool failed = jxr_block_failed(status); - if (x_out >= surface.width || y_out >= surface.height || failed) return; - if (params.output_plane_count > 1u) { - uint source_plane = params.output_plane < 3u ? params.output_plane : params.alpha_plane; - uint x = params.crop_x - / (params.output_plane == 1u || params.output_plane == 2u ? 2u : 1u) + x_out; - uint y_divisor = params.chroma_sampling == 1u - && (params.output_plane == 1u || params.output_plane == 2u) ? 2u : 1u; - uint y = params.crop_y / y_divisor + y_out; - int scaled; - bool alpha = params.output_plane >= 3u; - if (!jxr_scale( - jxr_read_plane(samples, planes[source_plane], x, y), - params.output_plane, alpha, params, scaled, status)) return; - output[jxr_output_index(surface, x_out, y_out, 0u, 1u)] = - (uchar)jxr_clamp_i32(scaled, 0, 255); - return; - } - uint x = params.crop_x + x_out; - uint y = params.crop_y + y_out; - for (uint channel = 0; channel < params.channels; ++channel) { - int value; - if (!jxr_formatted_integer( - samples, planes, channel, x, y, params, value, status)) return; - output[jxr_output_index(surface, x_out, y_out, channel, 1u)] = - (uchar)jxr_clamp_i32(value, 0, 255); - } -} - -__device__ __forceinline__ bool jxr_integer_store_value( - const int *samples, const JxrSamplePlaneAbi *planes, - JxrSurfacePlaneAbi surface, uint x_out, uint y_out, uint channel, - const JxrOutputAbi ¶ms, int &value, uint *status) { - uint x = params.crop_x + x_out; - uint y = params.crop_y + y_out; - if (params.output_plane_count > 1u) { - uint source_plane = params.output_plane < 3u ? params.output_plane : params.alpha_plane; - x = params.crop_x - / (params.output_plane == 1u || params.output_plane == 2u ? 2u : 1u) + x_out; - uint divisor = params.chroma_sampling == 1u - && (params.output_plane == 1u || params.output_plane == 2u) ? 2u : 1u; - y = params.crop_y / divisor + y_out; - bool alpha = params.output_plane >= 3u; - return jxr_scale( - jxr_read_plane(samples, planes[source_plane], x, y), - params.output_plane, alpha, params, value, status); - } - return jxr_formatted_integer(samples, planes, channel, x, y, params, value, status); -} - -#define JXR_INTEGER_KERNEL(NAME, BYTES, MINIMUM, MAXIMUM, STORE) \ -extern "C" __global__ void NAME( \ - const int *samples, const JxrSamplePlaneAbi *planes, \ - const JxrSurfacePlaneAbi *surfaces, uchar *output, uint *status, \ - JxrOutputAbi params, uint output_base) { \ - output += output_base; \ - uint x_out = blockIdx.x * blockDim.x + threadIdx.x; \ - uint y_out = blockIdx.y * blockDim.y + threadIdx.y; \ - JxrSurfacePlaneAbi surface = surfaces[params.output_plane]; \ - const bool failed = jxr_block_failed(status); \ - if (x_out >= surface.width || y_out >= surface.height || failed) return; \ - uint channels = params.output_plane_count > 1u ? 1u : surface.channels; \ - for (uint channel = 0; channel < channels; ++channel) { \ - int value; \ - if (!jxr_integer_store_value(samples, planes, surface, x_out, y_out, channel, params, value, status)) return; \ - uint index = jxr_output_index(surface, x_out, y_out, channel, BYTES); \ - STORE(output, index, (uint)jxr_clamp_i32(value, MINIMUM, MAXIMUM)); \ - } \ -} - -#define JXR_STORE16(output, index, value) jxr_store_u16(output, index, value) -#define JXR_STORE32(output, index, value) jxr_store_u32(output, index, value) -JXR_INTEGER_KERNEL(jxr_output_u16, 2u, 0, - (params.bit_depth == 2u ? 1023 : 65535), JXR_STORE16) -JXR_INTEGER_KERNEL(jxr_output_i16, 2u, -32768, 32767, JXR_STORE16) -JXR_INTEGER_KERNEL(jxr_output_i32, 4u, JXR_INT_MIN, JXR_INT_MAX, JXR_STORE32) - -extern "C" __global__ void jxr_output_f16( - const int *samples, const JxrSamplePlaneAbi *planes, - const JxrSurfacePlaneAbi *surfaces, uchar *output, uint *status, - JxrOutputAbi params, uint output_base) { - output += output_base; - uint x_out = blockIdx.x * blockDim.x + threadIdx.x; - uint y_out = blockIdx.y * blockDim.y + threadIdx.y; - JxrSurfacePlaneAbi surface = surfaces[0]; - const bool failed = jxr_block_failed(status); - if (x_out >= surface.width || y_out >= surface.height || failed) return; - uint x = params.crop_x + x_out; - uint y = params.crop_y + y_out; - uint alpha_bits = params.alpha_plane == JXR_UINT_MAX ? 0u - : jxr_f16_bits( - jxr_read_plane(samples, planes[params.alpha_plane], x, y), true, params, status); - for (uint channel = 0; channel < params.channels; ++channel) { - uint index = jxr_output_index(surface, x_out, y_out, channel, 2u); - if ((params.channel_layout == 12u || params.channel_layout == 13u) && channel == 3u) { - jxr_store_u16(output, index, 0u); - continue; - } - int sample; - bool alpha; - if (!jxr_component( - samples, planes, channel, x, y, params, sample, alpha, status)) return; - uint bits = jxr_f16_bits(sample, alpha, params, status); - if (params.premultiply_alpha != 0u && !alpha) { - uint sign = bits & 0x8000u; - bits = sign | jxr_unsigned_premultiply( - bits & 0x7fffu, - (alpha_bits & 0x8000u) == 0u ? alpha_bits & 0x7fffu : 0u, - 0x7fffu); - } - jxr_store_u16(output, index, bits); - } -} - -extern "C" __global__ void jxr_output_f32( - const int *samples, const JxrSamplePlaneAbi *planes, - const JxrSurfacePlaneAbi *surfaces, uchar *output, uint *status, - JxrOutputAbi params, uint output_base) { - output += output_base; - uint x_out = blockIdx.x * blockDim.x + threadIdx.x; - uint y_out = blockIdx.y * blockDim.y + threadIdx.y; - JxrSurfacePlaneAbi surface = surfaces[0]; - const bool failed = jxr_block_failed(status); - if (x_out >= surface.width || y_out >= surface.height || failed) return; - uint x = params.crop_x + x_out; - uint y = params.crop_y + y_out; - float alpha_value = 1.0f; - if (params.alpha_plane != JXR_UINT_MAX) { - alpha_value = jxr_clamp_f32(__uint_as_float(jxr_f32_bits( - jxr_read_plane(samples, planes[params.alpha_plane], x, y), true, params, status)), - 0.0f, 1.0f); - } - for (uint channel = 0; channel < params.channels; ++channel) { - uint index = jxr_output_index(surface, x_out, y_out, channel, 4u); - if ((params.channel_layout == 12u || params.channel_layout == 13u) && channel == 3u) { - jxr_store_u32(output, index, 0u); - continue; - } - int sample; - bool alpha; - if (!jxr_component( - samples, planes, channel, x, y, params, sample, alpha, status)) return; - float value = __uint_as_float(jxr_f32_bits(sample, alpha, params, status)); - if (params.premultiply_alpha != 0u && !alpha) value *= alpha_value; - jxr_store_u32(output, index, __float_as_uint(value)); - } -} - -extern "C" __global__ void jxr_output_bits( - const int *samples, const JxrSamplePlaneAbi *planes, - const JxrSurfacePlaneAbi *surfaces, uchar *output, uint *status, - JxrOutputAbi params, uint output_base) { - output += output_base; - uint byte_x = blockIdx.x * blockDim.x + threadIdx.x; - uint y_out = blockIdx.y * blockDim.y + threadIdx.y; - JxrSurfacePlaneAbi surface = surfaces[0]; - uint row_bytes = (surface.width + 7u) / 8u; - const bool failed = jxr_block_failed(status); - if (byte_x >= row_bytes || y_out >= surface.height || failed) return; - uchar byte = 0; - for (uint bit = 0; bit < 8u; ++bit) { - uint pixel_x = byte_x * 8u + bit; - if (pixel_x >= surface.width) break; - int value; - if (!jxr_scale(jxr_read_plane( - samples, planes[0], params.crop_x + pixel_x, params.crop_y + y_out), - 0u, false, params, value, status)) return; - uint packed = (uint)jxr_clamp_i32(value, 0, 1); - if (params.bit_black != 0u) packed = 1u - packed; - byte |= (uchar)(packed << (7u - bit)); - } - output[surface.byte_offset + y_out * surface.row_stride_bytes + byte_x] = byte; -} - -extern "C" __global__ void jxr_output_packed16( - const int *samples, const JxrSamplePlaneAbi *planes, - const JxrSurfacePlaneAbi *surfaces, uchar *output, uint *status, - JxrOutputAbi params, uint output_base) { - output += output_base; - uint x_out = blockIdx.x * blockDim.x + threadIdx.x; - uint y_out = blockIdx.y * blockDim.y + threadIdx.y; - JxrSurfacePlaneAbi surface = surfaces[0]; - const bool failed = jxr_block_failed(status); - if (x_out >= surface.width || y_out >= surface.height || failed) return; - int values[4]; - if (!jxr_converted( - samples, planes, params.crop_x + x_out, params.crop_y + y_out, - params, values, status)) return; - uint packed = 0u; - for (uint channel = 0; channel < 3u; ++channel) { - int value; - if (!jxr_scale(values[channel], channel, false, params, value, status)) return; - uint maximum = params.bit_depth == 10u && channel == 1u ? 63u : 31u; - uint shift = params.bit_depth == 10u - ? (channel == 0u ? 11u : (channel == 1u ? 5u : 0u)) - : (2u - channel) * 5u; - packed |= (uint)jxr_clamp_i32(value, 0, (int)maximum) << shift; - } - jxr_store_u16( - output, surface.byte_offset + y_out * surface.row_stride_bytes + x_out * 2u, packed); -} - -extern "C" __global__ void jxr_output_packed32( - const int *samples, const JxrSamplePlaneAbi *planes, - const JxrSurfacePlaneAbi *surfaces, uchar *output, uint *status, - JxrOutputAbi params, uint output_base) { - output += output_base; - uint x_out = blockIdx.x * blockDim.x + threadIdx.x; - uint y_out = blockIdx.y * blockDim.y + threadIdx.y; - JxrSurfacePlaneAbi surface = surfaces[0]; - const bool failed = jxr_block_failed(status); - if (x_out >= surface.width || y_out >= surface.height || failed) return; - int values[4]; - if (!jxr_converted( - samples, planes, params.crop_x + x_out, params.crop_y + y_out, - params, values, status)) return; - uint packed = 0u; - if (params.output_color == 6u) { - int scaled[3]; - uint exponent = 0u; - uint mantissa[3]; - uint local_exponent[3]; - for (uint channel = 0; channel < 3u; ++channel) { - if (!jxr_scale( - values[channel], channel, false, params, scaled[channel], status)) return; - if (scaled[channel] <= 0) { - mantissa[channel] = 0; - local_exponent[channel] = 0; - } else if ((scaled[channel] >> 7) > 1) { - mantissa[channel] = (uint)((scaled[channel] & 127) + 128); - local_exponent[channel] = (uint)(scaled[channel] >> 7); - } else { - mantissa[channel] = (uint)scaled[channel]; - local_exponent[channel] = 1u; - } - exponent = jxr_max_u32(exponent, local_exponent[channel]); - } - for (uint channel = 0; channel < 3u; ++channel) { - if (exponent > local_exponent[channel]) { - uint difference = exponent - local_exponent[channel]; - mantissa[channel] = difference >= 31u ? 0u - : (uint)((2u * mantissa[channel] + 1u) >> (difference + 1u)); - } - } - packed = jxr_min_u32(mantissa[0], 255u) - | (jxr_min_u32(mantissa[1], 255u) << 8u) - | (jxr_min_u32(mantissa[2], 255u) << 16u) - | (jxr_min_u32(exponent, 255u) << 24u); - } else { - for (uint channel = 0; channel < 3u; ++channel) { - int value; - if (!jxr_scale(values[channel], channel, false, params, value, status)) return; - packed |= (uint)jxr_clamp_i32(value, 0, 1023) << ((2u - channel) * 10u); - } - } - jxr_store_u32( - output, surface.byte_offset + y_out * surface.row_stride_bytes + x_out * 4u, packed); -} diff --git a/crates/jxr-cuda/src/kernels/output.rs b/crates/jxr-cuda/src/kernels/output.rs new file mode 100644 index 0000000..5adb778 --- /dev/null +++ b/crates/jxr-cuda/src/kernels/output.rs @@ -0,0 +1,562 @@ +// SPDX-License-Identifier: MIT OR Apache-2.0 + +use super::*; + +fn read_plane(samples: *const i32, plane: JxrSamplePlaneAbi, x: u32, y: u32) -> i32 { + device_load( + samples, + plane.sample_offset + (y - plane.origin_y) * plane.width + x - plane.origin_x, + ) +} + +fn clamped_plane(samples: *const i32, plane: JxrSamplePlaneAbi, x: i32, y: i32) -> i32 { + let local_x = (x - plane.origin_x as i32).clamp(0, plane.width as i32 - 1); + let local_y = (y - plane.origin_y as i32).clamp(0, plane.height as i32 - 1); + device_load( + samples, + plane.sample_offset + local_y as u32 * plane.width + local_x as u32, + ) +} + +fn upsample(previous: i32, current: i32, next: i32, centering: u32) -> Result<[i32; 2], MathError> { + let weights = match centering { + 0 => [4, 4, 0, 8], + 1 => [5, 3, 1, 7], + 2 => [6, 2, 2, 6], + 3 => [7, 1, 3, 5], + _ => [8, 0, 4, 4], + }; + Ok([ + math_try!(checked_i64_to_i32( + (previous as i64 * weights[2] + current as i64 * weights[3] + 4) >> 3 + )), + math_try!(checked_i64_to_i32( + (current as i64 * weights[0] + next as i64 * weights[1] + 4) >> 3 + )), + ]) +} + +fn chroma( + samples: *const i32, + plane: JxrSamplePlaneAbi, + x: u32, + y: u32, + p: JxrOutputAbi, +) -> Result { + if p.chroma_sampling == 3 { + return Ok(read_plane(samples, plane, x, y)); + } + let cx = (x >> 1) as i32; + if p.chroma_sampling == 2 { + let pair = math_try!(upsample( + clamped_plane(samples, plane, cx - 1, y as i32), + clamped_plane(samples, plane, cx, y as i32), + clamped_plane(samples, plane, cx + 1, y as i32), + p.chroma_centering_x, + )); + return Ok(pair[(x & 1) as usize]); + } + let cy = (y >> 1) as i32; + let mut vertical = [0; 3]; + for column in -1..=1 { + let pair = math_try!(upsample( + clamped_plane(samples, plane, cx + column, cy - 1), + clamped_plane(samples, plane, cx + column, cy), + clamped_plane(samples, plane, cx + column, cy + 1), + p.chroma_centering_y, + )); + vertical[(column + 1) as usize] = pair[(y & 1) as usize]; + } + let horizontal = math_try!(upsample( + vertical[0], + vertical[1], + vertical[2], + p.chroma_centering_x + )); + Ok(horizontal[(x & 1) as usize]) +} + +fn converted( + samples: *const i32, + planes: *const JxrSamplePlaneAbi, + x: u32, + y: u32, + p: JxrOutputAbi, +) -> Result<[i32; 4], MathError> { + let mut input = [0; 4]; + input[0] = read_plane(samples, device_load(planes, 0), x, y); + if p.component_count > 1 { + if p.internal_color == 1 { + input[1] = math_try!(chroma(samples, device_load(planes, 1), x, y, p)); + input[2] = math_try!(chroma(samples, device_load(planes, 2), x, y, p)); + } else { + for i in 1..p.component_count.min(4) { + input[i as usize] = read_plane(samples, device_load(planes, i), x, y); + } + } + } + let mut values = input; + if p.internal_color == 0 && p.output_color == 2 { + values[1] = input[0]; + values[2] = input[0]; + } else if p.internal_color == 1 && (p.output_color == 2 || p.output_color == 6) { + let temporary = math_try!(checked_sub(0, input[1])); + let green = math_try!(checked_sub(input[0], temporary >> 1)); + let red = math_try!(checked_sub( + math_try!(checked_add(temporary, green)), + (input[2] >> 1) + (input[2] & 1), + )); + values[0] = red; + values[1] = green; + values[2] = math_try!(checked_add(input[2], red)); + if p.bit_depth >= 8 && p.bit_depth <= 10 && p.red_blue_not_swapped == 0 { + let red = values[0]; + values[0] = values[2]; + values[2] = red; + } + } else if p.internal_color == 5 && p.output_color == 3 { + let black = math_try!(checked_add(input[3], input[0] >> 1)); + let magenta = math_try!(checked_sub( + math_try!(checked_sub(black, input[0])), + input[1] >> 1 + )); + let cyan = math_try!(checked_add( + math_try!(checked_add(input[1], magenta)), + input[2] >> 1 + )); + values = [cyan, magenta, math_try!(checked_sub(cyan, input[2])), black]; + } else if p.internal_color == 5 && p.output_color == 4 { + values = [input[1], input[2], input[3], input[0]]; + } + Ok(values) +} + +fn scale(sample: i32, component: u32, alpha: bool, p: JxrOutputAbi) -> Result { + let depth = p.bit_depth; + let scaled = if alpha { p.alpha_scaled } else { p.scaled }; + let shift_bits = if alpha { + p.alpha_shift_bits + } else { + p.shift_bits + }; + let mut bias = match depth { + 1 => 128, + 2 | 9 => 512, + 3 => 32768 >> shift_bits, + 8 => 16, + 10 => 32, + _ => 0, + }; + if !alpha && p.output_color == 3 { + bias = if component < 3 { + bias >> 1 + } else { + -(bias >> 1) + }; + } + if !alpha && p.output_color == 6 { + bias = 0; + } + let scale = scaled * 3; + let shifted_bias = math_try!(checked_i64_to_i32(bias as i64 * (1i64 << scale))); + let biased = math_try!(checked_add(sample, shifted_bias)); + let rounding = if scaled == 0 { + 0 + } else if depth == 0 || depth == 3 { + 4 + } else { + 3 + }; + let rounded = math_try!(checked_add(biased, rounding)); + let downshift = scale + if depth == 10 && component != 1 { 1 } else { 0 }; + let downscaled = rounded >> downshift; + if depth == 3 || depth == 4 || depth == 6 { + checked_i64_to_i32(downscaled as i64 * (1i64 << shift_bits)) + } else { + Ok(downscaled) + } +} + +fn f16_bits(sample: i32, alpha: bool, p: JxrOutputAbi) -> Result { + let scaled = math_try!(scale(sample, 0, alpha, p)); + let sign = if scaled < 0 { 0x8000 } else { 0 }; + Ok(sign | (scaled as i64).unsigned_abs().min(32767) as u16) +} + +fn f32_bits(sample: i32, alpha: bool, p: JxrOutputAbi) -> Result { + let scaled = math_try!(scale(sample, 0, alpha, p)); + let length = if alpha { + p.alpha_mantissa_length + } else { + p.mantissa_length + }; + let bias = if alpha { + p.alpha_exponent_bias_bits + } else { + p.exponent_bias_bits + } as i32; + let sign = if scaled < 0 { 0x80000000 } else { 0 }; + let magnitude = (scaled as i64).unsigned_abs(); + let implicit = 1u64 << length; + let mut exponent = (magnitude >> length) as i64; + let mut mantissa = (magnitude & (implicit - 1)) | implicit; + if exponent == 0 { + mantissa ^= implicit; + exponent = 1; + } + exponent = exponent - bias as i64 + 127; + while mantissa < implicit && exponent > 1 && mantissa > 0 { + exponent -= 1; + mantissa <<= 1; + } + if mantissa < implicit { + exponent = 0; + } else { + mantissa ^= implicit; + } + mantissa <<= 23 - length; + if exponent < 0 || exponent > 255 || mantissa > 0x7fffff { + return Err(1); + } + Ok(sign | ((exponent as u32) << 23) | mantissa as u32) +} + +fn is_alpha(channel: u32, p: JxrOutputAbi) -> bool { + p.alpha_plane != u32::MAX && channel == p.channels - 1 +} + +fn component( + samples: *const i32, + planes: *const JxrSamplePlaneAbi, + channel: u32, + x: u32, + y: u32, + p: JxrOutputAbi, +) -> Result { + if is_alpha(channel, p) { + return Ok(read_plane( + samples, + device_load(planes, p.alpha_plane), + x, + y, + )); + } + if p.output_color == 7 { + return Ok(read_plane(samples, device_load(planes, channel), x, y)); + } + if (p.channel_layout == 12 || p.channel_layout == 13) && channel == 3 { + return Ok(0); + } + let values = math_try!(converted(samples, planes, x, y, p)); + let source = if (p.channel_layout == 6 || p.channel_layout == 7 || p.channel_layout == 13) + && channel < 3 + { + 2 - channel + } else { + channel + }; + Ok(values[source as usize]) +} + +fn formatted_integer( + samples: *const i32, + planes: *const JxrSamplePlaneAbi, + channel: u32, + x: u32, + y: u32, + p: JxrOutputAbi, +) -> Result { + if (p.channel_layout == 12 || p.channel_layout == 13) && channel == 3 { + return Ok(0); + } + let alpha = is_alpha(channel, p); + let sample = math_try!(component(samples, planes, channel, x, y, p)); + let mut value = math_try!(scale(sample, channel, alpha, p)); + if p.premultiply_alpha != 0 && !alpha && p.alpha_plane != u32::MAX { + let alpha_sample = read_plane(samples, device_load(planes, p.alpha_plane), x, y); + let alpha_value = math_try!(scale(alpha_sample, 0, true, p)); + if p.bit_depth == 1 { + value = math_try!(premultiply_unsigned( + value.clamp(0, 255) as u32, + alpha_value.clamp(0, 255) as u32, + 255, + )) as i32; + } else if p.bit_depth == 2 || p.bit_depth == 3 { + value = math_try!(premultiply_unsigned( + value.clamp(0, 65535) as u32, + alpha_value.clamp(0, 65535) as u32, + 65535, + )) as i32; + } else if p.bit_depth == 4 { + value = math_try!(premultiply_signed(value, alpha_value, 32767)); + } else if p.bit_depth == 6 { + value = math_try!(premultiply_signed(value, alpha_value, i32::MAX)); + } + } + Ok(value) +} + +fn integer_value( + samples: *const i32, + planes: *const JxrSamplePlaneAbi, + channel: u32, + x_out: u32, + y_out: u32, + p: JxrOutputAbi, +) -> Result { + if p.output_plane_count > 1 { + let chroma = p.output_plane == 1 || p.output_plane == 2; + let source = if p.output_plane < 3 { + p.output_plane + } else { + p.alpha_plane + }; + let x = p.crop_x / if chroma { 2 } else { 1 } + x_out; + let y = p.crop_y + / if p.chroma_sampling == 1 && chroma { + 2 + } else { + 1 + } + + y_out; + scale( + read_plane(samples, device_load(planes, source), x, y), + p.output_plane, + p.output_plane >= 3, + p, + ) + } else { + formatted_integer( + samples, + planes, + channel, + p.crop_x + x_out, + p.crop_y + y_out, + p, + ) + } +} + +fn output_index(surface: JxrSurfacePlaneAbi, x: u32, y: u32, channel: u32, bytes: u32) -> u32 { + surface.byte_offset + y * surface.row_stride_bytes + (x * surface.channels + channel) * bytes +} + +fn store_word(output: *mut u8, index: u32, value: u32, bytes: u32) { + for byte in 0..bytes { + store(output, index + byte, (value >> (byte * 8)) as u8); + } +} + +pub fn write( + samples: *const i32, + planes: *const JxrSamplePlaneAbi, + surface: JxrSurfacePlaneAbi, + output: *mut u8, + p: JxrOutputAbi, + x_out: u32, + y_out: u32, + kind: u32, +) -> Result<(), MathError> { + let x = p.crop_x + x_out; + let y = p.crop_y + y_out; + if kind == 0 { + let mut byte = 0u8; + for bit in 0..8 { + let pixel_x = x_out * 8 + bit; + if pixel_x >= surface.width { + break; + } + let sample = read_plane(samples, device_load(planes, 0), p.crop_x + pixel_x, y); + let mut packed = math_try!(scale(sample, 0, false, p)).clamp(0, 1) as u32; + if p.bit_black != 0 { + packed = 1 - packed; + } + byte |= (packed << (7 - bit)) as u8; + } + store( + output, + surface.byte_offset + y_out * surface.row_stride_bytes + x_out, + byte, + ); + } else if kind <= 4 { + let bytes = if kind == 1 { + 1 + } else if kind == 4 { + 4 + } else { + 2 + }; + let channels = if p.output_plane_count > 1 { + 1 + } else { + surface.channels + }; + for channel in 0..channels { + let value = math_try!(integer_value(samples, planes, channel, x_out, y_out, p)); + let (minimum, maximum) = match kind { + 1 => (0, 255), + 2 => (0, if p.bit_depth == 2 { 1023 } else { 65535 }), + 3 => (-32768, 32767), + _ => (i32::MIN, i32::MAX), + }; + store_word( + output, + output_index(surface, x_out, y_out, channel, bytes), + value.clamp(minimum, maximum) as u32, + bytes, + ); + } + } else if kind == 5 || kind == 6 { + let alpha_sample = if p.alpha_plane == u32::MAX { + 0 + } else { + read_plane(samples, device_load(planes, p.alpha_plane), x, y) + }; + let alpha_bits = if kind == 5 && p.alpha_plane != u32::MAX { + math_try!(f16_bits(alpha_sample, true, p)) + } else { + 0 + }; + let alpha_value = if kind == 6 && p.alpha_plane != u32::MAX { + f32::from_bits(math_try!(f32_bits(alpha_sample, true, p))).clamp(0.0, 1.0) + } else { + 1.0 + }; + let bytes = if kind == 5 { 2 } else { 4 }; + for channel in 0..p.channels { + let index = output_index(surface, x_out, y_out, channel, bytes); + if (p.channel_layout == 12 || p.channel_layout == 13) && channel == 3 { + store_word(output, index, 0, bytes); + continue; + } + let alpha = is_alpha(channel, p); + let sample = math_try!(component(samples, planes, channel, x, y, p)); + let bits = if kind == 5 { + let bits = math_try!(f16_bits(sample, alpha, p)); + if p.premultiply_alpha != 0 && !alpha { + math_try!(premultiply_sign_magnitude_15(bits, alpha_bits)) as u32 + } else { + bits as u32 + } + } else { + let mut value = f32::from_bits(math_try!(f32_bits(sample, alpha, p))); + if p.premultiply_alpha != 0 && !alpha { + value *= alpha_value; + } + value.to_bits() + }; + store_word(output, index, bits, bytes); + } + } else { + let values = math_try!(converted(samples, planes, x, y, p)); + let mut packed = 0; + if kind == 8 && p.output_color == 6 { + let mut scaled = [0; 3]; + for channel in 0..3 { + scaled[channel as usize] = + math_try!(scale(values[channel as usize], channel, false, p)); + } + packed = pack_rgbe(scaled); + } else { + for channel in 0..3 { + let value = math_try!(scale(values[channel as usize], channel, false, p)); + let maximum = if kind == 8 { + 1023 + } else if p.bit_depth == 10 && channel == 1 { + 63 + } else { + 31 + }; + let shift = if kind == 8 { + (2 - channel) * 10 + } else if p.bit_depth == 10 { + if channel == 0 { + 11 + } else if channel == 1 { + 5 + } else { + 0 + } + } else { + (2 - channel) * 5 + }; + packed |= (value.clamp(0, maximum) as u32) << shift; + } + } + let bytes = if kind == 7 { 2 } else { 4 }; + store_word( + output, + surface.byte_offset + y_out * surface.row_stride_bytes + x_out * bytes, + packed, + bytes, + ); + } + Ok(()) +} + +fn premultiply_unsigned(value: u32, alpha: u32, maximum: u32) -> Result { + if maximum == 0 { + return Err(1); + } + Ok( + ((value.min(maximum) as u64 * alpha.min(maximum) as u64 + (maximum / 2) as u64) + / maximum as u64) as u32, + ) +} + +fn premultiply_signed(value: i32, alpha: i32, maximum: i32) -> Result { + if maximum <= 0 { + return Err(1); + } + let magnitude = (value as i64).unsigned_abs(); + let result = ((magnitude * alpha.clamp(0, maximum) as u64 + (maximum / 2) as u64) + / maximum as u64) as i64; + checked_i64_to_i32(if value < 0 { -result } else { result }) +} + +fn premultiply_sign_magnitude_15(value: u16, alpha: u16) -> Result { + let magnitude = math_try!(premultiply_unsigned( + (value & 0x7fff) as u32, + if alpha & 0x8000 == 0 { + (alpha & 0x7fff) as u32 + } else { + 0 + }, + 0x7fff + )); + Ok((value & 0x8000) | magnitude as u16) +} + +fn pack_rgbe(rgb: [i32; 3]) -> u32 { + let mut exponent = 0; + let mut mantissa = [0u32; 3]; + let mut local_exponent = [0u32; 3]; + for channel in 0..3 { + let value = rgb[channel]; + if value <= 0 { + mantissa[channel] = 0; + local_exponent[channel] = 0; + } else if value >> 7 > 1 { + mantissa[channel] = ((value & 127) + 128) as u32; + local_exponent[channel] = (value >> 7) as u32; + } else { + mantissa[channel] = value as u32; + local_exponent[channel] = 1; + } + exponent = exponent.max(local_exponent[channel]); + } + for channel in 0..3 { + if exponent > local_exponent[channel] { + let difference = exponent - local_exponent[channel]; + mantissa[channel] = if difference >= 31 { + 0 + } else { + (2 * mantissa[channel] + 1) >> (difference + 1) + }; + } + } + mantissa[0].min(255) + | (mantissa[1].min(255) << 8) + | (mantissa[2].min(255) << 16) + | (exponent.min(255) << 24) +} diff --git a/crates/jxr-cuda/src/kernels/overlap.cu b/crates/jxr-cuda/src/kernels/overlap.cu deleted file mode 100644 index 6d5a0e3..0000000 --- a/crates/jxr-cuda/src/kernels/overlap.cu +++ /dev/null @@ -1,242 +0,0 @@ -// SPDX-License-Identifier: MIT OR Apache-2.0 - -__device__ __forceinline__ bool jxr_inverse_rotate( - int v[2], uint *status, uint code) { - int temporary; - return jxr_add(v[1], 1, temporary, status, code) - && jxr_sub(v[0], temporary >> 1, v[0], status, code) - && jxr_add(v[0], 1, temporary, status, code) - && jxr_add(v[1], temporary >> 1, v[1], status, code); -} - -__device__ __forceinline__ bool jxr_inverse_scale( - int v[2], uint *status, uint code) { - int temporary; - return jxr_add(v[0], v[1], v[0], status, code) - && jxr_sub(v[0] >> 1, v[1], v[1], status, code) - && jxr_narrow((jxr_long)v[1] * 3ll, temporary, status, code) - && jxr_add(v[0], temporary >> 3, v[0], status, code) - && jxr_narrow((jxr_long)v[0] * 3ll, temporary, status, code) - && jxr_add(v[1], temporary >> 4, v[1], status, code) - && jxr_add(v[1], v[0] >> 7, v[1], status, code) - && jxr_sub(v[1], v[0] >> 10, v[1], status, code); -} - -__device__ __forceinline__ bool jxr_inverse_hadamard_post( - int v[4], uint *status, uint code) { - int temporary; - return jxr_sub(v[1], v[2], v[1], status, code) - && jxr_mul3_round(v[3], 4, 3, temporary, status, code) - && jxr_add(v[0], temporary, v[0], status, code) - && jxr_sub(v[3], v[1] >> 1, v[3], status, code) - && jxr_sub(v[0], v[1], temporary, status, code) - && jxr_sub(temporary >> 1, v[2], v[2], status, code) - && (temporary = v[2], v[2] = v[3], v[3] = temporary, true) - && jxr_sub(v[0], v[3], v[0], status, code) - && jxr_add(v[1], v[2], v[1], status, code); -} - -__device__ __forceinline__ bool jxr_inverse_todd_odd_post( - int v[4], uint *status, uint code) { - int first; - int second; - int temporary; - return jxr_add(v[3], v[0], v[3], status, code) - && jxr_sub(v[2], v[1], v[2], status, code) - && (first = v[3] >> 1, second = v[2] >> 1, true) - && jxr_sub(v[0], first, v[0], status, code) - && jxr_add(v[1], second, v[1], status, code) - && jxr_mul3_round(v[1], 6, 3, temporary, status, code) - && jxr_sub(v[0], temporary, v[0], status, code) - && jxr_mul3_round(v[0], 2, 2, temporary, status, code) - && jxr_add(v[1], temporary, v[1], status, code) - && jxr_mul3_round(v[1], 4, 3, temporary, status, code) - && jxr_sub(v[0], temporary, v[0], status, code) - && jxr_sub(v[1], second, v[1], status, code) - && jxr_add(v[0], first, v[0], status, code) - && jxr_add(v[2], v[1], v[2], status, code) - && jxr_sub(v[3], v[0], v[3], status, code); -} - -__device__ __forceinline__ bool jxr_overlap4( - int v[4], uint *status, uint code) { - int temporary; - int pair[2]; - return jxr_add(v[0], v[3], v[0], status, code) - && jxr_add(v[1], v[2], v[1], status, code) - && jxr_add(v[0], 1, temporary, status, code) - && jxr_sub(v[3], temporary >> 1, v[3], status, code) - && jxr_add(v[1], 1, temporary, status, code) - && jxr_sub(v[2], temporary >> 1, v[2], status, code) - && (pair[0] = v[0], pair[1] = v[3], true) - && jxr_inverse_scale(pair, status, code) - && (v[0] = pair[0], v[3] = pair[1], pair[0] = v[1], pair[1] = v[2], true) - && jxr_inverse_scale(pair, status, code) - && (v[1] = pair[0], v[2] = pair[1], true) - && jxr_mul3_round(v[3], 4, 3, temporary, status, code) - && jxr_add(v[0], temporary, v[0], status, code) - && jxr_mul3_round(v[2], 4, 3, temporary, status, code) - && jxr_add(v[1], temporary, v[1], status, code) - && jxr_sub(v[3], v[0] >> 1, v[3], status, code) - && jxr_sub(v[2], v[1] >> 1, v[2], status, code) - && jxr_add(v[0], v[3], v[0], status, code) - && jxr_add(v[1], v[2], v[1], status, code) - && jxr_narrow(-(jxr_long)v[3], v[3], status, code) - && jxr_narrow(-(jxr_long)v[2], v[2], status, code) - && (pair[0] = v[2], pair[1] = v[3], true) - && jxr_inverse_rotate(pair, status, code) - && (v[2] = pair[0], v[3] = pair[1], true) - && jxr_add(v[0], 1, temporary, status, code) - && jxr_add(v[3], temporary >> 1, v[3], status, code) - && jxr_add(v[1], 1, temporary, status, code) - && jxr_add(v[2], temporary >> 1, v[2], status, code) - && jxr_sub(v[0], v[3], v[0], status, code) - && jxr_sub(v[1], v[2], v[1], status, code); -} - -__device__ __forceinline__ bool jxr_overlap2( - int v[2], uint *status, uint code) { - int temporary; - return jxr_add(v[0], 2, temporary, status, code) - && jxr_add(v[1], temporary >> 2, v[1], status, code) - && jxr_add(v[1], 1, temporary, status, code) - && jxr_add(v[0], temporary >> 1, v[0], status, code) - && jxr_add(v[0], v[1] >> 5, v[0], status, code) - && jxr_add(v[0], v[1] >> 9, v[0], status, code) - && jxr_add(v[0], v[1] >> 13, v[0], status, code) - && jxr_add(v[0], 2, temporary, status, code) - && jxr_add(v[1], temporary >> 2, v[1], status, code); -} - -__device__ __forceinline__ bool jxr_overlap2x2( - int v[4], uint *status, uint code) { - int temporary; - return jxr_add(v[0], v[3], v[0], status, code) - && jxr_add(v[1], v[2], v[1], status, code) - && jxr_add(v[0], 1, temporary, status, code) - && jxr_sub(v[3], temporary >> 1, v[3], status, code) - && jxr_add(v[1], 1, temporary, status, code) - && jxr_sub(v[2], temporary >> 1, v[2], status, code) - && jxr_add(v[0], 2, temporary, status, code) - && jxr_add(v[1], temporary >> 2, v[1], status, code) - && jxr_add(v[1], 1, temporary, status, code) - && jxr_add(v[0], temporary >> 1, v[0], status, code) - && jxr_add(v[0], v[1] >> 5, v[0], status, code) - && jxr_add(v[0], v[1] >> 9, v[0], status, code) - && jxr_add(v[0], v[1] >> 13, v[0], status, code) - && jxr_add(v[0], 2, temporary, status, code) - && jxr_add(v[1], temporary >> 2, v[1], status, code) - && jxr_add(v[0], 1, temporary, status, code) - && jxr_add(v[3], temporary >> 1, v[3], status, code) - && jxr_add(v[1], 1, temporary, status, code) - && jxr_add(v[2], temporary >> 1, v[2], status, code) - && jxr_sub(v[0], v[3], v[0], status, code) - && jxr_sub(v[1], v[2], v[1], status, code); -} - -__device__ __forceinline__ bool jxr_group4x4( - int values[16], JxrUint4 indices, uint kind, uint *status, uint code) { - int group[4] = { - values[indices.x], values[indices.y], values[indices.z], values[indices.w]}; - bool ok = kind == 0u ? jxr_t2x2h(group, 0, status, code) - : (kind == 1u ? jxr_inverse_todd_odd_post(group, status, code) - : jxr_inverse_hadamard_post(group, status, code)); - if (!ok) return false; - values[indices.x] = group[0]; - values[indices.y] = group[1]; - values[indices.z] = group[2]; - values[indices.w] = group[3]; - return true; -} - -__device__ __forceinline__ bool jxr_overlap4x4( - int values[16], uint *status, uint code) { - const JxrUint4 groups[4] = { - {0, 3, 12, 15}, {1, 2, 13, 14}, {4, 7, 8, 11}, {5, 6, 9, 10}}; - for (uint i = 0; i < 4; ++i) { - if (!jxr_group4x4(values, groups[i], 0, status, code)) return false; - } - const uint rotations[8] = {13, 12, 9, 8, 7, 3, 6, 2}; - for (uint i = 0; i < 4; ++i) { - int pair[2] = {values[rotations[i * 2]], values[rotations[i * 2 + 1]]}; - if (!jxr_inverse_rotate(pair, status, code)) return false; - values[rotations[i * 2]] = pair[0]; - values[rotations[i * 2 + 1]] = pair[1]; - } - if (!jxr_group4x4(values, jxr_uint4(10, 11, 14, 15), 1, status, code)) return false; - const uint scales[8] = {0, 15, 1, 14, 4, 11, 5, 10}; - for (uint i = 0; i < 4; ++i) { - int pair[2] = {values[scales[i * 2]], values[scales[i * 2 + 1]]}; - if (!jxr_inverse_scale(pair, status, code)) return false; - values[scales[i * 2]] = pair[0]; - values[scales[i * 2 + 1]] = pair[1]; - } - for (uint i = 0; i < 4; ++i) { - if (!jxr_group4x4(values, groups[i], 2, status, code)) return false; - } - return true; -} - -__device__ __forceinline__ void jxr_apply_overlap( - int *samples, JxrOverlapWorkAbi work, uint *status) { - if (work.kind == 0u) { - int values[16]; - for (uint y = 0; y < 4; ++y) { - for (uint x = 0; x < 4; ++x) { - values[y * 4 + x] = samples[work.first + y * work.second + x]; - } - } - if (!jxr_overlap4x4(values, status, 2u)) return; - for (uint y = 0; y < 4; ++y) { - for (uint x = 0; x < 4; ++x) { - samples[work.first + y * work.second + x] = values[y * 4 + x]; - } - } - } else if (work.kind == 1u) { - int values[4] = { - samples[work.first], samples[work.first + work.second], - samples[work.first + work.second * 2u], samples[work.first + work.second * 3u]}; - if (!jxr_overlap4(values, status, 2u)) return; - for (uint i = 0; i < 4; ++i) samples[work.first + work.second * i] = values[i]; - } else if (work.kind == 2u || work.kind == 3u) { - int values[4] = { - samples[work.first], samples[work.first + 1u], - samples[work.first + work.second], samples[work.first + work.second + 1u]}; - bool ok = work.kind == 2u ? jxr_overlap4(values, status, 2u) - : jxr_overlap2x2(values, status, 2u); - if (!ok) return; - samples[work.first] = values[0]; - samples[work.first + 1u] = values[1]; - samples[work.first + work.second] = values[2]; - samples[work.first + work.second + 1u] = values[3]; - } else if (work.kind == 4u) { - int values[2] = {samples[work.first], samples[work.second]}; - if (!jxr_overlap2(values, status, 2u)) return; - samples[work.first] = values[0]; - samples[work.second] = values[1]; - } else if (work.kind == 5u) { - int result; - if (jxr_sub(samples[work.first], samples[work.second], result, status, 2u)) { - samples[work.first] = result; - } - } else if (work.kind == 6u) { - int result; - if (jxr_add(samples[work.first], samples[work.second], result, status, 2u)) { - samples[work.first] = result; - } - } -} - -extern "C" __global__ void jxr_first_overlap( - int *samples, const JxrOverlapWorkAbi *work, uint *status, uint work_count) { - const uint gid = blockIdx.x * blockDim.x + threadIdx.x; - const bool failed = jxr_block_failed(status); - if (gid < work_count && !failed) jxr_apply_overlap(samples, work[gid], status); -} - -extern "C" __global__ void jxr_second_overlap( - int *samples, const JxrOverlapWorkAbi *work, uint *status, uint work_count) { - const uint gid = blockIdx.x * blockDim.x + threadIdx.x; - const bool failed = jxr_block_failed(status); - if (gid < work_count && !failed) jxr_apply_overlap(samples, work[gid], status); -} diff --git a/crates/jxr-cuda/src/kernels/reconstruction.rs b/crates/jxr-cuda/src/kernels/reconstruction.rs new file mode 100644 index 0000000..40299a9 --- /dev/null +++ b/crates/jxr-cuda/src/kernels/reconstruction.rs @@ -0,0 +1,225 @@ +// SPDX-License-Identifier: MIT OR Apache-2.0 + +use super::*; +use jxr_math::{overlap::*, transform::*}; + +pub fn first_transform( + packed: *const i32, + macroblocks: *const JxrMacroblockAbi, + low_plane: *mut i32, + plane: JxrPlaneAbi, + gid: u32, +) -> Result<(), MathError> { + let metadata = device_load(macroblocks, plane.macroblock_offset + gid); + let count = plane.block_columns * plane.block_rows; + let mut low = [0; 16]; + low[0] = math_try!(dequantize( + device_load(packed, metadata.coefficient_offset), + metadata.quantizer_dc, + )); + if metadata.bands != 0 { + for block in 1..count { + let source = metadata.coefficient_offset + + if metadata.bands >= 2 { + block * 16 + } else { + block + }; + low[block as usize] = math_try!(dequantize( + device_load(packed, source), + metadata.quantizer_low_pass + )); + } + } + if plane.block_columns == 4 { + math_try!(inverse_core_transform(&mut low)); + } else if plane.block_rows == 2 { + let mut values = [low[0], low[1], low[2], low[3]]; + math_try!(inverse_chroma_420(&mut values)); + low[..4].copy_from_slice(&values); + } else { + let mut values = [ + low[0], low[1], low[2], low[3], low[4], low[5], low[6], low[7], + ]; + math_try!(inverse_chroma_422(&mut values)); + low[..8].copy_from_slice(&values); + } + if plane.scale_after_first_transform != 0 { + for block in 0..count { + low[block as usize] = math_try!(checked_mul(low[block as usize], 2)); + } + } + let base_x = (metadata.coded_x - plane.macroblock_origin_x) * plane.block_columns; + let base_y = (metadata.coded_y - plane.macroblock_origin_y) * plane.block_rows; + for row in 0..plane.block_rows { + for column in 0..plane.block_columns { + store( + low_plane, + plane.low_offset + (base_y + row) * plane.low_width + base_x + column, + low[(row * plane.block_columns + column) as usize], + ); + } + } + Ok(()) +} + +pub fn predict_high( + high: *mut i32, + metadata: JxrMacroblockAbi, + plane: JxrPlaneAbi, +) -> Result<(), MathError> { + let count = plane.block_columns * plane.block_rows; + if metadata.bands >= 2 { + if metadata.hp_prediction == 1 { + for block in 1..count { + if block % plane.block_columns != 0 { + for step in 1..4 { + let coefficient = step * 4; + let destination = block * 16 + coefficient; + store( + high, + destination, + math_try!(checked_add( + device_load(high, destination), + device_load(high, destination - 16), + )), + ); + } + } + } + } else if metadata.hp_prediction == 2 { + for block in plane.block_columns..count { + for coefficient in 1..4 { + let destination = block * 16 + coefficient; + let source = destination - plane.block_columns * 16; + store( + high, + destination, + math_try!(checked_add( + device_load(high, destination), + device_load(high, source) + )), + ); + } + } + } + } + Ok(()) +} + +pub fn second_transform( + high: *const i32, + low_plane: *const i32, + samples: *mut i32, + metadata: JxrMacroblockAbi, + plane: JxrPlaneAbi, + block: u32, +) -> Result<(), MathError> { + let mut coefficients = [0; 16]; + for coefficient in 1..16 { + coefficients[coefficient as usize] = math_try!(dequantize( + device_load(high, block * 16 + coefficient), + metadata.quantizer_high_pass, + )); + } + let local_x = metadata.coded_x - plane.macroblock_origin_x; + let local_y = metadata.coded_y - plane.macroblock_origin_y; + let low_x = local_x * plane.block_columns + block % plane.block_columns; + let low_y = local_y * plane.block_rows + block / plane.block_columns; + coefficients[0] = device_load( + low_plane, + plane.low_offset + low_y * plane.low_width + low_x, + ); + math_try!(inverse_core_transform(&mut coefficients)); + let output_x = low_x * 4; + let output_y = low_y * 4; + for row in 0..4 { + for column in 0..4 { + store( + samples, + plane.sample_offset + (output_y + row) * plane.sample_width + output_x + column, + coefficients[(row * 4 + column) as usize], + ); + } + } + Ok(()) +} + +pub fn overlap(samples: *mut i32, work: JxrOverlapWorkAbi) -> Result<(), MathError> { + if work.kind == 0 { + let mut values = [0; 16]; + for y in 0..4 { + for x in 0..4 { + values[(y * 4 + x) as usize] = + device_load(samples, work.first + y * work.second + x); + } + } + math_try!(inverse_overlap_filter_4x4(&mut values)); + for y in 0..4 { + for x in 0..4 { + store( + samples, + work.first + y * work.second + x, + values[(y * 4 + x) as usize], + ); + } + } + } else if work.kind <= 3 { + let indices = if work.kind == 1 { + [ + work.first, + work.first + work.second, + work.first + work.second * 2, + work.first + work.second * 3, + ] + } else { + [ + work.first, + work.first + 1, + work.first + work.second, + work.first + work.second + 1, + ] + }; + let mut values = [ + device_load(samples, indices[0]), + device_load(samples, indices[1]), + device_load(samples, indices[2]), + device_load(samples, indices[3]), + ]; + if work.kind == 3 { + math_try!(inverse_overlap_filter_2x2(&mut values)); + } else { + math_try!(inverse_overlap_filter_4(&mut values)); + } + for i in 0..4 { + store(samples, indices[i], values[i]); + } + } else if work.kind == 4 { + let mut values = [ + device_load(samples, work.first), + device_load(samples, work.second), + ]; + math_try!(inverse_overlap_filter_2(&mut values)); + store(samples, work.first, values[0]); + store(samples, work.second, values[1]); + } else if work.kind == 5 { + store( + samples, + work.first, + math_try!(checked_sub( + device_load(samples, work.first), + device_load(samples, work.second), + )), + ); + } else if work.kind == 6 { + store( + samples, + work.first, + math_try!(checked_add( + device_load(samples, work.first), + device_load(samples, work.second), + )), + ); + } + Ok(()) +} diff --git a/crates/jxr-cuda/src/kernels/rust-toolchain.toml b/crates/jxr-cuda/src/kernels/rust-toolchain.toml new file mode 100644 index 0000000..13204e7 --- /dev/null +++ b/crates/jxr-cuda/src/kernels/rust-toolchain.toml @@ -0,0 +1,3 @@ +[toolchain] +channel = "nightly-2026-04-03" +components = ["rust-src", "rustc-dev", "llvm-tools"] diff --git a/crates/jxr-cuda/src/kernels/simt.toml b/crates/jxr-cuda/src/kernels/simt.toml new file mode 100644 index 0000000..cd8552f --- /dev/null +++ b/crates/jxr-cuda/src/kernels/simt.toml @@ -0,0 +1,13 @@ +[package] +name = "jxr_reconstruction" +version = "0.1.0" +edition = "2024" + +[workspace] +resolver = "2" + +[dependencies] +cuda-core = { git = "https://github.com/NVlabs/cuda-oxide.git", rev = "a9f964a956f397dd0b3c8db88a3ca5824186c261" } +cuda-device = { git = "https://github.com/NVlabs/cuda-oxide.git", rev = "a9f964a956f397dd0b3c8db88a3ca5824186c261" } +cuda-host = { git = "https://github.com/NVlabs/cuda-oxide.git", rev = "a9f964a956f397dd0b3c8db88a3ca5824186c261" } +jxr-math = { path = "__JXR_MATH__" } diff --git a/crates/jxr-cuda/src/plan/tests.rs b/crates/jxr-cuda/src/plan/tests.rs index 311d025..e352632 100644 --- a/crates/jxr-cuda/src/plan/tests.rs +++ b/crates/jxr-cuda/src/plan/tests.rs @@ -13,6 +13,28 @@ use jxr_core::{ use super::{CudaDecodePlan, build_arenas}; use crate::CudaError; +#[test] +#[ignore = "requires compiled CUDA Oxide kernels and an NVIDIA GPU"] +fn device_overflow_is_reported_and_does_not_poison_the_next_submission() { + let session = crate::CudaDecoderSession::system_default().unwrap(); + let mut arena = luma_arena(&[(0, 0)]); + let input = Arc::get_mut(&mut arena).unwrap(); + input.coefficients[0] = i32::MAX; + input.macroblocks.quantizers[0].dc = 2; + let prepared = prepared_plan(1, 1); + let plan = cuda_plan(arena, &prepared).unwrap(); + assert!(matches!( + session.decode_to_host(&plan), + Err(CudaError::KernelArithmetic { status: 1 }) + )); + let valid = cuda_plan(luma_arena(&[(0, 0)]), &prepared).unwrap(); + let decoded = session.decode_to_host(&valid).unwrap(); + assert_eq!( + decoded.samples, + jxr_core::DecodedSamples::U8(vec![128; 256]) + ); +} + #[test] fn rejects_invalid_prediction_before_any_cuda_resource_can_be_created() { let mut arena = luma_arena(&[(0, 0)]); diff --git a/crates/jxr-cuda/src/runtime.rs b/crates/jxr-cuda/src/runtime.rs index 9c031f3..89e4ae5 100644 --- a/crates/jxr-cuda/src/runtime.rs +++ b/crates/jxr-cuda/src/runtime.rs @@ -8,7 +8,7 @@ use std::{ use cudarc::{ driver::{CudaContext, CudaFunction, CudaModule, CudaStream}, - nvrtc::{CompileOptions, compile_ptx_with_opts}, + nvrtc::Ptx, }; use crate::{CudaError, buffer_pool::CudaBufferPool, kernels, upload_cache::CudaUploadCache}; @@ -59,32 +59,17 @@ impl CudaRuntime { fn build_inner(device_ordinal: usize) -> Result, CudaError> { let context = CudaContext::new(device_ordinal)?; - let (major, minor) = context.compute_capability()?; + let (major, _) = context.compute_capability()?; if major < 5 { return Err(CudaError::UnsupportedDevice { reason: "JPEG XR CUDA reconstruction requires compute capability 5.0 or newer", }); } - let mut options = CompileOptions { - fmad: Some(false), - name: Some("jxr_reconstruction.cu".to_owned()), - ..CompileOptions::default() - }; - options.options.push("--std=c++17".to_owned()); - options - .options - .push(format!("--gpu-architecture=compute_{major}{minor}")); - let compiled = compile_ptx_with_opts(kernels::source(), options).map_err(|error| { - CudaError::RuntimeInitialization { + let module = context + .load_module(Ptx::from_src(kernels::PTX)) + .map_err(|error| CudaError::RuntimeInitialization { message: error.to_string(), - } - })?; - let module = - context - .load_module(compiled) - .map_err(|error| CudaError::RuntimeInitialization { - message: error.to_string(), - })?; + })?; let mut functions = HashMap::with_capacity(ENTRYPOINTS.len()); for entrypoint in ENTRYPOINTS { let function = module.load_function(entrypoint).map_err(|error| { diff --git a/crates/jxr-cuda/src/session.rs b/crates/jxr-cuda/src/session.rs index 5476887..6e3f1f4 100644 --- a/crates/jxr-cuda/src/session.rs +++ b/crates/jxr-cuda/src/session.rs @@ -28,14 +28,14 @@ impl core::fmt::Debug for CudaDecoderSession { } impl CudaDecoderSession { - /// Probe for both dynamically loaded CUDA driver and NVRTC libraries. + /// Probe for compiled CUDA Oxide kernels and the dynamically loaded CUDA driver. /// /// This performs no CUDA initialization and is safe on CPU-only hosts. #[must_use] pub fn is_available() -> bool { // SAFETY: These cudarc probes only attempt to open candidate shared // libraries and do not retain or dereference foreign handles. - unsafe { cudarc::driver::sys::is_culib_present() && cudarc::nvrtc::sys::is_culib_present() } + cfg!(jxr_cuda_oxide) && unsafe { cudarc::driver::sys::is_culib_present() } } /// Create a reusable session on device ordinal zero. diff --git a/crates/jxr-math/Cargo.toml b/crates/jxr-math/Cargo.toml index 9b54576..ea4fe25 100644 --- a/crates/jxr-math/Cargo.toml +++ b/crates/jxr-math/Cargo.toml @@ -3,10 +3,12 @@ name = "jxr-math" description = "Exact integer building blocks for JPEG XR decoding" version.workspace = true edition.workspace = true -rust-version.workspace = true +# Shared device code must support CUDA Oxide's pinned compiler. +rust-version = "1.96" license.workspace = true repository.workspace = true homepage.workspace = true +links = "jxr_math" [lints] workspace = true diff --git a/crates/jxr-math/build.rs b/crates/jxr-math/build.rs index 516c50c..8f42853 100644 --- a/crates/jxr-math/build.rs +++ b/crates/jxr-math/build.rs @@ -3,6 +3,15 @@ use std::{env, fmt::Write, fs, io, path::PathBuf}; const INPUT: &str = "data/reconstruction.abi"; fn main() -> io::Result<()> { + println!( + "cargo::metadata=manifest_dir={}", + PathBuf::from(required_env("CARGO_MANIFEST_DIR")?).display() + ); + println!("cargo:rustc-check-cfg=cfg(jxr_math_device)"); + println!("cargo:rerun-if-env-changed=JXR_MATH_DEVICE"); + if env::var_os("JXR_MATH_DEVICE").is_some() { + println!("cargo:rustc-cfg=jxr_math_device"); + } println!("cargo:rerun-if-changed={INPUT}"); let source = fs::read_to_string(INPUT)?; let permutation = parse_permutation(&source)?; @@ -61,20 +70,24 @@ fn generate( writeln!(generated, "\";")?; writeln!( generated, - "pub const CUDA_RECONSTRUCTION_CONSTANTS: &str = r\"" - )?; - writeln!( - generated, - "__device__ __constant__ unsigned int JXR_INVERSE_PERMUTATION[16] = {{{values}}};" + "/// Canonical C-layout descriptors shared with CUDA Oxide kernels." )?; + writeln!(generated, "pub mod reconstruction_abi {{")?; for abi in abi_structs { - writeln!(generated, "struct {} {{", abi.name)?; + writeln!(generated, " /// Generated reconstruction descriptor.")?; + writeln!(generated, " #[repr(C)]")?; + writeln!( + generated, + " #[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]" + )?; + writeln!(generated, " pub struct {} {{", abi.name)?; for field in &abi.fields { - writeln!(generated, " unsigned int {field};")?; + writeln!(generated, " /// Canonical ABI word.")?; + writeln!(generated, " pub {field}: u32,")?; } - writeln!(generated, "}};")?; + writeln!(generated, " }}")?; } - writeln!(generated, "\";")?; + writeln!(generated, "}}")?; Ok(generated) } diff --git a/crates/jxr-math/data/reconstruction.abi b/crates/jxr-math/data/reconstruction.abi index cd7519c..3fd001e 100644 --- a/crates/jxr-math/data/reconstruction.abi +++ b/crates/jxr-math/data/reconstruction.abi @@ -1,4 +1,4 @@ -# Canonical reconstruction constants shared by Rust, CUDA C, and Metal. +# Canonical reconstruction constants shared by Rust, CUDA Oxide, and Metal. inverse_permutation = 0,8,4,13,2,15,3,14,1,12,5,9,7,11,6,10 abi.JxrMacroblockAbi = coefficient_offset:u32,quantizer_dc:u32,quantizer_low_pass:u32,quantizer_high_pass:u32,bands:u32,hp_prediction:u32,coded_x:u32,coded_y:u32 abi.JxrPlaneAbi = macroblock_offset:u32,macroblock_count:u32,block_columns:u32,block_rows:u32,macroblock_origin_x:u32,macroblock_origin_y:u32,low_offset:u32,low_width:u32,sample_offset:u32,sample_width:u32,sample_height:u32,scale_after_first_transform:u32 diff --git a/crates/jxr-math/src/alpha.rs b/crates/jxr-math/src/alpha.rs index 7b2ec41..05282d1 100644 --- a/crates/jxr-math/src/alpha.rs +++ b/crates/jxr-math/src/alpha.rs @@ -5,19 +5,21 @@ use crate::arithmetic::MathError; /// Premultiply an unsigned channel with nearest rounding. pub fn premultiply_unsigned(value: u32, alpha: u32, maximum: u32) -> Result { if maximum == 0 { - return Err(MathError::overflow("zero alpha normalization range")); + return Err(crate::arithmetic::overflow( + "zero alpha normalization range", + )); } let value = value.min(maximum); let alpha = alpha.min(maximum); let product = u64::from(value) * u64::from(alpha); u32::try_from((product + u64::from(maximum / 2)) / u64::from(maximum)) - .map_err(|_| MathError::overflow("unsigned alpha premultiplication")) + .map_err(|_| crate::arithmetic::overflow("unsigned alpha premultiplication")) } /// Premultiply a signed channel by a nonnegative signed alpha channel. pub fn premultiply_signed(value: i32, alpha: i32, alpha_maximum: i32) -> Result { if alpha_maximum <= 0 { - return Err(MathError::overflow( + return Err(crate::arithmetic::overflow( "invalid signed alpha normalization range", )); } @@ -27,9 +29,9 @@ pub fn premultiply_signed(value: i32, alpha: i32, alpha_maximum: i32) -> Result< let maximum = u64::try_from(alpha_maximum).expect("validated alpha maximum is positive"); let magnitude = (product + maximum / 2) / maximum; let signed = i64::try_from(magnitude) - .map_err(|_| MathError::overflow("signed alpha premultiplication"))?; + .map_err(|_| crate::arithmetic::overflow("signed alpha premultiplication"))?; let signed = if value < 0 { -signed } else { signed }; - i32::try_from(signed).map_err(|_| MathError::overflow("signed alpha premultiplication")) + i32::try_from(signed).map_err(|_| crate::arithmetic::overflow("signed alpha premultiplication")) } /// Premultiply one T.832 BD16F sign-magnitude value by another used as alpha. diff --git a/crates/jxr-math/src/arithmetic.rs b/crates/jxr-math/src/arithmetic.rs index 89b8787..46cabc3 100644 --- a/crates/jxr-math/src/arithmetic.rs +++ b/crates/jxr-math/src/arithmetic.rs @@ -1,13 +1,22 @@ //! Checked and saturating integer arithmetic with explicit rounding semantics. +#[cfg(not(jxr_math_device))] use core::fmt; /// An integer operation could not be represented in the requested type. +#[cfg(not(jxr_math_device))] #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub struct MathError { operation: &'static str, } +// CUDA Oxide currently requires primitive enum payloads. Device callers record +// the failed stage in their status word; host errors retain operation strings. +/// Numeric arithmetic error used by the standalone CUDA Oxide device build. +#[cfg(jxr_math_device)] +pub type MathError = u32; + +#[cfg(not(jxr_math_device))] impl MathError { /// Create an overflow error for `operation`. #[must_use] @@ -22,33 +31,56 @@ impl MathError { } } +#[cfg(not(jxr_math_device))] impl fmt::Display for MathError { fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { write!(formatter, "integer overflow during {}", self.operation) } } +/// Construct an arithmetic error for host or device code. +#[must_use] +#[inline] +pub(crate) const fn overflow(operation: &'static str) -> MathError { + #[cfg(not(jxr_math_device))] + { + MathError::overflow(operation) + } + #[cfg(jxr_math_device)] + { + let _ = operation; + 1 + } +} + /// Add two signed coefficients, returning an error on overflow. +#[inline] pub fn checked_add(left: i32, right: i32) -> Result { left.checked_add(right) - .ok_or_else(|| MathError::overflow("addition")) + .ok_or(crate::arithmetic::overflow("addition")) } /// Subtract two signed coefficients, returning an error on overflow. +#[inline] pub fn checked_sub(left: i32, right: i32) -> Result { left.checked_sub(right) - .ok_or_else(|| MathError::overflow("subtraction")) + .ok_or(crate::arithmetic::overflow("subtraction")) } /// Multiply two signed coefficients, returning an error on overflow. +#[inline] pub fn checked_mul(left: i32, right: i32) -> Result { left.checked_mul(right) - .ok_or_else(|| MathError::overflow("multiplication")) + .ok_or(crate::arithmetic::overflow("multiplication")) } /// Convert an intermediate to a coefficient, returning an error on overflow. +#[inline] pub fn checked_i64_to_i32(value: i64) -> Result { - i32::try_from(value).map_err(|_| MathError::overflow("i64 to i32 conversion")) + match i32::try_from(value) { + Ok(value) => Ok(value), + Err(_) => Err(crate::arithmetic::overflow("i64 to i32 conversion")), + } } #[cfg(test)] diff --git a/crates/jxr-math/src/lib.rs b/crates/jxr-math/src/lib.rs index e84e5d9..816be83 100644 --- a/crates/jxr-math/src/lib.rs +++ b/crates/jxr-math/src/lib.rs @@ -7,6 +7,16 @@ #![no_std] #![forbid(unsafe_code)] +// CUDA Oxide cannot lower the nested ControlFlow layout produced by `?`. +macro_rules! math_try { + ($value:expr) => { + match $value { + Ok(value) => value, + Err(error) => return Err(error), + } + }; +} + pub mod alpha; pub mod arithmetic; pub mod color; diff --git a/crates/jxr-math/src/overlap.rs b/crates/jxr-math/src/overlap.rs index 631caab..fcc9860 100644 --- a/crates/jxr-math/src/overlap.rs +++ b/crates/jxr-math/src/overlap.rs @@ -6,49 +6,72 @@ use crate::{ }; /// Apply the exact four-sample JPEG XR inverse overlap postfilter in place. +#[inline] pub fn inverse_overlap_filter_4(values: &mut [i32; 4]) -> Result<(), MathError> { - values[0] = checked_add(values[0], values[3])?; - values[1] = checked_add(values[1], values[2])?; - values[3] = checked_sub(values[3], checked_add(values[0], 1)? >> 1)?; - values[2] = checked_sub(values[2], checked_add(values[1], 1)? >> 1)?; + values[0] = math_try!(checked_add(values[0], values[3])); + values[1] = math_try!(checked_add(values[1], values[2])); + values[3] = math_try!(checked_sub( + values[3], + math_try!(checked_add(values[0], 1)) >> 1 + )); + values[2] = math_try!(checked_sub( + values[2], + math_try!(checked_add(values[1], 1)) >> 1 + )); let mut pair = [values[0], values[3]]; - inverse_scale(&mut pair)?; + math_try!(inverse_scale(&mut pair)); [values[0], values[3]] = pair; pair = [values[1], values[2]]; - inverse_scale(&mut pair)?; + math_try!(inverse_scale(&mut pair)); [values[1], values[2]] = pair; - values[0] = checked_add(values[0], checked_add(checked_mul(values[3], 3)?, 4)? >> 3)?; - values[1] = checked_add(values[1], checked_add(checked_mul(values[2], 3)?, 4)? >> 3)?; - values[3] = checked_sub(values[3], values[0] >> 1)?; - values[2] = checked_sub(values[2], values[1] >> 1)?; - values[0] = checked_add(values[0], values[3])?; - values[1] = checked_add(values[1], values[2])?; - values[3] = checked_neg(values[3])?; - values[2] = checked_neg(values[2])?; + values[0] = math_try!(checked_add( + values[0], + math_try!(checked_add(math_try!(checked_mul(values[3], 3)), 4)) >> 3 + )); + values[1] = math_try!(checked_add( + values[1], + math_try!(checked_add(math_try!(checked_mul(values[2], 3)), 4)) >> 3 + )); + values[3] = math_try!(checked_sub(values[3], values[0] >> 1)); + values[2] = math_try!(checked_sub(values[2], values[1] >> 1)); + values[0] = math_try!(checked_add(values[0], values[3])); + values[1] = math_try!(checked_add(values[1], values[2])); + values[3] = math_try!(checked_neg(values[3])); + values[2] = math_try!(checked_neg(values[2])); pair = [values[2], values[3]]; - inverse_rotate(&mut pair)?; + math_try!(inverse_rotate(&mut pair)); [values[2], values[3]] = pair; - values[3] = checked_add(values[3], checked_add(values[0], 1)? >> 1)?; - values[2] = checked_add(values[2], checked_add(values[1], 1)? >> 1)?; - values[0] = checked_sub(values[0], values[3])?; - values[1] = checked_sub(values[1], values[2])?; + values[3] = math_try!(checked_add( + values[3], + math_try!(checked_add(values[0], 1)) >> 1 + )); + values[2] = math_try!(checked_add( + values[2], + math_try!(checked_add(values[1], 1)) >> 1 + )); + values[0] = math_try!(checked_sub(values[0], values[3])); + values[1] = math_try!(checked_sub(values[1], values[2])); Ok(()) } /// Apply the exact 4-by-4 JPEG XR inverse overlap postfilter in place. +#[inline] pub fn inverse_overlap_filter_4x4(values: &mut [i32; 16]) -> Result<(), MathError> { - for indices in [[0, 3, 12, 15], [1, 2, 13, 14], [4, 7, 8, 11], [5, 6, 9, 10]] { - apply_group(values, indices, |group| inverse_hadamard_2x2(group, 0))?; + let groups = [[0, 3, 12, 15], [1, 2, 13, 14], [4, 7, 8, 11], [5, 6, 9, 10]]; + for &indices in &groups { + math_try!(apply_group(values, indices, inverse_hadamard_zero)); } - for indices in [[13, 12], [9, 8], [7, 3], [6, 2]] { - apply_pair(values, indices, inverse_rotate)?; + let rotations = [[13, 12], [9, 8], [7, 3], [6, 2]]; + for &indices in &rotations { + math_try!(apply_pair(values, indices, inverse_rotate)); } - apply_group(values, [10, 11, 14, 15], inverse_todd_odd_post)?; - for indices in [[0, 15], [1, 14], [4, 11], [5, 10]] { - apply_pair(values, indices, inverse_scale)?; + math_try!(apply_group(values, [10, 11, 14, 15], inverse_todd_odd_post)); + let scales = [[0, 15], [1, 14], [4, 11], [5, 10]]; + for &indices in &scales { + math_try!(apply_pair(values, indices, inverse_scale)); } - for indices in [[0, 3, 12, 15], [1, 2, 13, 14], [4, 7, 8, 11], [5, 6, 9, 10]] { - apply_group(values, indices, inverse_hadamard_post)?; + for &indices in &groups { + math_try!(apply_group(values, indices, inverse_hadamard_post)); } Ok(()) } @@ -56,111 +79,187 @@ pub fn inverse_overlap_filter_4x4(values: &mut [i32; 16]) -> Result<(), MathErro /// Apply the exact 2-by-2 JPEG XR inverse overlap postfilter in place. /// /// This is the first-level interior operator for subsampled chroma. +#[inline] pub fn inverse_overlap_filter_2x2(values: &mut [i32; 4]) -> Result<(), MathError> { - values[0] = checked_add(values[0], values[3])?; - values[1] = checked_add(values[1], values[2])?; - values[3] = checked_sub(values[3], checked_add(values[0], 1)? >> 1)?; - values[2] = checked_sub(values[2], checked_add(values[1], 1)? >> 1)?; - values[1] = checked_add(values[1], checked_add(values[0], 2)? >> 2)?; - values[0] = checked_add(values[0], checked_add(values[1], 1)? >> 1)?; - values[0] = checked_add(values[0], values[1] >> 5)?; - values[0] = checked_add(values[0], values[1] >> 9)?; - values[0] = checked_add(values[0], values[1] >> 13)?; - values[1] = checked_add(values[1], checked_add(values[0], 2)? >> 2)?; - values[3] = checked_add(values[3], checked_add(values[0], 1)? >> 1)?; - values[2] = checked_add(values[2], checked_add(values[1], 1)? >> 1)?; - values[0] = checked_sub(values[0], values[3])?; - values[1] = checked_sub(values[1], values[2])?; + values[0] = math_try!(checked_add(values[0], values[3])); + values[1] = math_try!(checked_add(values[1], values[2])); + values[3] = math_try!(checked_sub( + values[3], + math_try!(checked_add(values[0], 1)) >> 1 + )); + values[2] = math_try!(checked_sub( + values[2], + math_try!(checked_add(values[1], 1)) >> 1 + )); + values[1] = math_try!(checked_add( + values[1], + math_try!(checked_add(values[0], 2)) >> 2 + )); + values[0] = math_try!(checked_add( + values[0], + math_try!(checked_add(values[1], 1)) >> 1 + )); + values[0] = math_try!(checked_add(values[0], values[1] >> 5)); + values[0] = math_try!(checked_add(values[0], values[1] >> 9)); + values[0] = math_try!(checked_add(values[0], values[1] >> 13)); + values[1] = math_try!(checked_add( + values[1], + math_try!(checked_add(values[0], 2)) >> 2 + )); + values[3] = math_try!(checked_add( + values[3], + math_try!(checked_add(values[0], 1)) >> 1 + )); + values[2] = math_try!(checked_add( + values[2], + math_try!(checked_add(values[1], 1)) >> 1 + )); + values[0] = math_try!(checked_sub(values[0], values[3])); + values[1] = math_try!(checked_sub(values[1], values[2])); Ok(()) } /// Apply the exact two-sample JPEG XR inverse overlap postfilter in place. /// /// This is the first-level boundary operator for subsampled chroma. +#[inline] pub fn inverse_overlap_filter_2(values: &mut [i32; 2]) -> Result<(), MathError> { - values[1] = checked_add(values[1], checked_add(values[0], 2)? >> 2)?; - values[0] = checked_add(values[0], checked_add(values[1], 1)? >> 1)?; - values[0] = checked_add(values[0], values[1] >> 5)?; - values[0] = checked_add(values[0], values[1] >> 9)?; - values[0] = checked_add(values[0], values[1] >> 13)?; - values[1] = checked_add(values[1], checked_add(values[0], 2)? >> 2)?; + values[1] = math_try!(checked_add( + values[1], + math_try!(checked_add(values[0], 2)) >> 2 + )); + values[0] = math_try!(checked_add( + values[0], + math_try!(checked_add(values[1], 1)) >> 1 + )); + values[0] = math_try!(checked_add(values[0], values[1] >> 5)); + values[0] = math_try!(checked_add(values[0], values[1] >> 9)); + values[0] = math_try!(checked_add(values[0], values[1] >> 13)); + values[1] = math_try!(checked_add( + values[1], + math_try!(checked_add(values[0], 2)) >> 2 + )); Ok(()) } +#[inline] fn inverse_rotate(values: &mut [i32; 2]) -> Result<(), MathError> { - values[0] = checked_sub(values[0], checked_add(values[1], 1)? >> 1)?; - values[1] = checked_add(values[1], checked_add(values[0], 1)? >> 1)?; + values[0] = math_try!(checked_sub( + values[0], + math_try!(checked_add(values[1], 1)) >> 1 + )); + values[1] = math_try!(checked_add( + values[1], + math_try!(checked_add(values[0], 1)) >> 1 + )); Ok(()) } +#[inline] fn inverse_scale(values: &mut [i32; 2]) -> Result<(), MathError> { - values[0] = checked_add(values[0], values[1])?; - values[1] = checked_sub(values[0] >> 1, values[1])?; - values[0] = checked_add(values[0], checked_mul(values[1], 3)? >> 3)?; - values[1] = checked_add(values[1], checked_mul(values[0], 3)? >> 4)?; - values[1] = checked_add(values[1], values[0] >> 7)?; - values[1] = checked_sub(values[1], values[0] >> 10)?; + values[0] = math_try!(checked_add(values[0], values[1])); + values[1] = math_try!(checked_sub(values[0] >> 1, values[1])); + values[0] = math_try!(checked_add( + values[0], + math_try!(checked_mul(values[1], 3)) >> 3 + )); + values[1] = math_try!(checked_add( + values[1], + math_try!(checked_mul(values[0], 3)) >> 4 + )); + values[1] = math_try!(checked_add(values[1], values[0] >> 7)); + values[1] = math_try!(checked_sub(values[1], values[0] >> 10)); Ok(()) } +#[inline] fn inverse_hadamard_post(values: &mut [i32; 4]) -> Result<(), MathError> { - values[1] = checked_sub(values[1], values[2])?; - values[0] = checked_add(values[0], checked_add(checked_mul(values[3], 3)?, 4)? >> 3)?; - values[3] = checked_sub(values[3], values[1] >> 1)?; - values[2] = checked_sub(checked_sub(values[0], values[1])? >> 1, values[2])?; - values.swap(2, 3); - values[0] = checked_sub(values[0], values[3])?; - values[1] = checked_add(values[1], values[2])?; + values[1] = math_try!(checked_sub(values[1], values[2])); + values[0] = math_try!(checked_add( + values[0], + math_try!(checked_add(math_try!(checked_mul(values[3], 3)), 4)) >> 3 + )); + values[3] = math_try!(checked_sub(values[3], values[1] >> 1)); + values[2] = math_try!(checked_sub( + math_try!(checked_sub(values[0], values[1])) >> 1, + values[2] + )); + [values[2], values[3]] = [values[3], values[2]]; + values[0] = math_try!(checked_sub(values[0], values[3])); + values[1] = math_try!(checked_add(values[1], values[2])); Ok(()) } +#[inline] fn inverse_todd_odd_post(values: &mut [i32; 4]) -> Result<(), MathError> { - values[3] = checked_add(values[3], values[0])?; - values[2] = checked_sub(values[2], values[1])?; + values[3] = math_try!(checked_add(values[3], values[0])); + values[2] = math_try!(checked_sub(values[2], values[1])); let first = values[3] >> 1; let second = values[2] >> 1; - values[0] = checked_sub(values[0], first)?; - values[1] = checked_add(values[1], second)?; - values[0] = checked_sub(values[0], checked_add(checked_mul(values[1], 3)?, 6)? >> 3)?; - values[1] = checked_add(values[1], checked_add(checked_mul(values[0], 3)?, 2)? >> 2)?; - values[0] = checked_sub(values[0], checked_add(checked_mul(values[1], 3)?, 4)? >> 3)?; - values[1] = checked_sub(values[1], second)?; - values[0] = checked_add(values[0], first)?; - values[2] = checked_add(values[2], values[1])?; - values[3] = checked_sub(values[3], values[0])?; + values[0] = math_try!(checked_sub(values[0], first)); + values[1] = math_try!(checked_add(values[1], second)); + values[0] = math_try!(checked_sub( + values[0], + math_try!(checked_add(math_try!(checked_mul(values[1], 3)), 6)) >> 3 + )); + values[1] = math_try!(checked_add( + values[1], + math_try!(checked_add(math_try!(checked_mul(values[0], 3)), 2)) >> 2 + )); + values[0] = math_try!(checked_sub( + values[0], + math_try!(checked_add(math_try!(checked_mul(values[1], 3)), 4)) >> 3 + )); + values[1] = math_try!(checked_sub(values[1], second)); + values[0] = math_try!(checked_add(values[0], first)); + values[2] = math_try!(checked_add(values[2], values[1])); + values[3] = math_try!(checked_sub(values[3], values[0])); Ok(()) } +#[inline] +fn inverse_hadamard_zero(values: &mut [i32; 4]) -> Result<(), MathError> { + inverse_hadamard_2x2(values, 0) +} + +#[inline] fn apply_group( values: &mut [i32; 16], indices: [usize; 4], operation: impl FnOnce(&mut [i32; 4]) -> Result<(), MathError>, ) -> Result<(), MathError> { - let mut group = indices.map(|index| values[index]); - operation(&mut group)?; - for (index, value) in indices.into_iter().zip(group) { - values[index] = value; + let mut group = [ + values[indices[0]], + values[indices[1]], + values[indices[2]], + values[indices[3]], + ]; + math_try!(operation(&mut group)); + for i in 0..4 { + values[indices[i]] = group[i]; } Ok(()) } +#[inline] fn apply_pair( values: &mut [i32; 16], indices: [usize; 2], operation: impl FnOnce(&mut [i32; 2]) -> Result<(), MathError>, ) -> Result<(), MathError> { - let mut pair = indices.map(|index| values[index]); - operation(&mut pair)?; - for (index, value) in indices.into_iter().zip(pair) { - values[index] = value; + let mut pair = [values[indices[0]], values[indices[1]]]; + math_try!(operation(&mut pair)); + for i in 0..2 { + values[indices[i]] = pair[i]; } Ok(()) } +#[inline] fn checked_neg(value: i32) -> Result { value .checked_neg() - .ok_or_else(|| MathError::overflow("overlap negation")) + .ok_or(crate::arithmetic::overflow("overlap negation")) } #[cfg(test)] diff --git a/crates/jxr-math/src/transform.rs b/crates/jxr-math/src/transform.rs index fa5b489..7bdba08 100644 --- a/crates/jxr-math/src/transform.rs +++ b/crates/jxr-math/src/transform.rs @@ -4,34 +4,42 @@ use crate::arithmetic::{MathError, checked_add, checked_mul, checked_sub}; use crate::tables::INVERSE_PERMUTATION; /// Apply the reversible JPEG XR 4-by-4 inverse core transform in place. +#[inline] pub fn inverse_core_transform(coefficients: &mut [i32; 16]) -> Result<(), MathError> { inverse_permute(coefficients); - transform_group(coefficients, [0, 1, 4, 5], |values| t2x2h(values, 1))?; - transform_group(coefficients, [2, 3, 6, 7], inverse_todd)?; - transform_group(coefficients, [8, 12, 9, 13], inverse_todd)?; - transform_group(coefficients, [10, 11, 14, 15], inverse_todd_odd)?; - transform_group(coefficients, [0, 3, 12, 15], |values| t2x2h(values, 0))?; - transform_group(coefficients, [5, 6, 9, 10], |values| t2x2h(values, 0))?; - transform_group(coefficients, [1, 2, 13, 14], |values| t2x2h(values, 0))?; - transform_group(coefficients, [4, 7, 8, 11], |values| t2x2h(values, 0)) + math_try!(transform_group(coefficients, [0, 1, 4, 5], t2x2h_round)); + math_try!(transform_group(coefficients, [2, 3, 6, 7], inverse_todd)); + math_try!(transform_group(coefficients, [8, 12, 9, 13], inverse_todd)); + math_try!(transform_group( + coefficients, + [10, 11, 14, 15], + inverse_todd_odd + )); + math_try!(transform_group(coefficients, [0, 3, 12, 15], t2x2h_zero)); + math_try!(transform_group(coefficients, [5, 6, 9, 10], t2x2h_zero)); + math_try!(transform_group(coefficients, [1, 2, 13, 14], t2x2h_zero)); + transform_group(coefficients, [4, 7, 8, 11], t2x2h_zero) } /// Apply the four-value inverse Hadamard lifting operator used by transforms and postfilters. +#[inline] pub fn inverse_hadamard_2x2(values: &mut [i32; 4], rounding: i32) -> Result<(), MathError> { t2x2h(values, rounding) } /// Apply the first-level inverse transform for one YUV 4:2:0 chroma macroblock. +#[inline] pub fn inverse_chroma_420(coefficients: &mut [i32; 4]) -> Result<(), MathError> { - t2x2h(coefficients, 0)?; - coefficients.swap(1, 2); + math_try!(t2x2h(coefficients, 0)); + [coefficients[1], coefficients[2]] = [coefficients[2], coefficients[1]]; Ok(()) } /// Apply the first-level inverse transform for one YUV 4:2:2 chroma macroblock. +#[inline] pub fn inverse_chroma_422(coefficients: &mut [i32; 8]) -> Result<(), MathError> { let mut pair = [coefficients[0], coefficients[4]]; - inverse_t2pt(&mut pair)?; + math_try!(inverse_t2pt(&mut pair)); coefficients[0] = pair[0]; coefficients[4] = pair[1]; @@ -41,9 +49,9 @@ pub fn inverse_chroma_422(coefficients: &mut [i32; 8]) -> Result<(), MathError> coefficients[2], coefficients[3], ]; - t2x2h(&mut first, 0)?; + math_try!(t2x2h(&mut first, 0)); coefficients[..4].copy_from_slice(&first); - coefficients.swap(1, 2); + [coefficients[1], coefficients[2]] = [coefficients[2], coefficients[1]]; let mut second = [ coefficients[4], @@ -51,90 +59,163 @@ pub fn inverse_chroma_422(coefficients: &mut [i32; 8]) -> Result<(), MathError> coefficients[5], coefficients[7], ]; - t2x2h(&mut second, 0)?; + math_try!(t2x2h(&mut second, 0)); coefficients[4] = second[0]; coefficients[6] = second[1]; coefficients[5] = second[2]; coefficients[7] = second[3]; - coefficients.swap(5, 6); + [coefficients[5], coefficients[6]] = [coefficients[6], coefficients[5]]; Ok(()) } +#[inline] fn inverse_t2pt(values: &mut [i32; 2]) -> Result<(), MathError> { - values[0] = checked_sub(values[0], checked_add(values[1], 1)? >> 1)?; - values[1] = checked_add(values[1], values[0])?; + values[0] = math_try!(checked_sub( + values[0], + math_try!(checked_add(values[1], 1)) >> 1 + )); + values[1] = math_try!(checked_add(values[1], values[0])); Ok(()) } +#[inline] fn transform_group( coefficients: &mut [i32; 16], indices: [usize; 4], operation: impl FnOnce(&mut [i32; 4]) -> Result<(), MathError>, ) -> Result<(), MathError> { - let mut values = indices.map(|index| coefficients[index]); - operation(&mut values)?; - for (index, value) in indices.into_iter().zip(values) { - coefficients[index] = value; + let mut values = [ + coefficients[indices[0]], + coefficients[indices[1]], + coefficients[indices[2]], + coefficients[indices[3]], + ]; + math_try!(operation(&mut values)); + for i in 0..4 { + coefficients[indices[i]] = values[i]; } Ok(()) } +#[inline] fn inverse_permute(coefficients: &mut [i32; 16]) { let input = *coefficients; - for (source, destination) in INVERSE_PERMUTATION.into_iter().enumerate() { - coefficients[destination] = input[source]; - } + // Constant subscripts let CUDA Oxide keep the permutation and temporary + // blocks in registers instead of materializing per-thread stack arrays. + coefficients[INVERSE_PERMUTATION[0]] = input[0]; + coefficients[INVERSE_PERMUTATION[1]] = input[1]; + coefficients[INVERSE_PERMUTATION[2]] = input[2]; + coefficients[INVERSE_PERMUTATION[3]] = input[3]; + coefficients[INVERSE_PERMUTATION[4]] = input[4]; + coefficients[INVERSE_PERMUTATION[5]] = input[5]; + coefficients[INVERSE_PERMUTATION[6]] = input[6]; + coefficients[INVERSE_PERMUTATION[7]] = input[7]; + coefficients[INVERSE_PERMUTATION[8]] = input[8]; + coefficients[INVERSE_PERMUTATION[9]] = input[9]; + coefficients[INVERSE_PERMUTATION[10]] = input[10]; + coefficients[INVERSE_PERMUTATION[11]] = input[11]; + coefficients[INVERSE_PERMUTATION[12]] = input[12]; + coefficients[INVERSE_PERMUTATION[13]] = input[13]; + coefficients[INVERSE_PERMUTATION[14]] = input[14]; + coefficients[INVERSE_PERMUTATION[15]] = input[15]; +} + +#[inline] +fn t2x2h_round(values: &mut [i32; 4]) -> Result<(), MathError> { + t2x2h(values, 1) +} + +#[inline] +fn t2x2h_zero(values: &mut [i32; 4]) -> Result<(), MathError> { + t2x2h(values, 0) } +#[inline] fn t2x2h(values: &mut [i32; 4], rounding: i32) -> Result<(), MathError> { - values[0] = checked_add(values[0], values[3])?; - values[1] = checked_sub(values[1], values[2])?; - let difference = checked_sub(values[0], values[1])?; - let first = checked_add(difference, rounding)? >> 1; + values[0] = math_try!(checked_add(values[0], values[3])); + values[1] = math_try!(checked_sub(values[1], values[2])); + let difference = math_try!(checked_sub(values[0], values[1])); + let first = math_try!(checked_add(difference, rounding)) >> 1; let second = values[2]; - values[2] = checked_sub(first, values[3])?; - values[3] = checked_sub(first, second)?; - values[0] = checked_sub(values[0], values[3])?; - values[1] = checked_add(values[1], values[2])?; + values[2] = math_try!(checked_sub(first, values[3])); + values[3] = math_try!(checked_sub(first, second)); + values[0] = math_try!(checked_sub(values[0], values[3])); + values[1] = math_try!(checked_add(values[1], values[2])); Ok(()) } +#[inline] fn inverse_todd(values: &mut [i32; 4]) -> Result<(), MathError> { - values[1] = checked_add(values[1], values[3])?; - values[0] = checked_sub(values[0], values[2])?; - values[3] = checked_sub(values[3], values[1] >> 1)?; - values[2] = checked_add(values[2], checked_add(values[0], 1)? >> 1)?; - values[0] = checked_sub(values[0], checked_add(checked_mul(values[1], 3)?, 4)? >> 3)?; - values[1] = checked_add(values[1], checked_add(checked_mul(values[0], 3)?, 4)? >> 3)?; - values[2] = checked_sub(values[2], checked_add(checked_mul(values[3], 3)?, 4)? >> 3)?; - values[3] = checked_add(values[3], checked_add(checked_mul(values[2], 3)?, 4)? >> 3)?; - values[2] = checked_sub(values[2], checked_add(values[1], 1)? >> 1)?; - values[3] = checked_sub(checked_add(values[0], 1)? >> 1, values[3])?; - values[1] = checked_add(values[1], values[2])?; - values[0] = checked_sub(values[0], values[3])?; + values[1] = math_try!(checked_add(values[1], values[3])); + values[0] = math_try!(checked_sub(values[0], values[2])); + values[3] = math_try!(checked_sub(values[3], values[1] >> 1)); + values[2] = math_try!(checked_add( + values[2], + math_try!(checked_add(values[0], 1)) >> 1 + )); + values[0] = math_try!(checked_sub( + values[0], + math_try!(checked_add(math_try!(checked_mul(values[1], 3)), 4)) >> 3 + )); + values[1] = math_try!(checked_add( + values[1], + math_try!(checked_add(math_try!(checked_mul(values[0], 3)), 4)) >> 3 + )); + values[2] = math_try!(checked_sub( + values[2], + math_try!(checked_add(math_try!(checked_mul(values[3], 3)), 4)) >> 3 + )); + values[3] = math_try!(checked_add( + values[3], + math_try!(checked_add(math_try!(checked_mul(values[2], 3)), 4)) >> 3 + )); + values[2] = math_try!(checked_sub( + values[2], + math_try!(checked_add(values[1], 1)) >> 1 + )); + values[3] = math_try!(checked_sub( + math_try!(checked_add(values[0], 1)) >> 1, + values[3] + )); + values[1] = math_try!(checked_add(values[1], values[2])); + values[0] = math_try!(checked_sub(values[0], values[3])); Ok(()) } +#[inline] fn inverse_todd_odd(values: &mut [i32; 4]) -> Result<(), MathError> { - values[3] = checked_add(values[3], values[0])?; - values[2] = checked_sub(values[2], values[1])?; + values[3] = math_try!(checked_add(values[3], values[0])); + values[2] = math_try!(checked_sub(values[2], values[1])); let first = values[3] >> 1; let second = values[2] >> 1; - values[0] = checked_sub(values[0], first)?; - values[1] = checked_add(values[1], second)?; - values[0] = checked_sub(values[0], checked_add(checked_mul(values[1], 3)?, 3)? >> 3)?; - values[1] = checked_add(values[1], checked_add(checked_mul(values[0], 3)?, 3)? >> 2)?; - values[0] = checked_sub(values[0], checked_add(checked_mul(values[1], 3)?, 4)? >> 3)?; - values[1] = checked_sub(values[1], second)?; - values[0] = checked_add(values[0], first)?; - values[2] = checked_add(values[2], values[1])?; - values[3] = checked_sub(values[3], values[0])?; - values[1] = values[1] - .checked_neg() - .ok_or_else(|| MathError::overflow("inverse Todd odd negation"))?; - values[2] = values[2] - .checked_neg() - .ok_or_else(|| MathError::overflow("inverse Todd odd negation"))?; + values[0] = math_try!(checked_sub(values[0], first)); + values[1] = math_try!(checked_add(values[1], second)); + values[0] = math_try!(checked_sub( + values[0], + math_try!(checked_add(math_try!(checked_mul(values[1], 3)), 3)) >> 3 + )); + values[1] = math_try!(checked_add( + values[1], + math_try!(checked_add(math_try!(checked_mul(values[0], 3)), 3)) >> 2 + )); + values[0] = math_try!(checked_sub( + values[0], + math_try!(checked_add(math_try!(checked_mul(values[1], 3)), 4)) >> 3 + )); + values[1] = math_try!(checked_sub(values[1], second)); + values[0] = math_try!(checked_add(values[0], first)); + values[2] = math_try!(checked_add(values[2], values[1])); + values[3] = math_try!(checked_sub(values[3], values[0])); + values[1] = math_try!( + values[1] + .checked_neg() + .ok_or(crate::arithmetic::overflow("inverse Todd odd negation")) + ); + values[2] = math_try!( + values[2] + .checked_neg() + .ok_or(crate::arithmetic::overflow("inverse Todd odd negation")) + ); Ok(()) } diff --git a/crates/jxr-math/tests/math_contract.rs b/crates/jxr-math/tests/math_contract.rs index a3691ca..62b02a4 100644 --- a/crates/jxr-math/tests/math_contract.rs +++ b/crates/jxr-math/tests/math_contract.rs @@ -9,9 +9,17 @@ fn quantizer_rejects_zero_and_dequantize_overflow() { } #[test] -fn cuda_reconstruction_abi_is_generated_from_the_canonical_schema() { - let source = jxr_math::tables::CUDA_RECONSTRUCTION_CONSTANTS; - assert!(source.contains("JXR_INVERSE_PERMUTATION[16]")); - assert!(source.contains("struct JxrMacroblockAbi")); - assert!(source.contains("struct JxrOutputAbi")); +fn reconstruction_descriptors_match_the_canonical_word_layout() { + use core::mem::{offset_of, size_of}; + use jxr_math::tables::{self, reconstruction_abi::*}; + assert_eq!( + size_of::(), + tables::ABI_JXRMACROBLOCKABI_SIZE + ); + assert_eq!(size_of::(), tables::ABI_JXRPLANEABI_SIZE); + assert_eq!(size_of::(), tables::ABI_JXROUTPUTABI_SIZE); + assert_eq!( + offset_of!(JxrOutputAbi, reserved0), + tables::ABI_JXROUTPUTABI_RESERVED0_OFFSET + ); } diff --git a/crates/jxr/tests/cuda_contracts.rs b/crates/jxr/tests/cuda_contracts.rs index f015157..9d5da54 100644 --- a/crates/jxr/tests/cuda_contracts.rs +++ b/crates/jxr/tests/cuda_contracts.rs @@ -71,7 +71,7 @@ fn auto_without_an_attached_cuda_session_remains_on_cpu() { } #[test] -#[ignore = "requires a compatible NVIDIA GPU, CUDA driver, and NVRTC"] +#[ignore = "requires a compatible NVIDIA GPU, CUDA driver, and compiled CUDA Oxide kernels"] fn cuda_submission_resident_readback_and_host_decode_match_cpu() { let bytes = decodable_minimal_raw_codestream(); let view = JxrView::parse(&bytes).unwrap(); @@ -131,7 +131,7 @@ fn cuda_submission_resident_readback_and_host_decode_match_cpu() { } #[test] -#[ignore = "requires a compatible NVIDIA GPU, CUDA driver, and NVRTC"] +#[ignore = "requires a compatible NVIDIA GPU, CUDA driver, and compiled CUDA Oxide kernels"] fn cuda_dense_batch_and_checked_destination_match_individual_output() { let bytes = decodable_minimal_raw_codestream(); let view = JxrView::parse(&bytes).unwrap(); @@ -162,7 +162,7 @@ fn cuda_dense_batch_and_checked_destination_match_individual_output() { } #[test] -#[ignore = "requires a compatible NVIDIA GPU, CUDA driver, and NVRTC"] +#[ignore = "requires a compatible NVIDIA GPU, CUDA driver, and compiled CUDA Oxide kernels"] fn high_level_cuda_batch_preserves_native_groups_and_resident_outputs() { let source: Arc<[u8]> = decodable_minimal_raw_codestream().into(); let session = jxr::cuda::CudaDecoderSession::system_default().unwrap(); diff --git a/docs/architecture.md b/docs/architecture.md index 6f610a4..9496846 100644 --- a/docs/architecture.md +++ b/docs/architecture.md @@ -87,7 +87,7 @@ fn main() -> Result<(), Box> { For CUDA, replace the session with `jxr::cuda::CudaDecoderSession::new(0)?` (device 0), attach it with `with_cuda_session(&session)`, and request -`BackendRequest::Cuda`. The CUDA Driver API and NVRTC bindings load at runtime. +`BackendRequest::Cuda`. The CUDA driver loads embedded CUDA Oxide PTX at runtime. Enabling a feature alone does not attach a session. Keep the session alive for as long as the decoder borrows it. @@ -146,10 +146,10 @@ container writer does not accept decoded pixels or synthesize entropy syntax. `jxr-math` is `no_std` and owns checked integer transforms, overlap operators, dequantization, HP prediction, interpolation, color, alpha, and RGBE arithmetic. -Metal uses backend-local MSL and CUDA uses backend-local CUDA C compiled through -NVRTC; their current -HP phase consumes compact CPU-selected directions and resolves prediction with -race-free per-block prefix sums before dequantization. Transform constants and +Metal uses backend-local MSL. CUDA uses Rust kernels compiled by CUDA Oxide, +calling the shared Rust transforms and overlap filters directly. The HP phase +consumes compact CPU-selected directions and resolves prediction within each +macroblock before dequantization. Transform constants and parameter indices are generated from `data/reconstruction.abi`, which also generates the Rust bindings used by both device encoders. All accelerator manifests describe the complete packaged Main-profile pipeline. diff --git a/docs/cuda-backend.md b/docs/cuda-backend.md index f688464..365f6a9 100644 --- a/docs/cuda-backend.md +++ b/docs/cuda-backend.md @@ -26,7 +26,7 @@ workspace: device allocations rather than exposing Metal types or duplicating its platform adapter. - `jxr-math/data/reconstruction.abi` is the canonical transform permutation and - descriptor schema. Its build script now emits both MSL and CUDA C declarations; + descriptor schema. Its build script emits MSL and CUDA Oxide Rust descriptors; Rust layout assertions keep the host descriptors tied to the same source. - Existing tests place scalar math in `jxr-math`, syntax/reconstruction/output behavior in `jxr-native`, public routing and batches in `jxr`, and official @@ -36,8 +36,7 @@ workspace: - Regular CI already exercises formatting, workspace Clippy/tests, minimal features, strict rustdoc, and fuzz-package compilation on Linux and macOS. Metal hardware validation is a serialized self-hosted workflow. CUDA follows - that split: dynamic-loading builds stay in regular CI, while compilation by - NVRTC, device execution, the 517-case differential suite, and benchmarks are + that split: dynamic-loading builds stay in regular CI, while CUDA Oxide compilation, device execution, the 517-case differential suite, and benchmarks are confined to the serialized NVIDIA runner. - The existing pathology benchmarks separate preparation, submission, synchronization, and readback while checking output checksums. The CUDA @@ -48,64 +47,50 @@ The Annex-A writer remains outside every accelerator path. It still accepts an already encoded codestream and packages metadata and optional encoded alpha; no pixel-to-codestream encoder is introduced. -## Dependency decision (September 2026) - -`jxr-cuda` pins `cudarc` 0.19.10 with only `std`, `driver`, `nvrtc`, -`dynamic-loading`, and the CUDA 11.4 ABI baseline enabled. The decision is based -on the following primary project and vendor documentation: - -- [`cudarc` 0.19.10 API documentation](https://docs.rs/cudarc/0.19.10/cudarc/) - documents safe context, stream, event, slice, Driver API, and NVRTC wrappers; - its dynamic-loading mode requires no CUDA libraries at build time. The project - is MIT OR Apache-2.0 and its - [release history](https://github.com/chelsea0x3b/cudarc/releases) shows current - CUDA 13.x maintenance. -- [`cust` 0.3.2 documentation](https://docs.rs/cust/0.3.2/cust/) describes a - capable Driver API wrapper, but requires CUDA development libraries on the - build system. The upstream - [Rust-CUDA README](https://github.com/Rust-GPU/rust-cuda/blob/main/README.md) - also characterizes the rebooted project as early development with expected - bugs and safety issues. Its Rust-device toolchain would add nightly/LLVM and - build-time CUDA requirements that conflict with this workspace's portable - all-feature checks. -- The original - [RustaCUDA repository](https://github.com/bheisler/RustaCUDA) is no longer the - maintained direction; its ecosystem moved toward `cust`/Rust-CUDA. -- NVIDIA documents NVRTC as the supported in-process CUDA C++ to PTX boundary - and states that its PTX can be loaded through the Driver API - ([NVRTC guide](https://docs.nvidia.com/cuda/nvrtc/)). NVIDIA also defines - streams and events as the core asynchronous execution primitives - ([CUDA Programming Guide](https://docs.nvidia.com/cuda/cuda-programming-guide/02-basics/asynchronous-execution.html)) - and documents stream-ordered allocation and memory-pool capability discovery - in the - [Driver API](https://docs.nvidia.com/cuda/cuda-driver-api/group__CUDA__MALLOC__ASYNC.html). - -The selected dependency adds no codec, image, neural-network, BLAS, or CUDA -Runtime API layer. Disabling default `cudarc` features avoids cuBLAS, cuDNN, -cuRAND, and their transitive/runtime surfaces. The lockfile adds `cudarc`, -ISC-licensed `libloading`, `cfg-if`, and the Windows loader shim; every license -is compatible with distribution under this workspace's terms. This keeps -compile time and binary surface substantially smaller than a general GPU -framework. - -Maintenance risk remains concentrated in one FFI wrapper and NVIDIA's Driver -API/NVRTC compatibility. The version is pinned so an upstream ABI or safety -change cannot silently alter builds. The dependency's dual license matches the -workspace. Runtime loader failures and missing symbols become explicit session -initialization errors; CPU-only builds and ordinary `jxr` builds do not link, -load, or probe CUDA. - -Enabling the feature embeds CUDA C source and compiles the small Rust/FFI -adapter, but runs neither `nvcc` nor NVRTC at build time. Creating a session -does incur one NVRTC compile and one Driver API module load for that device; -cloning and reusing the session amortizes that startup cost. Linux and Windows -use the NVIDIA libraries found by the platform loader, while macOS can compile -the Rust feature but reports CUDA unavailable at runtime. Dynamic loading keeps -the toolkit out of CPU deployments, but inherits the operating system's library -search-path trust boundary: production processes must not search user-writable -directories for `cuda`/`nvcuda` or `nvrtc`. Unsafe FFI representation is kept in -the focused backend and `cudarc`; checked Rust constructors validate buffer -extents, contexts, ABI widths, and kernel indexing metadata before submission. +## CUDA Oxide build and runtime + +All reconstruction kernels are Rust compiled to PTX by +[NVIDIA CUDA Oxide](https://github.com/NVlabs/cuda-oxide). The device project +pins the same CUDA Oxide revision and `nightly-2026-04-03` toolchain used by +this organization's `j2k` backend. The main workspace retains its existing +Rust toolchain. `jxr-math` supports the device compiler's Rust 1.96 baseline. +Transforms and overlap filters call the allocation-free +`jxr-math` implementation directly; device descriptors are generated from +`jxr-math/data/reconstruction.abi`. The isolated device build uses numeric +arithmetic errors, explicit error propagation, and indexed array operations to +accommodate current CUDA Oxide code-generation limits. Host errors retain their +operation descriptions. + +On Linux, install the pinned CUDA Oxide compiler and its documented build +prerequisites before building an executable CUDA backend. This includes a +working libclang and its builtin headers; custom installations may require +`LIBCLANG_PATH`, `LD_LIBRARY_PATH`, and `BINDGEN_EXTRA_CLANG_ARGS`. +The build script +stages its standalone device project under `OUT_DIR`, invokes +`cargo oxide build --arch sm_50 --no-fmad`, and embeds the resulting PTX. +`JXR_CUDA_OXIDE_ARCH` selects another target architecture. Compiler failures +are build errors; no CUDA C implementation is retained. + +```sh +JXR_REQUIRE_CUDA_OXIDE=1 cargo test -p jxr --features cuda --test cuda_contracts -- --ignored --test-threads=1 +JXR_REQUIRE_CUDA_OXIDE=1 cargo test -p jxr-test-support --features cuda --test cuda_hardware -- --ignored --test-threads=1 +``` + +For builds without the compiler, including cross-platform CI, the Rust +adapter and its device-free contracts still compile, but +`CudaDecoderSession::is_available()` returns false. Set +`JXR_REQUIRE_CUDA_OXIDE=1` in hardware CI to make missing compiler support an +error. Alternatively, `JXR_CUDA_OXIDE_PTX=/absolute/path/jxr_reconstruction.ptx` +embeds a previously compiled module, without invoking the device compiler. +The PTX target must be compatible with the deployment GPU. + +`cudarc` 0.19.10 supplies dynamically loaded Driver API contexts, streams, +events, allocations, and launches. Its `nvrtc` Cargo feature remains enabled +because this version gates its `Ptx` container and `load_module` method behind +that feature. JXR never invokes NVRTC or loads its library. Session creation +loads the embedded Rust-generated PTX through the CUDA driver; it performs no +source compilation. Deployment requires the NVIDIA driver, not NVRTC or the +CUDA Oxide compiler. The default target requires compute capability 5.0. ## Architecture and semantics @@ -125,7 +110,7 @@ ABI offset before allocation or submission. Five ordered phases implement: 5. chroma reconstruction, color conversion, alpha handling, crop, clipping, and typed or packed output storage. -CUDA C kernels and Rust descriptors are generated from the same canonical ABI +CUDA Oxide kernels and host descriptors follow the same canonical ABI schema used by Metal. Checked 64-bit intermediates set a device status word on any operation that would exceed the scalar `i32` contract. That status is read only after the completion event; it is never converted into a CPU retry. @@ -158,7 +143,7 @@ submission is likewise returned to the caller and is never retried on CPU. The regular CI matrix compiles and lints the CUDA feature on CPU-only Linux and macOS hosts through dynamic loading. `.github/workflows/cuda-hardware.yml` is a serialized, manually dispatched self-hosted NVIDIA workflow. It records the GPU -and toolkit, runs ignored lifecycle/destination tests, differentially compares +and CUDA Oxide version, runs ignored lifecycle/destination tests, differentially compares all 517 in-scope T.834 cases and ROI boundary variants with CPU, compares the complete suite with T.835, and retains reports.