diff --git a/.github/workflows/rust-host.yml b/.github/workflows/rust-host.yml new file mode 100644 index 000000000000..ad93cbf4b65b --- /dev/null +++ b/.github/workflows/rust-host.yml @@ -0,0 +1,43 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +name: Rust host checks + +on: + push: + branches: + - main + - "pull-request/[0-9]+" + paths: + - "rust/**" + - ".github/workflows/rust-host.yml" + workflow_dispatch: + +permissions: + contents: read + +concurrency: + group: ${{ github.workflow }}-${{ github.ref }} + cancel-in-progress: true + +jobs: + host: + runs-on: ubuntu-latest + defaults: + run: + working-directory: rust + shell: bash --noprofile --norc -euo pipefail {0} + env: + RUSTDOCFLAGS: -D warnings + steps: + - uses: actions/checkout@v6 + with: + persist-credentials: false + - name: Check the pinned Rust workspace + run: | + rustup toolchain install --no-self-update + rustup show active-toolchain + cargo fmt --all -- --check + cargo test --workspace --locked + cargo clippy --workspace --all-targets --locked -- -D warnings + cargo doc --workspace --no-deps --locked diff --git a/docs/index.rst b/docs/index.rst index e1edae43a940..8906c9170e1c 100644 --- a/docs/index.rst +++ b/docs/index.rst @@ -7,6 +7,7 @@ CUDA Core Compute Libraries cccl/index python/index + rust/index contributors/index maintainers/index infrastructure/index @@ -18,6 +19,8 @@ make CUDA C++ and Python more delightful. - :doc:`Python Libraries ` +- :doc:`Rust Development ` + - :doc:`Contributor Docs ` - :doc:`Maintainer Docs ` diff --git a/docs/rust/index.rst b/docs/rust/index.rst new file mode 100644 index 000000000000..2846d81e9413 --- /dev/null +++ b/docs/rust/index.rst @@ -0,0 +1,119 @@ +.. Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. +.. SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +Rust development +================ + +The Rust workspace contains the native host compiler foundation for ``cuda.coop``. +Its private ``cuda-coop-codegen`` crate specializes CUB block operations and returns +linkable device LTO-IR with exact temporary-storage requirements. It supports +compiler integration development; an application-facing Rust device API is still +under development. The crate is unpublished and its interface can change. + +The workspace is separate from CCCL's CMake build. Building the existing C++ +libraries does not require Rust, and the provider does not depend on Python. + +Supported provider forms +------------------------ + +Each request describes a complete, one-dimensional block, an ``i32`` payload, +a positive compile-time item count per thread, and one of these forms: + +.. list-table:: + :header-rows: 1 + :widths: 25 45 30 + + * - Operation + - Behavior + - Payload requirement + * - Full Load + - Fills every payload slot in blocked order + - Writable storage + * - Default-filled Load + - Loads a valid prefix and fills the suffix with a supplied value + - Writable storage + * - Full Store + - Writes every payload slot to memory + - Every slot initialized + * - Partial Store + - Writes a valid prefix; leaves the destination suffix unchanged + - Every slot initialized + +Both Direct and Transpose algorithms are supported. Direct accesses no shared +scratch. Transpose Store redistributes blocked payloads across lanes into striped +order, retaining initialized values in every slot, and uses the exact CUB storage +layout. Compiler integrations must enforce collective participation, memory +ownership, valid bounds, and synchronization before scratch reuse. Full and +masked Load produce fully initialized payloads; no-default masked Load is outside +this provider's current scope. + +Compile a provider +------------------ + +Use the pinned toolchain in ``rust/rust-toolchain.toml``. Host tests, lint checks, +and rustdoc builds need no CUDA installation. Running native compilation requires +a Linux CUDA toolkit with NVRTC LTO support; CUDA 13.3 is the initial validation +toolkit. A GPU and CUDA context are not required. + +From the CCCL checkout: + +.. code-block:: shell + + cd rust + cargo test --workspace + cargo clippy --workspace --all-targets -- -D warnings + RUSTDOCFLAGS="-D warnings" cargo doc --workspace --no-deps + cargo run -p cuda-coop-codegen --example specialize -- \ + /usr/local/cuda-13.3 compute_120 /tmp/block_load.ltoir + +The example below is the executed Cargo source. It obtains its CCCL source root +from its build location, passes an explicit toolkit root, compiles a filled Load, +and writes the resulting LTO-IR. It does not link or launch a device kernel. + +.. literalinclude:: ../../rust/cuda-coop-codegen/examples/specialize.rs + :language: rust + :start-after: //! Compile a host-side provider without a CUDA context or a device consumer. + +Native regression checks are opt-in: + +.. code-block:: shell + + CUDA_HOME=/usr/local/cuda-13.3 cargo test --workspace -- --include-ignored + +``CUDA_COOP_TEST_TARGET`` can select another supported NVRTC target; its default +is ``compute_120``. These checks compile the supported Direct/Transpose forms at +item counts 1 and 4, compare repeated LTO output, and check exact layout decoding +and specialization identity. They do not execute GPU kernels. + +Compiler integration and API reference +-------------------------------------- + +``Provider::from_checkout`` snapshots headers from the explicitly selected CCCL +checkout and toolkit, then loads NVRTC from that toolkit's ``lib64`` directory. +Those native libraries are trusted executable compiler inputs. No Python wheel, +working-directory header lookup, or runtime query kernel is involved. Reuse the +provider and batch requests for a common target into one NVRTC compilation. + +The same NVRTC program produces both LTO-IR and each CUB ``TempStorage`` layout. +The result distinguishes scratch-free specializations from an empty C++ type's +nonzero ``sizeof``. Input identities include the request, header contents, +provider source, loaded NVRTC/builtins bytes, version, target, and compiler options. +Symbols include a specialization digest. Request planning must assign each +specialization to exactly one batch per final link: a batch is one indivisible +LTO-IR buffer, so partially overlapping batches would duplicate strong definitions. + +Rustdoc is the API reference, including the generated ABI's pointer, validity, +mutation, and collective requirements. Open ``rust/target/doc/cuda_coop_codegen/index.html`` +after the documentation command above, or run ``cargo doc --no-deps --open`` from +``rust/``. Device discovery, lowering, linking, launch enforcement, and Cargo +invalidation of external compiler inputs belong to the compiler integration. + +The Oxide adapter's ``oxide::serve_inputs`` answers +``--oxide-library-inputs-v1`` with the selected header roots and NVRTC library +paths, without loading CUDA. A provider build script dispatches this query +alongside ``--oxide-library-request-v1``. Device registration getters use +``cuda_device::library::registration!("REGISTRATION_ENV_NAME")`` so the frontend +tracks those files, directory contents and selected toolkit before emitting +Cargo dependency information. Host checks and rustdoc do not query CUDA inputs. +Cargo's timestamp freshness rules apply; replacements with unchanged or older +timestamps may require an explicit rebuild. diff --git a/rust/.gitignore b/rust/.gitignore new file mode 100644 index 000000000000..b83d22266ac8 --- /dev/null +++ b/rust/.gitignore @@ -0,0 +1 @@ +/target/ diff --git a/rust/AGENTS.md b/rust/AGENTS.md new file mode 100644 index 000000000000..6493480b108a --- /dev/null +++ b/rust/AGENTS.md @@ -0,0 +1,63 @@ +# Rust `cuda.coop` guidance + +Read the repository's root guidance first. Apply this file to the Rust +workspace, its examples and tests, and `docs/rust/`. + +## Design tenets + +The design tenets of `cuda.coop` are **performant, consistent, delightful**. + +- **Performant:** target negligible overhead over CUB through compile-time + planning, inlining and LTO. Avoid hidden runtime work. Inspect the final + device artifact and measure representative execution before making claims + about register use, spills or overhead. +- **Consistent:** preserve common operation names, group-first arguments, + defaults, CUB behavior, mutation and result ownership. Document necessary + Rust syntax and safety adaptations. Additional capabilities extend shared + behavior without changing it. +- **Delightful:** use idiomatic Rust, useful defaults, inferred element types, + reusable generic kernels, ordinary Cargo builds and actionable errors. + +## API and examples + +- Use the single `cuda_coop` API, normally with `use cuda_coop as coop;`. + State the supported Oxide compiler/toolchain explicitly. Do not duplicate + primitive namespaces to separate built-in operators from static callables. +- Make the main per-thread item count a kernel const generic named + `ITEMS_PER_THREAD`, propagated through helpers and host specialization. + Keep operation-specific auxiliary and output extents independent. +- Infer payload element types from typed producers or assignments while + preserving widths. Investigate inference failures instead of adding a + redundant element type to an ordinary example. Document full signatures. +- Load fills its payload in place and returns unit. Preserve fresh results + where the common operation requires them, and reflect Store mutation in + borrowing and validity. Undefined, partial and root-only results must not + become unrestricted initialized Rust values. +- Preserve hierarchy units, supported groups, storage sharing and explicit + synchronization obligations. A group descriptor or Rust borrow alone is + not a proof of collective participation or shared-memory safety. +- Required collective completion must not depend solely on `Drop` running. + +## Ownership and validation + +Keep CUB specialization, exact scratch layout queries and operation semantics +in the native host provider. Keep compiler discovery, lowering, allocation, +device linking and launch enforcement in Oxide. Host compiler libraries and +filesystem dependencies must stay out of device code. CMake consumers do not +acquire a Rust requirement. + +Every public item needs useful rustdoc: arguments and generic parameters, +defaults, supported types/groups, participation, layout, mutation and result +validity, scratch/synchronization, errors and examples. Every public unsafe +item needs a concrete `# Safety` contract. Use rustdoc as the API reference +and the CCCL Rust guide for complete programs and explanations. + +Prefer end-user primitive tests, with item counts 1 and 4 where supported. +Use focused compiler/provider fixtures for independently reviewable compiler +capabilities. Keep displayed examples and their executed source identical. +Distinguish host tests, rustdoc/doctests, device compilation, GPU execution, +sanitizers and current-head CI. A sanitizer pass is not a Rust soundness proof. + +Run formatting, applicable lints, warning-free documentation and meaningful +targeted tests. Correct foundational defects at their owning layer and +requalify dependents instead of hiding a redesign in a later primitive family. diff --git a/rust/Cargo.lock b/rust/Cargo.lock new file mode 100644 index 000000000000..b1c2103ebcbe --- /dev/null +++ b/rust/Cargo.lock @@ -0,0 +1,208 @@ +# This file is automatically @generated by Cargo. +# It is not intended for manual editing. +version = 4 + +[[package]] +name = "block-buffer" +version = "0.10.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3078c7629b62d3f0439517fa394996acacc5cbc91c5a20d8c658e77abd503a71" +dependencies = [ + "generic-array", +] + +[[package]] +name = "cfg-if" +version = "1.0.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4e7648175b45a9a48536d676f68d918270699102aa8dab5496df06904c914600" + +[[package]] +name = "cpufeatures" +version = "0.2.17" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "59ed5838eebb26a2bb2e58f6d5b5316989ae9d08bab10e0e6d103e656d1b0280" +dependencies = [ + "libc", +] + +[[package]] +name = "crypto-common" +version = "0.1.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "78c8292055d1c1df0cce5d180393dc8cce0abec0a7102adb6c7b1eef6016d60a" +dependencies = [ + "generic-array", + "typenum", +] + +[[package]] +name = "cuda-coop-codegen" +version = "0.0.0" +dependencies = [ + "libloading", + "serde", + "serde_json", + "sha2", +] + +[[package]] +name = "digest" +version = "0.10.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9ed9a281f7bc9b7576e61468ba615a66a5c8cfdff42420a70aa82701a3b1e292" +dependencies = [ + "block-buffer", + "crypto-common", +] + +[[package]] +name = "generic-array" +version = "0.14.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "85649ca51fd72272d7821adaf274ad91c288277713d9c18820d8499a7ff69e9a" +dependencies = [ + "typenum", + "version_check", +] + +[[package]] +name = "itoa" +version = "1.0.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8f42a60cbdf9a97f5d2305f08a87dc4e09308d1276d28c869c684d7777685682" + +[[package]] +name = "libc" +version = "0.2.191" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "65cd87c9a6d60536580b316ad985726104730a5151c99845f2490428ef91240a" + +[[package]] +name = "libloading" +version = "0.8.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d7c4b02199fee7c5d21a5ae7d8cfa79a6ef5bb2fc834d6e9058e89c825efdc55" +dependencies = [ + "cfg-if", + "windows-link", +] + +[[package]] +name = "memchr" +version = "2.8.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cf8baf1c55e62ffcace7a9f06f4bd9cd3f0c4beb022d3b367256b91b87513d98" + +[[package]] +name = "proc-macro2" +version = "1.0.107" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "985e7ec9bb745e6ce6535b544d84d6cd6f7ad8bd711c398938ae983b91a766d9" +dependencies = [ + "unicode-ident", +] + +[[package]] +name = "quote" +version = "1.0.47" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1fbf4db142a473a8d80c26bbf18454ed458bf8d26c8219c331daecfdbd079001" +dependencies = [ + "proc-macro2", +] + +[[package]] +name = "serde" +version = "1.0.229" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4148590afebada386688f18773da617792bf2ef03ffc1e4cbd2b1d45b023e0ba" +dependencies = [ + "serde_core", + "serde_derive", +] + +[[package]] +name = "serde_core" +version = "1.0.229" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "67dca2c9c51e58a4791a4b1ed58308b39c64224d349a935ab5039aa360942a48" +dependencies = [ + "serde_derive", +] + +[[package]] +name = "serde_derive" +version = "1.0.229" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e7a5d71263a5a7d47b41f6b3f06ba276f10cc18b0931f1799f710578e2309348" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "serde_json" +version = "1.0.154" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e7e9cc8b1b85264074fbcc02a88680c4096b1e47df8f739dceb03bf482f04bd6" +dependencies = [ + "itoa", + "memchr", + "serde", + "serde_core", + "zmij", +] + +[[package]] +name = "sha2" +version = "0.10.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a7507d819769d01a365ab707794a4084392c824f54a7a6a7862f8c3d0892b283" +dependencies = [ + "cfg-if", + "cpufeatures", + "digest", +] + +[[package]] +name = "syn" +version = "3.0.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d78c8dee4c7bf0e14673097256fed6142ce9d3b85a408189d07482442145823b" +dependencies = [ + "proc-macro2", + "quote", + "unicode-ident", +] + +[[package]] +name = "typenum" +version = "1.20.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b6f5e870be6c3b371b77fe0ee0bafb859fa4964b4404c27de1d380043c4dda20" + +[[package]] +name = "unicode-ident" +version = "1.0.27" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a2c754d6c33795a1c324727428e5a7dedb5b06195f9890bdbcba760d3e246563" + +[[package]] +name = "version_check" +version = "0.9.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0b928f33d975fc6ad9f86c8f283853ad26bdd5b10b7f1542aa2fa15e2289105a" + +[[package]] +name = "windows-link" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f0805222e57f7521d6a62e36fa9163bc891acd422f971defe97d64e70d0a4fe5" + +[[package]] +name = "zmij" +version = "1.0.23" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "29666d0abbfad1e3dc4dcf6144730dd3a3ab225bbbdac83319345b1b44ccfc1b" diff --git a/rust/Cargo.toml b/rust/Cargo.toml new file mode 100644 index 000000000000..c90fafb1430a --- /dev/null +++ b/rust/Cargo.toml @@ -0,0 +1,18 @@ +# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. +# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +[workspace] +members = ["cuda-coop-codegen"] +resolver = "3" + +[workspace.package] +edition = "2024" +license = "Apache-2.0 WITH LLVM-exception" +rust-version = "1.85" + +[workspace.lints.rust] +missing_docs = "deny" +unsafe_op_in_unsafe_fn = "deny" + +[workspace.lints.rustdoc] +broken_intra_doc_links = "deny" diff --git a/rust/cuda-coop-codegen/Cargo.toml b/rust/cuda-coop-codegen/Cargo.toml new file mode 100644 index 000000000000..7a92fcaaded4 --- /dev/null +++ b/rust/cuda-coop-codegen/Cargo.toml @@ -0,0 +1,19 @@ +# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. +# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +[package] +name = "cuda-coop-codegen" +version = "0.0.0" +edition.workspace = true +license.workspace = true +rust-version.workspace = true +publish = false + +[dependencies] +libloading = "0.8" +serde = { version = "1", features = ["derive"] } +serde_json = "1" +sha2 = "0.10" + +[lints] +workspace = true diff --git a/rust/cuda-coop-codegen/examples/specialize.rs b/rust/cuda-coop-codegen/examples/specialize.rs new file mode 100644 index 000000000000..3a2cfe26364b --- /dev/null +++ b/rust/cuda-coop-codegen/examples/specialize.rs @@ -0,0 +1,35 @@ +// Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +//! Compile a host-side provider without a CUDA context or a device consumer. + +use cuda_coop_codegen::{Algorithm, Operation, Provider, Request}; +use std::{env, path::Path, time::Instant}; + +fn main() -> Result<(), Box> { + let arguments: Vec<_> = env::args().collect(); + if arguments.len() != 4 { + return Err("usage: specialize CUDA_TOOLKIT COMPUTE_TARGET OUTPUT_LTOIR".into()); + } + // This example is built inside the selected source checkout. Consumers + // supply their own explicit source root instead of relying on their cwd. + let checkout = Path::new(env!("CARGO_MANIFEST_DIR")).join("../.."); + let started = Instant::now(); + let provider = Provider::from_checkout(checkout, &arguments[1])?; + let initialized = started.elapsed(); + let request = Request::new(Operation::LoadFilled, Algorithm::Transpose, 128, 4)?; + let started = Instant::now(); + let artifact = provider.compile(&[request], &arguments[2])?; + let compiled = started.elapsed(); + std::fs::write(&arguments[3], &artifact.ltoir)?; + println!("symbol: {}", artifact.specializations[0].symbol); + println!( + "CUB layout: {:?}; scratch: {:?}; scratch free: {}", + artifact.specializations[0].cub_layout, + artifact.specializations[0].scratch, + artifact.specializations[0].scratch_free + ); + println!("initialization: {initialized:?}; compilation: {compiled:?}"); + print!("{}", artifact.compiler_log); + Ok(()) +} diff --git a/rust/cuda-coop-codegen/src/headers.rs b/rust/cuda-coop-codegen/src/headers.rs new file mode 100644 index 000000000000..dd55da52337e --- /dev/null +++ b/rust/cuda-coop-codegen/src/headers.rs @@ -0,0 +1,95 @@ +// Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +//! Snapshot compiler inputs so identity describes the actual compilation bytes. + +use crate::{Error, Result}; +use sha2::{Digest, Sha256}; +use std::{ + collections::BTreeMap, + ffi::CString, + fs, + path::{Path, PathBuf}, +}; + +pub(crate) struct Headers { + pub(crate) entries: Vec<(CString, CString)>, + pub(crate) digest: [u8; 32], +} + +impl Headers { + pub(crate) fn roots(checkout: &Path, toolkit: &Path) -> [(PathBuf, &'static str); 4] { + [ + (checkout.join("cub/cub"), "cub/"), + (checkout.join("libcudacxx/include"), ""), + (checkout.join("thrust/thrust"), "thrust/"), + (toolkit.join("include"), ""), + ] + } + + pub(crate) fn read(checkout: &Path, toolkit: &Path) -> Result { + let mut files = BTreeMap::new(); + // First entry wins: the selected checkout always owns CCCL headers, + // including when the toolkit also bundles an older CCCL installation. + for (root, prefix) in Self::roots(checkout, toolkit) { + visit(&root, &root, prefix, &mut files)?; + } + let mut hash = Sha256::new(); + let mut entries = Vec::with_capacity(files.len()); + for (name, contents) in files { + hash_field(&mut hash, name.as_bytes()); + hash_field(&mut hash, &contents); + let content = CString::new(contents) + .map_err(|_| Error::message(format!("NUL byte in header: {name}")))?; + entries.push((CString::new(name)?, content)); + } + Ok(Self { + entries, + digest: hash.finalize().into(), + }) + } +} + +fn visit( + root: &Path, + directory: &Path, + prefix: &str, + files: &mut BTreeMap>, +) -> Result<()> { + for entry in fs::read_dir(directory).map_err(|e| Error::input(directory, e))? { + let entry = entry?; + let path = entry.path(); + if path.is_dir() { + // Header-directory symlinks would need cycle handling. The include + // root itself may be a symlink (as in normal CUDA installations). + if entry.file_type()?.is_symlink() { + return Err(Error::message(format!( + "header directory symlink is unsupported: {}", + path.display() + ))); + } + visit(root, &path, prefix, files)?; + } else if path.is_file() { + let relative = path.strip_prefix(root).expect("child of include root"); + let name = format!( + "{prefix}{}", + relative + .to_str() + .ok_or_else(|| Error::message(format!( + "non-UTF-8 header path: {}", + path.display() + )))? + .replace('\\', "/") + ); + if let std::collections::btree_map::Entry::Vacant(slot) = files.entry(name) { + slot.insert(fs::read(&path).map_err(|e| Error::input(&path, e))?); + } + } + } + Ok(()) +} + +pub(crate) fn hash_field(hash: &mut Sha256, bytes: &[u8]) { + hash.update((bytes.len() as u64).to_le_bytes()); + hash.update(bytes); +} diff --git a/rust/cuda-coop-codegen/src/lib.rs b/rust/cuda-coop-codegen/src/lib.rs new file mode 100644 index 000000000000..872adf16c76a --- /dev/null +++ b/rust/cuda-coop-codegen/src/lib.rs @@ -0,0 +1,463 @@ +// Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +//! Native, host-only CUB specialization for Rust cooperative algorithms. +//! +//! This private development crate compiles i32 block Load/Store providers with +//! NVRTC. It returns linkable LTO-IR and the exact C++ scratch layout from the +//! **same compilation**, without a CUDA context, GPU or Python dependency. +//! It is a compiler building block, not a device API or a package release. +//! +//! # Supported specializations +//! +//! [`Request`] describes a one-dimensional complete block, positive fixed +//! items per thread, and [`Algorithm::Direct`] or [`Algorithm::Transpose`]. +//! Loads either fill the full tile or take an explicit default for the masked +//! suffix. Stores require all payload slots initialized, even for partial +//! stores. Transpose Store replaces each lane's blocked payload with striped +//! items from across the block. No-default masked loads are +//! intentionally absent from this provider surface. +//! +//! [`Provider::from_checkout`] snapshots the selected CCCL and CUDA headers; +//! [`Provider::compile`] uses those exact bytes and loaded NVRTC to emit an +//! [`Artifact`]. Each returned device ABI is unsafe to call: its documented +//! participation, pointer, capacity, validity and synchronization obligations +//! belong to the compiler integration. This crate allocates no device memory, +//! links no final kernel and supplies no launch policy. +//! +//! # Host example +//! +//! ```no_run +//! use cuda_coop_codegen::{Algorithm, Operation, Provider, Request}; +//! +//! # fn main() -> Result<(), Box> { +//! let provider = Provider::from_checkout("/path/to/cccl", "/usr/local/cuda")?; +//! let request = Request::new(Operation::LoadFull, Algorithm::Transpose, 128, 4)?; +//! let artifact = provider.compile(&[request], "compute_120")?; +//! std::fs::write("block_load.ltoir", &artifact.ltoir)?; +//! println!("{}: {:?}", artifact.specializations[0].symbol, artifact.specializations[0].scratch); +//! # Ok(()) } +//! ``` +//! +//! The example compiles on the host; executing it requires a compatible Linux +//! CUDA toolkit. The `specialize` example writes the artifact and prints its +//! identity/layout. `cargo doc --no-deps` generates this API reference. + +#![deny(unsafe_code)] + +mod headers; +#[allow(unsafe_code)] +mod nvrtc; +pub mod oxide; + +use headers::{Headers, hash_field}; +use sha2::{Digest, Sha256}; +use std::{fmt, path::Path}; + +/// A provider input, loader or compiler error, including NVRTC's compile log. +#[derive(Debug)] +pub struct Error(String); +impl Error { + pub(crate) fn message(message: impl Into) -> Self { + Self(message.into()) + } + pub(crate) fn input(path: &Path, error: std::io::Error) -> Self { + Self(format!("cannot read {}: {error}", path.display())) + } +} +impl fmt::Display for Error { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.write_str(&self.0) + } +} +impl std::error::Error for Error {} +impl From for Error { + fn from(error: std::io::Error) -> Self { + Self(error.to_string()) + } +} +impl From for Error { + fn from(error: std::ffi::NulError) -> Self { + Self(error.to_string()) + } +} +/// Result returned by native provider operations. +pub type Result = std::result::Result; + +/// CUB block layout algorithm; both variants produce/consume blocked payloads. +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub enum Algorithm { + /// Consecutive items per thread, without shared scratch or a reuse barrier. + Direct, + /// Coalesced access with shared-memory transposition to/from blocked layout. + Transpose, +} +impl Algorithm { + fn token(self) -> &'static str { + match self { + Self::Direct => "DIRECT", + Self::Transpose => "TRANSPOSE", + } + } +} + +/// Device operation and its validity controls, all with i32 elements. +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub enum Operation { + /// Load every slot. ABI: `(const int* memory, int* payload, unsigned char* scratch)`. + LoadFull, + /// Load a valid prefix and fill its suffix. ABI adds `(int valid, int default_value)`. + LoadFilled, + /// Store every slot. ABI: `(int* memory, int* payload, unsigned char* scratch)`. + StoreFull, + /// Store a prefix only. ABI adds `(int valid)`; other destination elements remain unchanged. + StorePartial, +} +impl Operation { + fn is_load(self) -> bool { + matches!(self, Self::LoadFull | Self::LoadFilled) + } + fn token(self) -> &'static str { + match self { + Self::LoadFull => "load_full", + Self::LoadFilled => "load_filled", + Self::StoreFull => "store_full", + Self::StorePartial => "store_partial", + } + } +} + +/// A validated i32, one-dimensional block specialization. +/// +/// This is an in-process implementation interface, not a serialized compiler +/// protocol. All fields are static specialization inputs. The generated ABI +/// uses runtime `valid` and `default_value` only when selected by [`Operation`]. +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub struct Request { + operation: Operation, + algorithm: Algorithm, + block_threads: u32, + items_per_thread: u32, +} +impl Request { + /// Validate an exact block shape and payload extent. + /// + /// `block_threads` must be 1..=1024 and `items_per_thread` positive. Their + /// product must fit i32, matching CUB's valid-prefix count. Target-specific + /// compiler restrictions are diagnosed by NVRTC during [`Provider::compile`]. + /// + /// # Errors + /// + /// Returns an error for zero/oversized blocks, zero item counts or a tile + /// extent that cannot be represented by CUB's signed valid-prefix count. + pub fn new( + operation: Operation, + algorithm: Algorithm, + block_threads: u32, + items_per_thread: u32, + ) -> Result { + if !(1..=1024).contains(&block_threads) + || items_per_thread == 0 + || u64::from(block_threads) * u64::from(items_per_thread) > i32::MAX as u64 + { + return Err(Error::message( + "require 1..=1024 block threads, positive items per thread and tile extent <= i32::MAX", + )); + } + Ok(Self { + operation, + algorithm, + block_threads, + items_per_thread, + }) + } + /// The operation selected by this request. + pub fn operation(self) -> Operation { + self.operation + } + /// The layout algorithm selected by this request. + pub fn algorithm(self) -> Algorithm { + self.algorithm + } + /// Exact number of threads in the one-dimensional block. + pub fn block_threads(self) -> u32 { + self.block_threads + } + /// Number of i32 payload elements per thread. + pub fn items_per_thread(self) -> u32 { + self.items_per_thread + } +} + +/// Byte size and minimum byte alignment of a storage region. +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub struct Layout { + /// Required byte capacity; zero means no shared scratch is accessed. + pub size: usize, + /// Positive power-of-two alignment in bytes. + pub alignment: usize, +} + +/// One CUB specialization and its exact storage requirements. +/// +/// # Device ABI obligations +/// +/// Every thread in the exact block executes the provider in converged control +/// flow. Memory is device-accessible and properly aligned; input/output spans +/// cover the full tile or selected valid prefix. Payloads contain the requested +/// number of i32 slots. Store payloads are fully initialized and may be mutated. +/// `memory` is the same block tile base pointer in every lane; the provider +/// applies no block-origin offset for the caller. +/// Load output and Store input must not overlap the memory tile or scratch. +/// The caller supplies disjoint collective ownership of the memory tile. +/// +/// Partial controls are block-uniform, with `0 <= valid <= tile_size`; generated +/// code traps on invalid counts. A filled Load's default is also block-uniform. +/// Transpose scratch satisfies [`Self::scratch`] and cannot alias another live +/// region. The caller synchronizes before reuse, including loop iterations; +/// the provider does not add a trailing reuse barrier. Direct ignores scratch, +/// so a null pointer is permitted for that argument only. All providers return +/// unit; loads fill payloads in place and stores write memory. +#[derive(Debug)] +pub struct Specialization { + /// SHA-256 of request/source, target/options, captured headers and native toolchain inputs. + pub identity: [u8; 32], + /// Name of the `extern "C"` device function in the LTO artifact. + pub symbol: String, + + /// Actual CUB TempStorage sizeof/alignof from the same NVRTC program. + pub cub_layout: Layout, + /// Required shared region; Direct uses zero bytes despite its empty C++ type's sizeof=1. + pub scratch: Layout, + /// Whether this specialization accesses no shared scratch or scratch-reuse barrier. + pub scratch_free: bool, + /// Whether Store redistributes payloads across lanes into striped order. + /// True only for Transpose stores; every slot remains initialized. + pub mutates_payload: bool, +} + +/// A batch emitted by one NVRTC compilation with one target and header snapshot. +/// +/// [`Self::specializations`] preserves request order. Its symbols include each +/// specialization's full input digest, independently of other batch members. +/// Request planning must assign each specialization to exactly one batch per +/// final link. A batch is indivisible; a symbol name cannot deduplicate its +/// strong definitions against a partially overlapping batch. +#[derive(Debug)] +pub struct Artifact { + /// SHA-256 input identity including ordered specialization identities. + pub identity: [u8; 32], + /// Per-request symbols and exact layouts, in request order. + pub specializations: Vec, + /// One LTO-IR buffer containing every requested specialization. + pub ltoir: Vec, + /// Generated C++ provider source, including each precise device ABI. + pub source: String, + /// Exact NVRTC options, in order. + pub options: Vec, + /// NVRTC diagnostics, including successful-compilation warnings. + pub compiler_log: String, +} + +/// Host compiler bound to one immutable header snapshot and CUDA toolkit. +/// +/// Construct once for a compilation and reuse for independent specializations. +/// No cache, environment-based header search, Driver API or GPU is involved. +/// Keep the selected toolkit installation unchanged while the provider is alive; +/// its native libraries are executable compiler inputs, not inert data. +pub struct Provider { + headers: Headers, + compiler: nvrtc::Compiler, +} +impl Provider { + /// Load Linux NVRTC from `toolkit/lib64` and snapshot header inputs. + /// + /// `checkout` must be a complete CCCL source checkout containing `cub/cub`, + /// `libcudacxx/include` and `thrust/thrust`. Paths are resolved explicitly; + /// neither the current directory nor an installed Python package is searched. + /// CCCL's selected headers take priority over the toolkit's bundled copies. + /// NVRTC and builtins are loaded by their explicit absolute paths, without + /// a link-time CUDA dependency. Header bytes are passed in memory to the + /// compiler, so later filesystem changes do not silently alter this + /// provider's input identity. + /// + /// # Errors + /// + /// Errors identify unreadable inputs, unsupported directory symlinks, missing + /// compiler libraries and incompatible NVRTC entry points. A new instance + /// observes changed headers/toolchain bytes. Linux is the initial supported host. + pub fn from_checkout(checkout: impl AsRef, toolkit: impl AsRef) -> Result { + let checkout = checkout + .as_ref() + .canonicalize() + .map_err(|e| Error::input(checkout.as_ref(), e))?; + let toolkit = toolkit + .as_ref() + .canonicalize() + .map_err(|e| Error::input(toolkit.as_ref(), e))?; + Ok(Self { + headers: Headers::read(&checkout, &toolkit)?, + compiler: nvrtc::Compiler::load(&toolkit)?, + }) + } + + /// Compile a nonempty batch for an explicit architecture such as `compute_120`. + /// + /// Every request uses the same captured headers, target, C++17 options and + /// LTO program. Duplicate requests are rejected, since each input identity + /// owns one definition. The order of requests is preserved in the output. + /// CUB layout name expressions are registered before compilation; literal + /// template arguments carry sizeof/alignof without another compilation or + /// query kernel. + /// + /// # Errors + /// + /// Empty/duplicate requests and malformed targets are rejected before native + /// compilation. Unsupported targets/instantiations return NVRTC's log. An + /// unrecognized layout name fails closed instead of guessing storage sizes. + pub fn compile(&self, requests: &[Request], target: &str) -> Result { + if !target.starts_with("compute_") + || !target[8..].chars().all(|c| c.is_ascii_alphanumeric()) + || target.len() <= 8 + { + return Err(Error::message( + "target must be an explicit NVRTC architecture such as compute_120", + )); + } + if requests.is_empty() { + return Err(Error::message("at least one specialization is required")); + } + for (index, request) in requests.iter().enumerate() { + if requests[..index].contains(request) { + return Err(Error::message("duplicate specialization in one batch")); + } + } + let options = vec![ + "--std=c++17".to_owned(), + format!("--gpu-architecture={target}"), + "--device-c".to_owned(), + "-dlto".to_owned(), + "--no-source-include".to_owned(), + ]; + let mut inputs = Sha256::new(); + for bytes in [ + b"cuda-coop-codegen-v1".as_slice(), + &self.headers.digest, + &self.compiler.identity, + include_bytes!("lib.rs"), + include_bytes!("headers.rs"), + include_bytes!("nvrtc.rs"), + ] { + hash_field(&mut inputs, bytes); + } + for option in &options { + hash_field(&mut inputs, option.as_bytes()); + } + let mut batch_identity = inputs.clone(); + let mut source = String::from( + "template \nextern __device__ unsigned char layout_probe;\n", + ); + let mut expressions = Vec::with_capacity(requests.len()); + let mut specializations = Vec::with_capacity(requests.len()); + for &request in requests { + let mut hash = inputs.clone(); + hash_field( + &mut hash, + self::source(request, "cuda_coop_specialization").as_bytes(), + ); + let identity: [u8; 32] = hash.finalize().into(); + hash_field(&mut batch_identity, &identity); + let suffix: String = identity.iter().map(|byte| format!("{byte:02x}")).collect(); + let symbol = format!("cuda_coop_{}_{}", request.operation.token(), suffix); + source.push_str(&self::source(request, &symbol)); + expressions.push(format!("&layout_probe")); + specializations.push(Specialization { + identity, + symbol, + cub_layout: Layout { + size: 0, + alignment: 1, + }, + scratch: Layout { + size: 0, + alignment: 1, + }, + scratch_free: request.algorithm == Algorithm::Direct, + mutates_payload: !request.operation.is_load() + && request.algorithm == Algorithm::Transpose, + }); + } + let output = self + .compiler + .compile(&source, &options, &expressions, &self.headers)?; + for (specialization, layout) in specializations.iter_mut().zip(output.layouts) { + specialization.cub_layout = layout; + if !specialization.scratch_free { + specialization.scratch = layout; + } + } + Ok(Artifact { + identity: batch_identity.finalize().into(), + specializations, + ltoir: output.ltoir, + source, + options, + compiler_log: output.log, + }) + } +} + +fn source(request: Request, symbol: &str) -> String { + let (header, primitive, algorithm, pointer) = if request.operation.is_load() { + ("load", "BlockLoad", "BLOCK_LOAD", "const int*") + } else { + ("store", "BlockStore", "BLOCK_STORE", "int*") + }; + let (controls, operands, check) = match request.operation { + Operation::LoadFull | Operation::StoreFull => ("", "", String::new()), + Operation::LoadFilled => ( + ", int valid, int default_value", + ", valid, default_value", + format!( + "if (valid < 0 || valid > {}) __trap();", + request.block_threads * request.items_per_thread + ), + ), + Operation::StorePartial => ( + ", int valid", + ", valid", + format!( + "if (valid < 0 || valid > {}) __trap();", + request.block_threads * request.items_per_thread + ), + ), + }; + let constructor = if request.algorithm == Algorithm::Direct { + "primitive operation;" + } else { + "primitive operation{*reinterpret_cast(scratch)};" + }; + let method = if request.operation.is_load() { + "Load" + } else { + "Store" + }; + format!( + r#"// Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception +#include +namespace {symbol}_detail {{ +using primitive = cub::{primitive}; +static_assert(sizeof(int) == 4 && alignof(int) == 4); +extern "C" __device__ void {symbol}({pointer} memory, int* payload, unsigned char* scratch{controls}) {{ + {check} + {constructor} + operation.{method}(memory, *reinterpret_cast(payload){operands}); +}} +}} +"#, + threads = request.block_threads, + items = request.items_per_thread, + choice = request.algorithm.token() + ) +} diff --git a/rust/cuda-coop-codegen/src/nvrtc.rs b/rust/cuda-coop-codegen/src/nvrtc.rs new file mode 100644 index 000000000000..fe0ca34e7658 --- /dev/null +++ b/rust/cuda-coop-codegen/src/nvrtc.rs @@ -0,0 +1,307 @@ +// Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +//! Native ABI boundary. Library handles outlive every copied entry point. + +use crate::{ + Error, Layout, Result, + headers::{Headers, hash_field}, +}; +use libloading::Library; +use sha2::{Digest, Sha256}; +use std::{ + ffi::{CStr, CString, c_char, c_int, c_void}, + fs, + path::{Path, PathBuf}, + ptr, +}; + +type Handle = *mut c_void; +type Status = c_int; +type Create = unsafe extern "C" fn( + *mut Handle, + *const c_char, + *const c_char, + c_int, + *const *const c_char, + *const *const c_char, +) -> Status; +type Destroy = unsafe extern "C" fn(*mut Handle) -> Status; +type AddExpression = unsafe extern "C" fn(Handle, *const c_char) -> Status; +type Compile = unsafe extern "C" fn(Handle, c_int, *const *const c_char) -> Status; +type GetSize = unsafe extern "C" fn(Handle, *mut usize) -> Status; +type GetBytes = unsafe extern "C" fn(Handle, *mut c_char) -> Status; +type GetName = unsafe extern "C" fn(Handle, *const c_char, *mut *const c_char) -> Status; +type GetError = unsafe extern "C" fn(Status) -> *const c_char; + +pub(crate) struct Compiler { + _library: Library, + _builtins: Library, + pub(crate) identity: [u8; 32], + create: Create, + destroy: Destroy, + add_expression: AddExpression, + compile_program: Compile, + log_size: GetSize, + log: GetBytes, + lto_size: GetSize, + lto: GetBytes, + name: GetName, + error: GetError, +} +pub(crate) struct Output { + pub(crate) ltoir: Vec, + pub(crate) layouts: Vec, + pub(crate) log: String, +} + +impl Compiler { + pub(crate) fn input_paths(toolkit: &Path) -> [PathBuf; 2] { + ["libnvrtc-builtins.so", "libnvrtc.so"].map(|name| toolkit.join("lib64").join(name)) + } + + pub(crate) fn load(toolkit: &Path) -> Result { + if !cfg!(target_os = "linux") { + return Err(Error::message( + "the native provider currently supports Linux hosts", + )); + } + let mut hash = Sha256::new(); + let mut load = |input: &Path| -> Result { + let path = input.canonicalize().map_err(|e| Error::input(input, e))?; + let name = input + .file_name() + .and_then(|name| name.to_str()) + .expect("fixed library name"); + hash_field(&mut hash, name.as_bytes()); + hash_field( + &mut hash, + &fs::read(&path).map_err(|e| Error::input(&path, e))?, + ); + // SAFETY: the explicit toolkit is trusted native executable input; + // only documented NVRTC entry points are called with their C ABI. + unsafe { Library::new(&path) } + .map_err(|e| Error::message(format!("cannot load {}: {e}", path.display()))) + }; + let [builtins_path, library_path] = Self::input_paths(toolkit); + let builtins = load(&builtins_path)?; + let library = load(&library_path)?; + // SAFETY: symbols use signatures from nvrtc.h. Missing entry points + // return errors before a program is created; handles remain owned. + unsafe { + let version = library + .get:: Status>(b"nvrtcVersion\0") + .map_err(|e| Error::message(format!("missing NVRTC entry point: {e}")))?; + let (mut major, mut minor) = (0, 0); + if version(&mut major, &mut minor) != 0 { + return Err(Error::message("nvrtcVersion failed")); + } + hash_field(&mut hash, &major.to_le_bytes()); + hash_field(&mut hash, &minor.to_le_bytes()); + macro_rules! symbol { + ($name:literal, $ty:ty) => { + *library + .get::<$ty>(concat!($name, "\0").as_bytes()) + .map_err(|e| { + Error::message(format!("missing NVRTC entry point {}: {e}", $name)) + })? + }; + } + Ok(Self { + create: symbol!("nvrtcCreateProgram", Create), + destroy: symbol!("nvrtcDestroyProgram", Destroy), + add_expression: symbol!("nvrtcAddNameExpression", AddExpression), + compile_program: symbol!("nvrtcCompileProgram", Compile), + log_size: symbol!("nvrtcGetProgramLogSize", GetSize), + log: symbol!("nvrtcGetProgramLog", GetBytes), + lto_size: symbol!("nvrtcGetLTOIRSize", GetSize), + lto: symbol!("nvrtcGetLTOIR", GetBytes), + name: symbol!("nvrtcGetLoweredName", GetName), + error: symbol!("nvrtcGetErrorString", GetError), + identity: hash.finalize().into(), + _library: library, + _builtins: builtins, + }) + } + } + fn check(&self, status: Status, operation: &str) -> Result<()> { + if status == 0 { + return Ok(()); + } + // SAFETY: NVRTC returns a non-null static NUL-terminated error string. + let message = unsafe { CStr::from_ptr((self.error)(status)) }.to_string_lossy(); + Err(Error::message(format!("{operation}: {message} ({status})"))) + } + pub(crate) fn compile( + &self, + source: &str, + options: &[String], + expressions: &[String], + headers: &Headers, + ) -> Result { + let source = CString::new(source)?; + let options = options + .iter() + .map(|s| CString::new(s.as_str())) + .collect::, _>>()?; + let option_pointers: Vec<_> = options.iter().map(|s| s.as_ptr()).collect(); + let header_names: Vec<_> = headers + .entries + .iter() + .map(|(name, _)| name.as_ptr()) + .collect(); + let header_contents: Vec<_> = headers + .entries + .iter() + .map(|(_, contents)| contents.as_ptr()) + .collect(); + let count = c_int::try_from(headers.entries.len()) + .map_err(|_| Error::message("too many input headers"))?; + let mut program = Program { + raw: ptr::null_mut(), + compiler: self, + }; + // SAFETY: C strings and arrays remain live until after program destroy; + // NVRTC initializes the uniquely owned handle. + self.check( + unsafe { + (self.create)( + &mut program.raw, + source.as_ptr(), + c"cuda_coop.cu".as_ptr(), + count, + header_contents.as_ptr(), + header_names.as_ptr(), + ) + }, + "nvrtcCreateProgram", + )?; + let expressions = expressions + .iter() + .map(|s| CString::new(s.as_str())) + .collect::, _>>()?; + for expression in &expressions { + // SAFETY: program and terminated expression remain live through compilation. + self.check( + unsafe { (self.add_expression)(program.raw, expression.as_ptr()) }, + "nvrtcAddNameExpression", + )?; + } + // SAFETY: program and option strings are live for this synchronous call. + let status = unsafe { + (self.compile_program)( + program.raw, + option_pointers.len() as c_int, + option_pointers.as_ptr(), + ) + }; + let mut log = self.bytes(program.raw, self.log_size, self.log)?; + if log.last() == Some(&0) { + log.pop(); + } + let log = String::from_utf8_lossy(&log).into_owned(); + self.check(status, "nvrtcCompileProgram") + .map_err(|e| Error::message(format!("{e}\n{log}")))?; + let mut layouts = Vec::with_capacity(expressions.len()); + for expression in &expressions { + let mut lowered = ptr::null(); + // SAFETY: successful compilation owns the returned name until destroy. + self.check( + unsafe { (self.name)(program.raw, expression.as_ptr(), &mut lowered) }, + "nvrtcGetLoweredName", + )?; + if lowered.is_null() { + return Err(Error::message("NVRTC returned a null layout name")); + } + // SAFETY: successful query returned a program-owned C string. + layouts.push(decode_layout( + unsafe { CStr::from_ptr(lowered) } + .to_str() + .map_err(|_| Error::message("non-UTF-8 layout name"))?, + )?); + } + let ltoir = self.bytes(program.raw, self.lto_size, self.lto)?; + if ltoir.is_empty() { + return Err(Error::message("NVRTC emitted empty LTO-IR")); + } + Ok(Output { + ltoir, + layouts, + log, + }) + } + fn bytes(&self, program: Handle, size: GetSize, get: GetBytes) -> Result> { + let mut length = 0; + // SAFETY: program is live; NVRTC writes one size_t. + self.check(unsafe { size(program, &mut length) }, "NVRTC output size")?; + let mut bytes = vec![0; length]; + if length != 0 { + // SAFETY: exact reported byte capacity is allocated and writable. + self.check( + unsafe { get(program, bytes.as_mut_ptr().cast()) }, + "NVRTC output", + )?; + } + Ok(bytes) + } +} +struct Program<'a> { + raw: Handle, + compiler: &'a Compiler, +} +impl Drop for Program<'_> { + fn drop(&mut self) { + if !self.raw.is_null() { + // SAFETY: unique owner; the compiler library remains loaded. + unsafe { + (self.compiler.destroy)(&mut self.raw); + } + } + } +} +fn decode_layout(name: &str) -> Result { + let invalid = || Error::message(format!("unsupported NVRTC layout name: {name}")); + let values = name + .strip_prefix("_Z12layout_probeILy") + .and_then(|s| s.strip_suffix("EE")) + .ok_or_else(invalid)?; + let (size, alignment) = values.split_once("ELy").ok_or_else(invalid)?; + let size: usize = size.parse().map_err(|_| invalid())?; + let alignment: usize = alignment.parse().map_err(|_| invalid())?; + if size == 0 || !alignment.is_power_of_two() || size % alignment != 0 { + return Err(invalid()); + } + Ok(Layout { size, alignment }) +} +#[cfg(test)] +mod tests { + use super::*; + #[test] + fn layout_literals_and_rejection() { + assert_eq!( + decode_layout("_Z12layout_probeILy1ELy1EE").unwrap(), + Layout { + size: 1, + alignment: 1 + } + ); + assert_eq!( + decode_layout("_Z12layout_probeILy2064ELy16EE").unwrap(), + Layout { + size: 2064, + alignment: 16 + } + ); + for invalid in [ + "_Z12layout_probeILy0ELy1EE", + "_Z12layout_probeILy32ELy3EE", + "_Z12layout_probeILy33ELy16EE", + "_Z12layout_probeILy32ELy0EE", + "_Z12layout_probeILy18446744073709551616ELy1EE", + "other", + "_Z12layout_probeILy32ELy16EE.extra", + ] { + assert!(decode_layout(invalid).is_err(), "{invalid}"); + } + } +} diff --git a/rust/cuda-coop-codegen/src/oxide.rs b/rust/cuda-coop-codegen/src/oxide.rs new file mode 100644 index 000000000000..291e3f90004f --- /dev/null +++ b/rust/cuda-coop-codegen/src/oxide.rs @@ -0,0 +1,458 @@ +// Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +//! Version-one host transport for the Oxide library-operation marker. +//! +//! A dependency build script can serve as its already-built provider executable: +//! normal Cargo execution embeds [`registration`] in a device crate constant; +//! `--oxide-library-request-v1` execution passes stdin/stdout to [`serve`]; +//! `--oxide-library-inputs-v1` passes them to [`serve_inputs`] during macro expansion. +//! Use `cuda_device::library::registration!` in the getter to track these inputs. +//! The compiler reads the registration getter's constant MIR and invokes this +//! host process. This module never invokes Cargo or executes device Rust code. +//! The protocol supports the four i32 Load/Store forms documented by [`Operation`]. +//! Device calls retain all participation, validity and scratch obligations of +//! [`crate::Specialization`]. No allocation or synchronization policy is supplied. + +use crate::{Algorithm, Artifact, Error, Layout, Operation, Provider, Request, Result}; +use serde::{Deserialize, Serialize}; +use std::io::{Read, Write}; +use std::path::{Path, PathBuf}; + +const VERSION: u32 = 1; +const MAGIC: &[u8; 8] = b"OXDLIB01"; +const MAX_HEADER: usize = 1024 * 1024; +const MAX_LTO: usize = 256 * 1024 * 1024; + +#[derive(Serialize)] +struct Registration<'a> { + version: u32, + provider: &'a str, + executable: &'a Path, + operation: &'a str, + algorithm: &'a str, +} + +#[derive(Debug, Deserialize, Serialize, PartialEq)] +#[serde( + tag = "kind", + content = "bits", + rename_all = "snake_case", + deny_unknown_fields +)] +enum Scalar { + Signed(u16), + Unsigned(u16), + Float(u16), +} + +#[derive(Debug, Deserialize, Serialize, PartialEq)] +#[serde( + tag = "kind", + content = "scalar", + rename_all = "snake_case", + deny_unknown_fields +)] +enum Argument { + ConstPointer(Scalar), + MutPointer(Scalar), + Value(Scalar), +} + +#[derive(Debug, Deserialize)] +#[serde(deny_unknown_fields)] +struct BatchRequest { + version: u32, + target: String, + cuda_root: PathBuf, + requests: Vec, +} + +#[derive(Debug, Deserialize)] +#[serde(deny_unknown_fields)] +struct SpecializationRequest { + operation: String, + algorithm: String, + dtype: Scalar, + block: [u32; 3], + items_per_thread: u32, +} + +impl SpecializationRequest { + fn native(&self) -> Result { + if self.dtype != Scalar::Signed(32) { + return Err(Error::message( + "this provider supports signed i32 elements only", + )); + } + if self.block[1..] != [1, 1] { + return Err(Error::message( + "this provider requires the exact block shape [B,1,1]", + )); + } + let operation = match self.operation.as_str() { + "load_full" => Operation::LoadFull, + "load_filled" => Operation::LoadFilled, + "store_full" => Operation::StoreFull, + "store_partial" => Operation::StorePartial, + _ => { + return Err(Error::message(format!( + "unsupported operation {:?}", + self.operation + ))); + } + }; + let algorithm = match self.algorithm.as_str() { + "direct" => Algorithm::Direct, + "transpose" => Algorithm::Transpose, + _ => { + return Err(Error::message(format!( + "unsupported algorithm {:?}", + self.algorithm + ))); + } + }; + Request::new(operation, algorithm, self.block[0], self.items_per_thread) + } +} + +#[derive(Serialize)] +struct WireLayout { + size: u64, + alignment: u64, +} + +impl TryFrom for WireLayout { + type Error = Error; + fn try_from(layout: Layout) -> Result { + Ok(Self { + size: u64::try_from(layout.size) + .map_err(|_| Error::message("layout size exceeds u64"))?, + alignment: u64::try_from(layout.alignment) + .map_err(|_| Error::message("layout alignment exceeds u64"))?, + }) + } +} + +#[derive(Serialize)] +struct SpecializationResponse { + request_index: u32, + identity: [u8; 32], + symbol: String, + arguments: Vec, + cub_layout: WireLayout, + scratch: WireLayout, + scratch_free: bool, + mutates_payload: bool, +} + +#[derive(Serialize)] +struct BatchResponse { + version: u32, + batch_identity: [u8; 32], + specializations: Vec, +} + +fn arguments(operation: Operation) -> Vec { + let mut arguments = vec![ + if operation.is_load() { + Argument::ConstPointer(Scalar::Signed(32)) + } else { + Argument::MutPointer(Scalar::Signed(32)) + }, + Argument::MutPointer(Scalar::Signed(32)), + ]; + if matches!(operation, Operation::LoadFilled | Operation::StorePartial) { + arguments.push(Argument::Value(Scalar::Signed(32))); + } + if operation == Operation::LoadFilled { + arguments.push(Argument::Value(Scalar::Signed(32))); + } + arguments +} + +/// Encode a constant registration for one native recipe. +/// +/// `executable` is the live, already-built dependency build script, normally +/// obtained from `std::env::current_exe()`. Embed the returned string in the +/// device crate through `cargo::rustc-env`; a named getter must return that +/// constant directly. The marker uses the getter's function-item type. +/// +/// # Errors +/// +/// Rejects a nonabsolute executable or a path that cannot be encoded as JSON. +pub fn registration( + executable: &Path, + operation: Operation, + algorithm: Algorithm, +) -> Result { + if !executable.is_absolute() { + return Err(Error::message("provider executable path must be absolute")); + } + serde_json::to_string(&Registration { + version: VERSION, + provider: "cuda_coop", + executable, + operation: operation.token(), + algorithm: match algorithm { + Algorithm::Direct => "direct", + Algorithm::Transpose => "transpose", + }, + }) + .map_err(|error| Error::message(format!("cannot encode provider registration: {error}"))) +} + +#[derive(Deserialize)] +#[serde(deny_unknown_fields)] +struct InputRequest { + version: u32, + cuda_root: PathBuf, + target: String, +} + +#[derive(Serialize)] +struct InputResponse { + version: u32, + paths: Vec, + environment: Vec, +} + +/// Declare native inputs for Oxide's registration macro without loading CUDA. +/// +/// `checkout` selects the same CCCL tree as [`serve`]. The bounded version-one +/// JSON request contains `cuda_root` and `target`; the response lists header +/// directories and NVRTC libraries. The macro tracks files and directories +/// recursively before rustc emits dep-info, including new header files. +/// No GPU, native compiler invocation or library loading occurs here. +/// Cargo's timestamp freshness rules apply; replacements with unchanged or older +/// timestamps can remain undetected. Rebuild after restoring or replacing inputs +/// this way, including a symlink switched to an older installation. +/// +/// # Errors +/// +/// Rejects malformed, oversized or unsupported requests, nonabsolute toolkit +/// paths, invalid targets, and inaccessible checkout paths. +pub fn serve_inputs(checkout: &Path, input: impl Read, mut output: impl Write) -> Result<()> { + let mut bytes = Vec::new(); + input + .take((MAX_HEADER + 1) as u64) + .read_to_end(&mut bytes)?; + if bytes.len() > MAX_HEADER { + return Err(Error::message("provider input request exceeds 1 MiB")); + } + let request: InputRequest = serde_json::from_slice(&bytes) + .map_err(|error| Error::message(format!("invalid provider input request: {error}")))?; + if request.version != VERSION + || !request.cuda_root.is_absolute() + || !request.target.starts_with("compute_") + || request.target.len() <= 8 + || !request.target[8..] + .chars() + .all(|character| character.is_ascii_alphanumeric()) + { + return Err(Error::message( + "provider inputs require version 1, an absolute toolkit path and compute_XX target", + )); + } + let checkout = checkout + .canonicalize() + .map_err(|error| Error::input(checkout, error))?; + let paths = crate::headers::Headers::roots(&checkout, &request.cuda_root) + .into_iter() + .map(|(path, _)| path) + .chain(crate::nvrtc::Compiler::input_paths(&request.cuda_root)) + .collect(); + let response = serde_json::to_vec(&InputResponse { + version: VERSION, + paths, + environment: Vec::new(), + }) + .map_err(|error| Error::message(format!("cannot encode provider inputs: {error}")))?; + if response.len() > MAX_HEADER { + return Err(Error::message("provider input response exceeds 1 MiB")); + } + output.write_all(&response)?; + Ok(()) +} + +/// Compile one versioned stdin request and write its framed native response. +/// +/// `checkout` selects the CCCL checkout whose headers the provider captures. +/// The request supplies an explicit toolkit path and target. Stdout contains +/// only the frame; the caller reports errors on stderr and exits nonzero. +/// The 24-byte prefix is `OXDLIB01`, JSON header length, and raw LTO length +/// (both lengths are little-endian u64), followed by the two byte buffers. +/// Requests and response headers are limited to 1 MiB; LTO is limited to 256 MiB. +/// Runtime arguments are a scalar/raw-pointer tuple; scratch is inserted at +/// parameter two and uses a generic byte pointer after shared address conversion. +/// +/// # Errors +/// +/// Rejects malformed/version-mismatched requests, unsupported types/recipes, +/// invalid geometry, and duplicate specializations. Native compiler errors +/// retain NVRTC diagnostics. No response is written until compilation succeeds. +pub fn serve(checkout: &Path, input: impl Read, mut output: impl Write) -> Result<()> { + let mut bytes = Vec::new(); + input + .take((MAX_HEADER + 1) as u64) + .read_to_end(&mut bytes)?; + if bytes.len() > MAX_HEADER { + return Err(Error::message("provider request exceeds 1 MiB")); + } + let request: BatchRequest = serde_json::from_slice(&bytes) + .map_err(|error| Error::message(format!("invalid provider request: {error}")))?; + if request.version != VERSION { + return Err(Error::message(format!( + "unsupported provider protocol {}; expected {VERSION}", + request.version + ))); + } + let requests: Vec<_> = request + .requests + .iter() + .map(SpecializationRequest::native) + .collect::>()?; + let artifact = Provider::from_checkout(checkout, &request.cuda_root)? + .compile(&requests, &request.target)?; + write_response(&requests, artifact, &mut output) +} + +fn write_response(requests: &[Request], artifact: Artifact, output: &mut impl Write) -> Result<()> { + if artifact.ltoir.is_empty() || artifact.ltoir.len() > MAX_LTO { + return Err(Error::message( + "native LTO buffer must contain 1..=256 MiB bytes", + )); + } + let specializations = artifact + .specializations + .into_iter() + .zip(requests) + .enumerate() + .map(|(index, (entry, request))| { + Ok(SpecializationResponse { + request_index: u32::try_from(index) + .map_err(|_| Error::message("request index exceeds u32"))?, + identity: entry.identity, + symbol: entry.symbol, + arguments: arguments(request.operation()), + cub_layout: entry.cub_layout.try_into()?, + scratch: entry.scratch.try_into()?, + scratch_free: entry.scratch_free, + mutates_payload: entry.mutates_payload, + }) + }) + .collect::>>()?; + let header = serde_json::to_vec(&BatchResponse { + version: VERSION, + batch_identity: artifact.identity, + specializations, + }) + .map_err(|error| Error::message(format!("cannot encode provider response: {error}")))?; + if header.len() > MAX_HEADER { + return Err(Error::message("provider response header exceeds 1 MiB")); + } + output.write_all(MAGIC)?; + output.write_all(&(header.len() as u64).to_le_bytes())?; + output.write_all(&(artifact.ltoir.len() as u64).to_le_bytes())?; + output.write_all(&header)?; + output.write_all(&artifact.ltoir)?; + Ok(()) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn input_query_needs_no_installed_toolkit() { + let checkout = Path::new(env!("CARGO_MANIFEST_DIR")); + let toolkit = std::env::temp_dir().join("cuda-coop-input-query-no-toolkit"); + let request = serde_json::json!({"version":1,"cuda_root":toolkit,"target":"compute_120"}); + let mut bytes = Vec::new(); + serve_inputs(checkout, request.to_string().as_bytes(), &mut bytes).unwrap(); + let response: serde_json::Value = serde_json::from_slice(&bytes).unwrap(); + assert_eq!(response["version"], 1); + assert_eq!(response["paths"].as_array().unwrap().len(), 6); + assert!( + response["paths"] + .as_array() + .unwrap() + .iter() + .all(|path| Path::new(path.as_str().unwrap()).is_absolute()) + ); + assert!(response["environment"].as_array().unwrap().is_empty()); + for request in [ + serde_json::json!({"version":2,"cuda_root":toolkit,"target":"compute_120"}), + serde_json::json!({"version":1,"cuda_root":"relative","target":"compute_120"}), + serde_json::json!({"version":1,"cuda_root":toolkit,"target":"invalid"}), + ] { + assert!(serve_inputs(checkout, request.to_string().as_bytes(), Vec::new()).is_err()); + } + assert!(serve_inputs(checkout, vec![b' '; MAX_HEADER + 1].as_slice(), Vec::new()).is_err()); + } + + #[test] + fn registration_binds_native_recipe() { + let value: serde_json::Value = serde_json::from_str( + ®istration( + Path::new("/provider"), + Operation::LoadFilled, + Algorithm::Transpose, + ) + .unwrap(), + ) + .unwrap(); + assert_eq!(value["version"], 1); + assert_eq!(value["operation"], "load_filled"); + assert_eq!(value["algorithm"], "transpose"); + assert!( + registration( + Path::new("relative"), + Operation::LoadFull, + Algorithm::Direct + ) + .is_err() + ); + } + + #[test] + fn rejects_request_before_opening_toolkit() { + let base = serde_json::json!({"version":1,"target":"compute_120","cuda_root":"/absent", "requests":[{ + "operation":"load_full","algorithm":"direct","dtype":{"kind":"signed","bits":32},"block":[128,1,1],"items_per_thread":4 + }]}); + for (path, value, diagnostic) in [ + ( + "/version", + serde_json::json!(2), + "unsupported provider protocol", + ), + ( + "/requests/0/dtype/bits", + serde_json::json!(64), + "signed i32", + ), + ( + "/requests/0/block/1", + serde_json::json!(2), + "exact block shape", + ), + ( + "/requests/0/operation", + serde_json::json!("unknown"), + "unsupported operation", + ), + ] { + let mut request = base.clone(); + *request.pointer_mut(path).unwrap() = value; + let mut output = Vec::new(); + let error = serve( + Path::new("/absent"), + serde_json::to_vec(&request).unwrap().as_slice(), + &mut output, + ) + .unwrap_err(); + assert!(error.to_string().contains(diagnostic), "{error}"); + assert!(output.is_empty()); + } + } +} diff --git a/rust/cuda-coop-codegen/tests/native.rs b/rust/cuda-coop-codegen/tests/native.rs new file mode 100644 index 000000000000..4f8e16d38eb8 --- /dev/null +++ b/rust/cuda-coop-codegen/tests/native.rs @@ -0,0 +1,133 @@ +// Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +//! Real native compiler coverage; ordinary host tests need no CUDA installation. + +use cuda_coop_codegen::{Algorithm, Layout, Operation, Provider, Request}; +use std::{env, path::Path, time::Instant}; + +#[test] +fn rejects_invalid_shapes() { + for (threads, items) in [(0, 1), (1025, 1), (128, 0), (1024, u32::MAX)] { + assert!(Request::new(Operation::LoadFull, Algorithm::Direct, threads, items).is_err()); + } +} + +#[test] +#[ignore = "requires CUDA_HOME and a Linux CUDA toolkit with NVRTC LTO support"] +fn native_specializations_and_determinism() -> Result<(), Box> { + let toolkit = env::var("CUDA_HOME")?; + let target = env::var("CUDA_COOP_TEST_TARGET").unwrap_or_else(|_| "compute_120".into()); + let checkout = Path::new(env!("CARGO_MANIFEST_DIR")).join("../.."); + let start = Instant::now(); + let provider = Provider::from_checkout(&checkout, &toolkit)?; + eprintln!("provider initialization: {:?}", start.elapsed()); + let mut requests = Vec::new(); + for algorithm in [Algorithm::Direct, Algorithm::Transpose] { + for items in [1, 4] { + for operation in [ + Operation::LoadFull, + Operation::LoadFilled, + Operation::StoreFull, + Operation::StorePartial, + ] { + requests.push(Request::new(operation, algorithm, 128, items)?); + } + } + } + let start = Instant::now(); + let first = provider.compile(&requests, &target)?; + eprintln!( + "16-specialization batch: {:?}, {} bytes", + start.elapsed(), + first.ltoir.len() + ); + let second = provider.compile(&requests, &target)?; + assert_eq!(first.identity, second.identity); + assert_eq!(first.ltoir, second.ltoir, "repeated native output changed"); + assert_eq!(first.specializations.len(), requests.len()); + let mut identities = std::collections::HashSet::new(); + for ((entry, repeated), request) in first + .specializations + .iter() + .zip(&second.specializations) + .zip(&requests) + { + assert_eq!(entry.identity, repeated.identity); + assert_eq!(entry.symbol, repeated.symbol); + assert_eq!( + entry.mutates_payload, + request.algorithm() == Algorithm::Transpose + && matches!( + request.operation(), + Operation::StoreFull | Operation::StorePartial + ) + ); + assert!(identities.insert(entry.identity)); + eprintln!( + "{request:?}: CUB {:?}, scratch {:?}", + entry.cub_layout, entry.scratch + ); + if request.algorithm() == Algorithm::Direct { + assert!(entry.scratch_free); + assert_eq!( + entry.scratch, + Layout { + size: 0, + alignment: 1 + } + ); + assert_eq!( + entry.cub_layout, + Layout { + size: 1, + alignment: 1 + } + ); + } else { + assert!(!entry.scratch_free); + assert_eq!(entry.scratch, entry.cub_layout); + assert!(entry.scratch.size >= 128 * request.items_per_thread() as usize * 4); + } + } + // A specialization's symbol is independent of unrelated batch members. + let single = provider.compile(&requests[..1], &target)?; + assert_eq!( + single.specializations[0].identity, + first.specializations[0].identity + ); + assert_eq!( + single.specializations[0].symbol, + first.specializations[0].symbol + ); + assert_ne!(single.identity, first.identity); + assert!( + provider + .compile(&[], &target) + .unwrap_err() + .to_string() + .contains("at least one") + ); + assert!( + provider + .compile(&[requests[0]; 2], &target) + .unwrap_err() + .to_string() + .contains("duplicate") + ); + assert!( + provider + .compile(&requests[..1], "compute_120;injected") + .unwrap_err() + .to_string() + .contains("target must") + ); + assert!( + provider + .compile(&requests[..1], "compute_999999") + .unwrap_err() + .to_string() + .contains("nvrtcCompileProgram") + ); + Ok(()) +} diff --git a/rust/cuda-coop-codegen/tests/oxide-fixture/Cargo.toml b/rust/cuda-coop-codegen/tests/oxide-fixture/Cargo.toml new file mode 100644 index 000000000000..d2b9d7f747af --- /dev/null +++ b/rust/cuda-coop-codegen/tests/oxide-fixture/Cargo.toml @@ -0,0 +1,13 @@ +# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. +# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +[workspace] +members = ["api", "helper", "consumer"] +default-members = ["consumer"] +resolver = "3" + +[workspace.dependencies] +cuda-device = "0.2.1" +cuda-host = "0.2.1" +cuda-core = "0.4.0" +oxide-artifacts = { version = "0.2.1", features = ["object-read"] } diff --git a/rust/cuda-coop-codegen/tests/oxide-fixture/README.md b/rust/cuda-coop-codegen/tests/oxide-fixture/README.md new file mode 100644 index 000000000000..fb76581bf752 --- /dev/null +++ b/rust/cuda-coop-codegen/tests/oxide-fixture/README.md @@ -0,0 +1,38 @@ +# Native Oxide integration fixture + +This private unsafe fixture qualifies the compiler/provider boundary. It is not +an end-user collective API. The API crate embeds constant registrations from its +Cargo build script; the compiler invokes that same built executable for one +native request batch. A separate generic helper crate is reached from five +concrete kernel roots. No consumer supplies an LTO artifact list. + +From the CCCL checkout, with the agreed Oxide source and backend already built: + +```sh +python3 rust/cuda-coop-codegen/tests/run_oxide_fixture.py \ + --oxide-root /path/to/cuda-rust \ + --cargo-oxide /path/to/cargo-oxide \ + --backend /path/to/librustc_codegen_cuda.so \ + --cuda-root /path/to/cuda \ + --arch sm_120 \ + --output /path/to/new-evidence-directory +``` + +The runner copies the fixture into the new output directory and patches only +that copy to the supplied source checkout. It seeds dependency resolutions from +that checkout's `vecadd` lockfile. The selected nightly must match Oxide's pin. +The log and command receipt are retained beside the copied consumer. + +The 480 cases cover i32 Direct/Transpose Load and Store, items 1/4, one/three +blocks, full/default-filled Load, full/partial Store, post-Store payload values, +and Store → explicit block synchronization → Load of the same region. Every +block uses uniform origins, count, default and mode; per-thread payload arrays +are initialized before use. There is no no-default partial-validity claim. + +The fixture checks compiler-derived scratch sizes and exact block geometry, +rejects wrong shapes, thread counts and explicit byte counts, and runs with +`auto_shared`. Its helpers deliberately retain function boundaries to +exercise discovery; register/stack results are not an inlining baseline. The +materialized path uses `cargo oxide run --materialize-cubin` and the ordinary +embedded module loader. Convergent declarations require the coordinated O03a +compiler change. diff --git a/rust/cuda-coop-codegen/tests/oxide-fixture/api/Cargo.toml b/rust/cuda-coop-codegen/tests/oxide-fixture/api/Cargo.toml new file mode 100644 index 000000000000..0d5afeb1bb00 --- /dev/null +++ b/rust/cuda-coop-codegen/tests/oxide-fixture/api/Cargo.toml @@ -0,0 +1,14 @@ +# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. +# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +[package] +name = "coop-marker-fixture-api" +version = "0.0.0" +edition = "2024" +publish = false + +[dependencies] +cuda-device.workspace = true + +[build-dependencies] +cuda-coop-codegen = { path = "../../.." } diff --git a/rust/cuda-coop-codegen/tests/oxide-fixture/api/build.rs b/rust/cuda-coop-codegen/tests/oxide-fixture/api/build.rs new file mode 100644 index 000000000000..6c1abfb85736 --- /dev/null +++ b/rust/cuda-coop-codegen/tests/oxide-fixture/api/build.rs @@ -0,0 +1,41 @@ +// Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +use cuda_coop_codegen::{Algorithm, Operation, oxide}; + +fn main() { + let argument = std::env::args().nth(1); + if matches!( + argument.as_deref(), + Some("--oxide-library-request-v1" | "--oxide-library-inputs-v1") + ) { + let checkout = std::path::Path::new(env!("CCCL_SOURCE")); + let result = if argument.as_deref() == Some("--oxide-library-inputs-v1") { + oxide::serve_inputs(checkout, std::io::stdin().lock(), std::io::stdout().lock()) + } else { + oxide::serve(checkout, std::io::stdin().lock(), std::io::stdout().lock()) + }; + if let Err(error) = result { + eprintln!("{error}"); + std::process::exit(1); + } + return; + } + println!("cargo::rerun-if-changed=build.rs"); + println!("cargo::rerun-if-env-changed=CCCL_SOURCE"); + let executable = std::env::current_exe().unwrap(); + for (algorithm, tag) in [ + (Algorithm::Direct, "DIRECT"), + (Algorithm::Transpose, "TRANSPOSE"), + ] { + for (operation, name) in [ + (Operation::LoadFull, "LOAD_FULL"), + (Operation::LoadFilled, "LOAD_FILLED"), + (Operation::StoreFull, "STORE_FULL"), + (Operation::StorePartial, "STORE_PARTIAL"), + ] { + let registration = oxide::registration(&executable, operation, algorithm).unwrap(); + println!("cargo::rustc-env=COOP_{tag}_{name}={registration}"); + } + } +} diff --git a/rust/cuda-coop-codegen/tests/oxide-fixture/api/src/lib.rs b/rust/cuda-coop-codegen/tests/oxide-fixture/api/src/lib.rs new file mode 100644 index 000000000000..5a6d130add0c --- /dev/null +++ b/rust/cuda-coop-codegen/tests/oxide-fixture/api/src/lib.rs @@ -0,0 +1,95 @@ +// Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +//! Private unsafe fixture wrappers, with no public collective API contract. +#![no_std] + +macro_rules! registration { + ($name:ident, $variable:literal) => { + fn $name() -> &'static str { + cuda_device::library::registration!($variable) + } + }; +} +registration!(direct_load_full, "COOP_DIRECT_LOAD_FULL"); +registration!(direct_load_filled, "COOP_DIRECT_LOAD_FILLED"); +registration!(direct_store_full, "COOP_DIRECT_STORE_FULL"); +registration!(direct_store_partial, "COOP_DIRECT_STORE_PARTIAL"); +registration!(transpose_load_full, "COOP_TRANSPOSE_LOAD_FULL"); +registration!(transpose_load_filled, "COOP_TRANSPOSE_LOAD_FILLED"); +registration!(transpose_store_full, "COOP_TRANSPOSE_STORE_FULL"); +registration!(transpose_store_partial, "COOP_TRANSPOSE_STORE_PARTIAL"); + +/// # Safety +/// All threads participate with one block-wide input origin, valid count and +/// fill selection. The caller satisfies the exact CUB ABI and collective requirements. +#[inline(never)] +pub unsafe fn load( + input: *const i32, + values: *mut i32, + valid: i32, + filled: bool, +) { + unsafe { + if filled { + if TRANSPOSE { + cuda_device::library::call_v1::( + transpose_load_filled, + (input, values, valid, -7_i32), + ); + } else { + cuda_device::library::call_v1::( + direct_load_filled, + (input, values, valid, -7_i32), + ); + } + } else if TRANSPOSE { + cuda_device::library::call_v1::( + transpose_load_full, + (input, values), + ); + } else { + cuda_device::library::call_v1::( + direct_load_full, + (input, values), + ); + } + } +} + +/// # Safety +/// The payload is initialized. All threads participate with one block-wide +/// output origin, valid count and partial selection, satisfying the exact CUB ABI. +#[inline(never)] +pub unsafe fn store( + output: *mut i32, + values: *mut i32, + valid: i32, + partial: bool, +) { + unsafe { + if partial { + if TRANSPOSE { + cuda_device::library::call_v1::( + transpose_store_partial, + (output, values, valid), + ); + } else { + cuda_device::library::call_v1::( + direct_store_partial, + (output, values, valid), + ); + } + } else if TRANSPOSE { + cuda_device::library::call_v1::( + transpose_store_full, + (output, values), + ); + } else { + cuda_device::library::call_v1::( + direct_store_full, + (output, values), + ); + } + } +} diff --git a/rust/cuda-coop-codegen/tests/oxide-fixture/consumer/Cargo.toml b/rust/cuda-coop-codegen/tests/oxide-fixture/consumer/Cargo.toml new file mode 100644 index 000000000000..e2393ccec577 --- /dev/null +++ b/rust/cuda-coop-codegen/tests/oxide-fixture/consumer/Cargo.toml @@ -0,0 +1,15 @@ +# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. +# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +[package] +name = "coop-marker-fixture" +version = "0.0.0" +edition = "2024" +publish = false + +[dependencies] +cuda-device.workspace = true +cuda-host.workspace = true +cuda-core.workspace = true +oxide-artifacts.workspace = true +coop-marker-fixture-helper = { path = "../helper" } diff --git a/rust/cuda-coop-codegen/tests/oxide-fixture/consumer/src/main.rs b/rust/cuda-coop-codegen/tests/oxide-fixture/consumer/src/main.rs new file mode 100644 index 000000000000..eba29e10753e --- /dev/null +++ b/rust/cuda-coop-codegen/tests/oxide-fixture/consumer/src/main.rs @@ -0,0 +1,257 @@ +// Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +use cuda_device::{kernel, launch_bounds, launch_contract}; +use cuda_host::cuda_module; + +#[cuda_module] +mod kernels { + use super::*; + #[kernel] + #[launch_bounds(256)] + #[launch_contract(domain = 3, dynamic_shared = inferred)] + pub unsafe fn direct_1( + input: *const i32, + output: *mut i32, + post: *mut i32, + valid: i32, + mode: u32, + ) { + unsafe { + coop_marker_fixture_helper::roundtrip::<1, false>(input, output, post, valid, mode); + } + } + #[kernel] + #[launch_bounds(256)] + #[launch_contract(domain = 3, dynamic_shared = inferred)] + pub unsafe fn direct_4( + input: *const i32, + output: *mut i32, + post: *mut i32, + valid: i32, + mode: u32, + ) { + unsafe { + coop_marker_fixture_helper::roundtrip::<4, false>(input, output, post, valid, mode); + } + } + #[kernel] + #[launch_bounds(256)] + #[launch_contract(domain = 3, dynamic_shared = inferred)] + pub unsafe fn transpose_1( + input: *const i32, + output: *mut i32, + post: *mut i32, + valid: i32, + mode: u32, + ) { + unsafe { + coop_marker_fixture_helper::roundtrip::<1, true>(input, output, post, valid, mode); + } + } + #[kernel] + #[launch_bounds(256)] + #[launch_contract(domain = 3, dynamic_shared = inferred)] + pub unsafe fn transpose_4( + input: *const i32, + output: *mut i32, + post: *mut i32, + valid: i32, + mode: u32, + ) { + unsafe { + coop_marker_fixture_helper::roundtrip::<4, true>(input, output, post, valid, mode); + } + } + #[kernel] + #[launch_bounds(256)] + #[launch_contract(domain = 3, dynamic_shared = inferred)] + pub unsafe fn shared_helper_root( + input: *const i32, + output: *mut i32, + post: *mut i32, + valid: i32, + mode: u32, + ) { + unsafe { + coop_marker_fixture_helper::roundtrip::<4, true>(input, output, post, valid, mode); + } + } +} + +fn main() -> Result<(), Box> { + use cuda_core::{CudaContext, DeviceBuffer}; + use cuda_core::{LaunchConfig3D, LaunchContractError}; + use oxide_artifacts::{ArtifactPayloadKind, read_artifact_bundles_from_object_bytes}; + let executable = std::fs::read(std::env::current_exe()?)?; + let bundles = read_artifact_bundles_from_object_bytes(&executable)?; + let bundle = bundles + .iter() + .find(|b| b.payload(ArtifactPayloadKind::Cubin).is_some()) + .ok_or("missing embedded native artifact")?; + let cubin = bundle.payload(ArtifactPayloadKind::Cubin).unwrap(); + std::fs::write("fixture.cubin", cubin)?; + let context = CudaContext::new(0)?; + let stream = context.default_stream(); + // SAFETY: the embedded artifact is generated from these kernel declarations. + let module = unsafe { kernels::load(&context)? }; + macro_rules! check_plan { + ($prepare:ident, $bytes:expr, $alignment:expr) => {{ + let prepared = module.$prepare(LaunchConfig3D::auto_shared((1, 1, 1), (128, 1, 1)))?; + let facts = prepared.function().launch_facts().unwrap(); + assert_eq!(facts.exact_block, Some((128, 1, 1))); + assert_eq!(facts.dynamic_shared_bytes, $bytes); + assert_eq!(facts.dynamic_shared_alignment, $alignment); + assert_eq!(prepared.__raw_config().shared_mem_bytes, $bytes); + // Equal thread products do not imply the required exact shape. + for shape in [(64, 2, 1), (64, 1, 1)] { + assert!(matches!( + module.$prepare(LaunchConfig3D::auto_shared((1, 1, 1), shape)), + Err(LaunchContractError::BlockShapeMismatch { .. }) + )); + } + let required: u32 = $bytes; + for bytes in [required.saturating_sub(1), required + 1] { + if bytes != required { + assert!(matches!( + module.$prepare(LaunchConfig3D::new((1, 1, 1), (128, 1, 1), bytes)), + Err(LaunchContractError::DynamicSharedMemoryExactMismatch { .. }) + )); + } + } + }}; + } + check_plan!(prepare_direct_1, 0, 1); + check_plan!(prepare_direct_4, 0, 1); + check_plan!(prepare_transpose_1, 528, 16); + check_plan!(prepare_transpose_4, 2064, 16); + check_plan!(prepare_shared_helper_root, 2064, 16); + let mut cases = 0; + for (kernel, items, transpose) in [ + (0, 1, false), + (1, 4, false), + (2, 1, true), + (3, 4, true), + (4, 4, true), + ] { + let tile = 128 * items; + for blocks in [1, 3] { + let extent = tile * blocks; + let input: Vec = (0..extent).map(|i| i as i32 * 3 - 11).collect(); + let device_input = DeviceBuffer::from_host(&stream, &input)?; + for valid in [0, 1, tile - 1, tile, extent - 1, extent] { + for mode in 0..8_u32 { + let output = DeviceBuffer::from_host(&stream, &vec![-99_i32; extent])?; + let post = DeviceBuffer::from_host(&stream, &vec![-99_i32; extent])?; + let config = LaunchConfig3D::auto_shared((blocks as u32, 1, 1), (128, 1, 1)); + let input_ptr = device_input.cu_deviceptr() as *const i32; + let output_ptr = output.cu_deviceptr() as *mut i32; + let post_ptr = post.cu_deviceptr() as *mut i32; + unsafe { + match kernel { + 0 => module.direct_1( + &stream, + &module.prepare_direct_1(config)?, + input_ptr, + output_ptr, + post_ptr, + valid as i32, + mode, + )?, + 1 => module.direct_4( + &stream, + &module.prepare_direct_4(config)?, + input_ptr, + output_ptr, + post_ptr, + valid as i32, + mode, + )?, + 2 => module.transpose_1( + &stream, + &module.prepare_transpose_1(config)?, + input_ptr, + output_ptr, + post_ptr, + valid as i32, + mode, + )?, + 3 => module.transpose_4( + &stream, + &module.prepare_transpose_4(config)?, + input_ptr, + output_ptr, + post_ptr, + valid as i32, + mode, + )?, + _ => module.shared_helper_root( + &stream, + &module.prepare_shared_helper_root(config)?, + input_ptr, + output_ptr, + post_ptr, + valid as i32, + mode, + )?, + } + } + let loaded: Vec = input + .iter() + .enumerate() + .map(|(i, value)| { + (if mode & 1 != 0 && i >= valid { + -7 + } else { + *value + }) + 1 + }) + .collect(); + let expected_output: Vec = loaded + .iter() + .enumerate() + .map(|(i, value)| { + if mode & 2 != 0 && i >= valid { + -99 + } else { + *value + } + }) + .collect(); + let expected_post: Vec = (0..extent) + .map(|index| { + if mode & 4 != 0 { + return expected_output[index]; + } + let origin = index / tile * tile; + let lane = (index - origin) / items; + let item = index % items; + loaded[if transpose { + origin + item * 128 + lane + } else { + index + }] + }) + .collect(); + assert_eq!( + output.to_host_vec(&stream)?, + expected_output, + "output kernel={kernel} blocks={blocks} valid={valid} mode={mode}" + ); + assert_eq!( + post.to_host_vec(&stream)?, + expected_post, + "post-store kernel={kernel} blocks={blocks} valid={valid} mode={mode}" + ); + cases += 1; + } + } + } + } + assert_eq!(cases, 480); + println!( + "PASS {cases} native marker GPU cases with inferred scratch through five roots and sixteen provider specializations; {} CUBIN bytes", + cubin.len() + ); + Ok(()) +} diff --git a/rust/cuda-coop-codegen/tests/oxide-fixture/helper/Cargo.toml b/rust/cuda-coop-codegen/tests/oxide-fixture/helper/Cargo.toml new file mode 100644 index 000000000000..a997768dc07c --- /dev/null +++ b/rust/cuda-coop-codegen/tests/oxide-fixture/helper/Cargo.toml @@ -0,0 +1,12 @@ +# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. +# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +[package] +name = "coop-marker-fixture-helper" +version = "0.0.0" +edition = "2024" +publish = false + +[dependencies] +cuda-device.workspace = true +coop-marker-fixture-api = { path = "../api" } diff --git a/rust/cuda-coop-codegen/tests/oxide-fixture/helper/src/lib.rs b/rust/cuda-coop-codegen/tests/oxide-fixture/helper/src/lib.rs new file mode 100644 index 000000000000..025d76c5fb6d --- /dev/null +++ b/rust/cuda-coop-codegen/tests/oxide-fixture/helper/src/lib.rs @@ -0,0 +1,65 @@ +// Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +#![no_std] + +/// Roundtrip through one shared generic helper reached from several kernels. +/// +/// # Safety +/// Each block owns a disjoint tile. Launch exact 128-thread one-dimensional +/// blocks with sufficient scratch. Input/output cover every full tile; partial +/// counts are a global prefix within their combined extent. Every thread in +/// a block receives equal input/output origins, valid count and mode. The fixed +/// fill value is uniform, and each thread owns its separate initialized payload. +#[inline(never)] +pub unsafe fn roundtrip( + input: *const i32, + output: *mut i32, + post_store: *mut i32, + valid: i32, + mode: u32, +) { + let lane = cuda_device::thread::threadIdx_x() as usize; + let origin = cuda_device::thread::blockIdx_x() as usize * 128 * ITEMS_PER_THREAD; + let valid = (valid - origin as i32) + .max(0) + .min((128 * ITEMS_PER_THREAD) as i32); + let mut payload = [0_i32; ITEMS_PER_THREAD]; + unsafe { + coop_marker_fixture_api::load::( + input.add(origin), + payload.as_mut_ptr(), + valid, + mode & 1 != 0, + ); + cuda_device::thread::sync_threads(); + for item in &mut payload { + *item += 1; + } + coop_marker_fixture_api::store::( + output.add(origin), + payload.as_mut_ptr(), + valid, + mode & 2 != 0, + ); + for (i, item) in payload.iter().enumerate() { + post_store + .add(origin + lane * ITEMS_PER_THREAD + i) + .write(*item); + } + cuda_device::thread::sync_threads(); + if mode & 4 != 0 { + coop_marker_fixture_api::load::( + output.add(origin), + payload.as_mut_ptr(), + valid, + false, + ); + for (i, item) in payload.iter().enumerate() { + post_store + .add(origin + lane * ITEMS_PER_THREAD + i) + .write(*item); + } + } + } +} diff --git a/rust/cuda-coop-codegen/tests/run_oxide_fixture.py b/rust/cuda-coop-codegen/tests/run_oxide_fixture.py new file mode 100644 index 000000000000..6b357575759d --- /dev/null +++ b/rust/cuda-coop-codegen/tests/run_oxide_fixture.py @@ -0,0 +1,83 @@ +# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. +# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +"""Run coordinated source integration with an explicitly supplied Oxide checkout.""" + +import argparse +import json +import os +import shutil +import subprocess +from pathlib import Path + +parser = argparse.ArgumentParser(description=__doc__) +parser.add_argument("--oxide-root", required=True, type=Path) +parser.add_argument("--cargo-oxide", required=True, type=Path) +parser.add_argument("--backend", required=True, type=Path) +parser.add_argument("--output", required=True, type=Path) +parser.add_argument("--cuda-root", type=Path, default=Path("/usr/local/cuda")) +parser.add_argument("--arch", default="sm_120") +args = parser.parse_args() +source = Path(__file__).resolve().parent +checkout = source.parents[2] +output = args.output.resolve() +output.mkdir(parents=True, exist_ok=False) +fixture = output / "fixture" +shutil.copytree(source / "oxide-fixture", fixture) +# Only the generated copy receives cross-checkout paths. Ordinary source files +# and the reviewed manifests remain unchanged. +api_manifest = fixture / "api/Cargo.toml" +api_manifest.write_text( + api_manifest.read_text().replace( + 'path = "../../.."', f"path = {json.dumps(str(source.parent))}" + ) +) +oxide = args.oxide_root.resolve() +manifest = fixture / "Cargo.toml" +with manifest.open("a") as stream: + stream.write("\n[patch.crates-io]\n") + for name in ["cuda-device", "cuda-host", "cuda-macros", "oxide-artifacts"]: + stream.write( + f"{name} = {{ path = {json.dumps(str(oxide / 'cuda-oxide/crates' / name))} }}\n" + ) + for name in ["cuda-core", "cuda-core-derive", "cuda-bindings"]: + stream.write(f"{name} = {{ path = {json.dumps(str(oxide / name))} }}\n") +# Reuse exact same-base dependency resolutions; Cargo adds only fixture/provider +# packages. Do not refresh every example's lockfile for this integration test. +shutil.copyfile( + oxide / "cuda-oxide/crates/rustc-codegen-cuda/examples/vecadd/Cargo.lock", + fixture / "Cargo.lock", +) +env = os.environ.copy() +env.update( + CCCL_SOURCE=str(checkout), + CUDA_HOME=str(args.cuda_root.resolve()), + CUDA_PATH=str(args.cuda_root.resolve()), + CUDA_OXIDE_BACKEND=str(args.backend.resolve()), +) +command = [ + str(args.cargo_oxide.resolve()), + "oxide", + "run", + "--arch", + args.arch, + "--materialize-cubin", +] +with (output / "build-run.log").open("w") as log: + result = subprocess.run( + command, cwd=fixture / "consumer", env=env, stdout=log, stderr=subprocess.STDOUT + ) +(output / "command.json").write_text( + json.dumps( + { + "argv": command, + "cwd": str(fixture / "consumer"), + "oxide_root": str(oxide), + "backend": env["CUDA_OXIDE_BACKEND"], + "returncode": result.returncode, + }, + indent=2, + ) + + "\n" +) +raise SystemExit(result.returncode) diff --git a/rust/rust-toolchain.toml b/rust/rust-toolchain.toml new file mode 100644 index 000000000000..06d6cb6baf70 --- /dev/null +++ b/rust/rust-toolchain.toml @@ -0,0 +1,6 @@ +# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. +# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception +[toolchain] +channel = "nightly-2026-10-11" +components = ["clippy", "rustfmt"] +profile = "minimal"