From d1d7e5f87b3e4c89ed5a222f4f787f017c6d701b Mon Sep 17 00:00:00 2001 From: Trent Nelson Date: Mon, 19 Oct 2026 21:39:36 -0700 Subject: [PATCH 01/33] Add native Rust provider for CUB block load and store Create a host-only Cargo workspace that specializes selected CUB forms through NVRTC and reports exact scratch layouts with their LTO output. Bind specialization identities to source, headers and toolkit bytes. Include a runnable provider example, Rust guide and CPU checks. Signed-off-by: Trent Nelson --- .github/workflows/rust-host.yml | 43 ++ docs/index.rst | 3 + docs/rust/index.rst | 109 +++++ rust/.gitignore | 1 + rust/AGENTS.md | 63 +++ rust/Cargo.lock | 110 +++++ rust/Cargo.toml | 18 + rust/cuda-coop-codegen/Cargo.toml | 17 + rust/cuda-coop-codegen/examples/specialize.rs | 35 ++ rust/cuda-coop-codegen/src/headers.rs | 86 ++++ rust/cuda-coop-codegen/src/lib.rs | 462 ++++++++++++++++++ rust/cuda-coop-codegen/src/nvrtc.rs | 302 ++++++++++++ rust/cuda-coop-codegen/tests/native.rs | 133 +++++ rust/rust-toolchain.toml | 6 + 14 files changed, 1388 insertions(+) create mode 100644 .github/workflows/rust-host.yml create mode 100644 docs/rust/index.rst create mode 100644 rust/.gitignore create mode 100644 rust/AGENTS.md create mode 100644 rust/Cargo.lock create mode 100644 rust/Cargo.toml create mode 100644 rust/cuda-coop-codegen/Cargo.toml create mode 100644 rust/cuda-coop-codegen/examples/specialize.rs create mode 100644 rust/cuda-coop-codegen/src/headers.rs create mode 100644 rust/cuda-coop-codegen/src/lib.rs create mode 100644 rust/cuda-coop-codegen/src/nvrtc.rs create mode 100644 rust/cuda-coop-codegen/tests/native.rs create mode 100644 rust/rust-toolchain.toml diff --git a/.github/workflows/rust-host.yml b/.github/workflows/rust-host.yml new file mode 100644 index 000000000000..d5bde6e27fa7 --- /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 + 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..23776f9573fc --- /dev/null +++ b/docs/rust/index.rst @@ -0,0 +1,109 @@ +.. 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. 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..fdc6db4041d8 --- /dev/null +++ b/rust/Cargo.lock @@ -0,0 +1,110 @@ +# 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", + "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 = "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 = "sha2" +version = "0.10.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a7507d819769d01a365ab707794a4084392c824f54a7a6a7862f8c3d0892b283" +dependencies = [ + "cfg-if", + "cpufeatures", + "digest", +] + +[[package]] +name = "typenum" +version = "1.20.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b6f5e870be6c3b371b77fe0ee0bafb859fa4964b4404c27de1d380043c4dda20" + +[[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" 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..7476b2c60391 --- /dev/null +++ b/rust/cuda-coop-codegen/Cargo.toml @@ -0,0 +1,17 @@ +# 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" +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..7b9105133221 --- /dev/null +++ b/rust/cuda-coop-codegen/src/headers.rs @@ -0,0 +1,86 @@ +// 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}; + +pub(crate) struct Headers { + pub(crate) entries: Vec<(CString, CString)>, + pub(crate) digest: [u8; 32], +} + +impl Headers { + 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 [ + (checkout.join("cub/cub"), "cub/"), + (checkout.join("libcudacxx/include"), ""), + (checkout.join("thrust/thrust"), "thrust/"), + (toolkit.join("include"), ""), + ] { + 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..298510237eae --- /dev/null +++ b/rust/cuda-coop-codegen/src/lib.rs @@ -0,0 +1,462 @@ +// 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; + +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..52daf38a400d --- /dev/null +++ b/rust/cuda-coop-codegen/src/nvrtc.rs @@ -0,0 +1,302 @@ +// 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, + 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 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 = |name: &str| -> Result { + let path = toolkit + .join("lib64") + .join(name) + .canonicalize() + .map_err(|e| Error::input(&toolkit.join("lib64").join(name), e))?; + 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 = load("libnvrtc-builtins.so")?; + let library = load("libnvrtc.so")?; + // 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/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/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" From 19060fb569299480ffcd91ca86848b8f3605447b Mon Sep 17 00:00:00 2001 From: Trent Nelson Date: Mon, 19 Oct 2026 21:48:44 -0700 Subject: [PATCH 02/33] Keep the Rust CI toolchain installation explicit Install the pinned toolchain without updating rustup itself. Signed-off-by: Trent Nelson --- .github/workflows/rust-host.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/rust-host.yml b/.github/workflows/rust-host.yml index d5bde6e27fa7..ad93cbf4b65b 100644 --- a/.github/workflows/rust-host.yml +++ b/.github/workflows/rust-host.yml @@ -35,7 +35,7 @@ jobs: persist-credentials: false - name: Check the pinned Rust workspace run: | - rustup toolchain install + rustup toolchain install --no-self-update rustup show active-toolchain cargo fmt --all -- --check cargo test --workspace --locked From bcec7aee462ccdf650f8843d07befcf3cc8dee25 Mon Sep 17 00:00:00 2001 From: Trent Nelson Date: Mon, 19 Oct 2026 22:08:08 -0700 Subject: [PATCH 03/33] Connect native cooperative providers to Oxide library calls Serve native specialization batches through the versioned Oxide boundary and exercise Load and Store through a standalone Cargo consumer with a generic helper dependency. Signed-off-by: Trent Nelson --- rust/Cargo.lock | 98 +++++ rust/cuda-coop-codegen/Cargo.toml | 2 + rust/cuda-coop-codegen/src/lib.rs | 1 + rust/cuda-coop-codegen/src/oxide.rs | 355 ++++++++++++++++++ .../tests/oxide-fixture/Cargo.toml | 13 + .../tests/oxide-fixture/README.md | 37 ++ .../tests/oxide-fixture/api/Cargo.toml | 14 + .../tests/oxide-fixture/api/build.rs | 35 ++ .../tests/oxide-fixture/api/src/lib.rs | 95 +++++ .../tests/oxide-fixture/consumer/Cargo.toml | 15 + .../tests/oxide-fixture/consumer/src/main.rs | 197 ++++++++++ .../tests/oxide-fixture/helper/Cargo.toml | 12 + .../tests/oxide-fixture/helper/src/lib.rs | 65 ++++ .../tests/run_oxide_fixture.py | 83 ++++ 14 files changed, 1022 insertions(+) create mode 100644 rust/cuda-coop-codegen/src/oxide.rs create mode 100644 rust/cuda-coop-codegen/tests/oxide-fixture/Cargo.toml create mode 100644 rust/cuda-coop-codegen/tests/oxide-fixture/README.md create mode 100644 rust/cuda-coop-codegen/tests/oxide-fixture/api/Cargo.toml create mode 100644 rust/cuda-coop-codegen/tests/oxide-fixture/api/build.rs create mode 100644 rust/cuda-coop-codegen/tests/oxide-fixture/api/src/lib.rs create mode 100644 rust/cuda-coop-codegen/tests/oxide-fixture/consumer/Cargo.toml create mode 100644 rust/cuda-coop-codegen/tests/oxide-fixture/consumer/src/main.rs create mode 100644 rust/cuda-coop-codegen/tests/oxide-fixture/helper/Cargo.toml create mode 100644 rust/cuda-coop-codegen/tests/oxide-fixture/helper/src/lib.rs create mode 100644 rust/cuda-coop-codegen/tests/run_oxide_fixture.py diff --git a/rust/Cargo.lock b/rust/Cargo.lock index fdc6db4041d8..b1c2103ebcbe 100644 --- a/rust/Cargo.lock +++ b/rust/Cargo.lock @@ -41,6 +41,8 @@ name = "cuda-coop-codegen" version = "0.0.0" dependencies = [ "libloading", + "serde", + "serde_json", "sha2", ] @@ -64,6 +66,12 @@ dependencies = [ "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" @@ -80,6 +88,73 @@ dependencies = [ "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" @@ -91,12 +166,29 @@ dependencies = [ "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" @@ -108,3 +200,9 @@ 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/cuda-coop-codegen/Cargo.toml b/rust/cuda-coop-codegen/Cargo.toml index 7476b2c60391..7a92fcaaded4 100644 --- a/rust/cuda-coop-codegen/Cargo.toml +++ b/rust/cuda-coop-codegen/Cargo.toml @@ -11,6 +11,8 @@ publish = false [dependencies] libloading = "0.8" +serde = { version = "1", features = ["derive"] } +serde_json = "1" sha2 = "0.10" [lints] diff --git a/rust/cuda-coop-codegen/src/lib.rs b/rust/cuda-coop-codegen/src/lib.rs index 298510237eae..872adf16c76a 100644 --- a/rust/cuda-coop-codegen/src/lib.rs +++ b/rust/cuda-coop-codegen/src/lib.rs @@ -48,6 +48,7 @@ mod headers; #[allow(unsafe_code)] mod nvrtc; +pub mod oxide; use headers::{Headers, hash_field}; use sha2::{Digest, Sha256}; diff --git a/rust/cuda-coop-codegen/src/oxide.rs b/rust/cuda-coop-codegen/src/oxide.rs new file mode 100644 index 000000000000..09d57e186b6b --- /dev/null +++ b/rust/cuda-coop-codegen/src/oxide.rs @@ -0,0 +1,355 @@ +// 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`]. +//! 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}"))) +} + +/// 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 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/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..fc123a55f9d3 --- /dev/null +++ b/rust/cuda-coop-codegen/tests/oxide-fixture/README.md @@ -0,0 +1,37 @@ +# 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 current fixture uses explicit oversized scratch while O02 launch-facts +integration is pending. 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..76b2e828a01b --- /dev/null +++ b/rust/cuda-coop-codegen/tests/oxide-fixture/api/build.rs @@ -0,0 +1,35 @@ +// 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() { + if std::env::args().nth(1).as_deref() == Some("--oxide-library-request-v1") { + if let Err(error) = oxide::serve( + std::path::Path::new(env!("CCCL_SOURCE")), + std::io::stdin().lock(), + std::io::stdout().lock(), + ) { + 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..5f555bfd6953 --- /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 { + env!($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..ccf454ce08ac --- /dev/null +++ b/rust/cuda-coop-codegen/tests/oxide-fixture/consumer/src/main.rs @@ -0,0 +1,197 @@ +// Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +use cuda_device::kernel; +use cuda_host::cuda_module; + +#[cuda_module] +mod kernels { + use super::*; + #[kernel] + pub 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] + pub 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] + pub 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] + pub 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] + pub 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::simt::LaunchConfig; + use cuda_core::{CudaContext, DeviceBuffer}; + 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(); + let module = kernels::load(&context)?; + 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])?; + // Internal ABI qualification only. O02 integration replaces + // this fixed oversized allocation with finalized root facts. + let config = LaunchConfig { + grid_dim: (blocks as u32, 1, 1), + block_dim: (128, 1, 1), + shared_mem_bytes: 4096, + }; + 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, + config, + input_ptr, + output_ptr, + post_ptr, + valid as i32, + mode, + )?, + 1 => module.direct_4( + &stream, + config, + input_ptr, + output_ptr, + post_ptr, + valid as i32, + mode, + )?, + 2 => module.transpose_1( + &stream, + config, + input_ptr, + output_ptr, + post_ptr, + valid as i32, + mode, + )?, + 3 => module.transpose_4( + &stream, + config, + input_ptr, + output_ptr, + post_ptr, + valid as i32, + mode, + )?, + _ => module.shared_helper_root( + &stream, + 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 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) From 37ab28836f07213fae9e0040ba1ef3f441b976d6 Mon Sep 17 00:00:00 2001 From: Trent Nelson Date: Mon, 19 Oct 2026 22:19:14 -0700 Subject: [PATCH 04/33] Use inferred resources in the provider integration fixture --- .../tests/oxide-fixture/README.md | 5 +- .../tests/oxide-fixture/consumer/src/main.rs | 102 ++++++++++++++---- 2 files changed, 84 insertions(+), 23 deletions(-) diff --git a/rust/cuda-coop-codegen/tests/oxide-fixture/README.md b/rust/cuda-coop-codegen/tests/oxide-fixture/README.md index fc123a55f9d3..fb76581bf752 100644 --- a/rust/cuda-coop-codegen/tests/oxide-fixture/README.md +++ b/rust/cuda-coop-codegen/tests/oxide-fixture/README.md @@ -29,8 +29,9 @@ 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 current fixture uses explicit oversized scratch while O02 launch-facts -integration is pending. Its helpers deliberately retain function boundaries to +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 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 index ccf454ce08ac..eba29e10753e 100644 --- a/rust/cuda-coop-codegen/tests/oxide-fixture/consumer/src/main.rs +++ b/rust/cuda-coop-codegen/tests/oxide-fixture/consumer/src/main.rs @@ -1,38 +1,72 @@ // Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception -use cuda_device::kernel; +use cuda_device::{kernel, launch_bounds, launch_contract}; use cuda_host::cuda_module; #[cuda_module] mod kernels { use super::*; #[kernel] - pub fn direct_1(input: *const i32, output: *mut i32, post: *mut i32, valid: i32, mode: u32) { + #[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] - pub fn direct_4(input: *const i32, output: *mut i32, post: *mut i32, valid: i32, mode: u32) { + #[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] - pub fn transpose_1(input: *const i32, output: *mut i32, post: *mut i32, valid: i32, mode: u32) { + #[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] - pub fn transpose_4(input: *const i32, output: *mut i32, post: *mut i32, valid: i32, mode: u32) { + #[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] - pub fn shared_helper_root( + #[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, @@ -46,8 +80,8 @@ mod kernels { } fn main() -> Result<(), Box> { - use cuda_core::simt::LaunchConfig; 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)?; @@ -59,7 +93,39 @@ fn main() -> Result<(), Box> { std::fs::write("fixture.cubin", cubin)?; let context = CudaContext::new(0)?; let stream = context.default_stream(); - let module = kernels::load(&context)?; + // 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), @@ -77,13 +143,7 @@ fn main() -> Result<(), Box> { 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])?; - // Internal ABI qualification only. O02 integration replaces - // this fixed oversized allocation with finalized root facts. - let config = LaunchConfig { - grid_dim: (blocks as u32, 1, 1), - block_dim: (128, 1, 1), - shared_mem_bytes: 4096, - }; + 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; @@ -91,7 +151,7 @@ fn main() -> Result<(), Box> { match kernel { 0 => module.direct_1( &stream, - config, + &module.prepare_direct_1(config)?, input_ptr, output_ptr, post_ptr, @@ -100,7 +160,7 @@ fn main() -> Result<(), Box> { )?, 1 => module.direct_4( &stream, - config, + &module.prepare_direct_4(config)?, input_ptr, output_ptr, post_ptr, @@ -109,7 +169,7 @@ fn main() -> Result<(), Box> { )?, 2 => module.transpose_1( &stream, - config, + &module.prepare_transpose_1(config)?, input_ptr, output_ptr, post_ptr, @@ -118,7 +178,7 @@ fn main() -> Result<(), Box> { )?, 3 => module.transpose_4( &stream, - config, + &module.prepare_transpose_4(config)?, input_ptr, output_ptr, post_ptr, @@ -127,7 +187,7 @@ fn main() -> Result<(), Box> { )?, _ => module.shared_helper_root( &stream, - config, + &module.prepare_shared_helper_root(config)?, input_ptr, output_ptr, post_ptr, @@ -190,7 +250,7 @@ fn main() -> Result<(), Box> { } assert_eq!(cases, 480); println!( - "PASS {cases} native marker GPU cases through five roots and sixteen provider specializations; {} CUBIN bytes", + "PASS {cases} native marker GPU cases with inferred scratch through five roots and sixteen provider specializations; {} CUBIN bytes", cubin.len() ); Ok(()) From 39ec3f3ab686a394021b78b9bfabdad53bc25ec1 Mon Sep 17 00:00:00 2001 From: Trent Nelson Date: Mon, 19 Oct 2026 22:34:29 -0700 Subject: [PATCH 05/33] Add checked per-thread Rust payloads --- docs/rust/index.rst | 11 +- docs/rust/payload.rst | 77 ++++++++ rust/Cargo.lock | 4 + rust/Cargo.toml | 2 +- rust/cuda-coop/Cargo.toml | 13 ++ rust/cuda-coop/examples/payload.rs | 22 +++ rust/cuda-coop/src/alignment.rs | 65 +++++++ rust/cuda-coop/src/lib.rs | 40 ++++ rust/cuda-coop/src/thread_data.rs | 301 +++++++++++++++++++++++++++++ rust/cuda-coop/tests/payload.rs | 148 ++++++++++++++ 10 files changed, 680 insertions(+), 3 deletions(-) create mode 100644 docs/rust/payload.rst create mode 100644 rust/cuda-coop/Cargo.toml create mode 100644 rust/cuda-coop/examples/payload.rs create mode 100644 rust/cuda-coop/src/alignment.rs create mode 100644 rust/cuda-coop/src/lib.rs create mode 100644 rust/cuda-coop/src/thread_data.rs create mode 100644 rust/cuda-coop/tests/payload.rs diff --git a/docs/rust/index.rst b/docs/rust/index.rst index 23776f9573fc..72be476dcb99 100644 --- a/docs/rust/index.rst +++ b/docs/rust/index.rst @@ -4,8 +4,10 @@ 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 +The Rust workspace contains per-thread payload types and the native host compiler +foundation for ``cuda.coop``. The ``cuda-coop`` crate supplies checked +``ThreadData`` values; see :doc:`payload`. 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. @@ -107,3 +109,8 @@ mutation, and collective requirements. Open ``rust/target/doc/cuda_coop_codegen/ 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. + +.. toctree:: + :maxdepth: 1 + + payload diff --git a/docs/rust/payload.rst b/docs/rust/payload.rst new file mode 100644 index 000000000000..ae2d0f55d26c --- /dev/null +++ b/docs/rust/payload.rst @@ -0,0 +1,77 @@ +.. Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. +.. SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +Per-thread payloads +=================== + +The unpublished ``cuda-coop`` crate provides ``ThreadData`` and its sealed +``ThreadDataLike`` interface. It is a feature-free ``no_std`` crate with no +dependencies, and supports Rust 1.85 or newer. Oxide is the intended device +compiler integration. This crate currently supplies the payload vocabulary; +collectives, groups and scratch storage are still under development. + +Construct and access a payload +------------------------------ + +``ThreadData`` owns a fixed number +of per-thread slots. ``ITEMS`` must be positive. The element type ``T`` is inferred +from typed writes and preserves its width: the supported types are ``i8``, +``u8``, ``i16``, ``u16``, ``i32``, ``u32``, ``i64``, ``u64``, ``f32`` and ``f64``. + +New slots are unwritten. ``write(index, value)`` initializes or replaces exactly +one slot. ``get(index)`` returns a copied value in ``Some`` for initialized slots, +and ``None`` for unwritten or out-of-bounds slots. ``get_mut(index)`` applies the +same checks before returning an exclusive reference. + +Indexing reads or mutates an initialized slot. Reading or assigning through an +unwritten index panics; use ``write`` for first initialization. Out-of-bounds +indexing and writes also panic. ``len()`` counts all slots, including unwritten +ones. Copies, moves, swaps and conditional selection preserve the exact validity +of each slot. The payload exposes no unrestricted initialized slice. + +This complete host program uses a generic per-thread item count and an inferred +element type: + +.. literalinclude:: ../../rust/cuda-coop/examples/payload.rs + :language: rust + :start-after: //! Construct and inspect a per-thread payload on the host. + +Run it from the Rust workspace: + +.. code-block:: shell + + cargo run -p cuda-coop --example payload + +Generic payload helpers can use ``D: ThreadDataLike``, with ``D::Item`` for the +element type and ``D::EXTENT`` for the compile-time slot count. The trait's +``get`` and ``write`` methods preserve the same access rules. The trait is sealed +to the crate's supported payload implementations. + +Alignment and layout +-------------------- + +``ALIGN`` is a minimum alignment in bytes for the actual contiguous value +storage. Its default of 1 adds no requirement beyond the element type's natural +alignment. All powers of two from 1 through 2²⁹ are supported. Zero, non-powers of +two and larger values fail to compile; zero item counts fail when the constructor +is instantiated. Large alignments may increase object size through padding. + +Code with generic alignment uses the public bound +``Alignment: SupportedAlignment``. The representation keeps validity +metadata private and does not establish a public foreign-function ABI. Payload +alignment does not allocate scratch memory or establish collective participation +or synchronization. + +Validation and API reference +---------------------------- + +.. code-block:: shell + + cargo test -p cuda-coop + cargo clippy -p cuda-coop --all-targets -- -D warnings + RUSTDOCFLAGS="-D warnings" cargo doc -p cuda-coop --no-deps + +The tests cover checked access, arbitrary initialized slots, numeric inference, +copies and selection, value-array alignment, and rejected types, extents and +alignments. These host tests and doctests do not establish GPU collective safety. +Rustdoc at ``rust/target/doc/cuda_coop/index.html`` contains the full API reference. diff --git a/rust/Cargo.lock b/rust/Cargo.lock index fdc6db4041d8..650a0f45e4ce 100644 --- a/rust/Cargo.lock +++ b/rust/Cargo.lock @@ -36,6 +36,10 @@ dependencies = [ "typenum", ] +[[package]] +name = "cuda-coop" +version = "0.0.0" + [[package]] name = "cuda-coop-codegen" version = "0.0.0" diff --git a/rust/Cargo.toml b/rust/Cargo.toml index c90fafb1430a..c6fd608c3716 100644 --- a/rust/Cargo.toml +++ b/rust/Cargo.toml @@ -2,7 +2,7 @@ # SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception [workspace] -members = ["cuda-coop-codegen"] +members = ["cuda-coop", "cuda-coop-codegen"] resolver = "3" [workspace.package] diff --git a/rust/cuda-coop/Cargo.toml b/rust/cuda-coop/Cargo.toml new file mode 100644 index 000000000000..1ca43ea36394 --- /dev/null +++ b/rust/cuda-coop/Cargo.toml @@ -0,0 +1,13 @@ +# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. +# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +[package] +name = "cuda-coop" +version = "0.0.0" +edition.workspace = true +license.workspace = true +rust-version.workspace = true +publish = false + +[lints] +workspace = true diff --git a/rust/cuda-coop/examples/payload.rs b/rust/cuda-coop/examples/payload.rs new file mode 100644 index 000000000000..a7625da6c1ba --- /dev/null +++ b/rust/cuda-coop/examples/payload.rs @@ -0,0 +1,22 @@ +// Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +//! Construct and inspect a per-thread payload on the host. + +use cuda_coop as coop; + +fn payload(value: i16) -> coop::ThreadData { + let mut items = coop::ThreadData::<_, ITEMS_PER_THREAD>::new(); + items.write(0, value); + items +} + +fn main() { + let mut items = payload::<4>(7); + assert_eq!(items.get(0), Some(7_i16)); + assert_eq!(items.get(1), None); + items[0] += 2; + items.write(3, 15); + assert_eq!(items.get(0), Some(9)); + assert_eq!(items.get(3), Some(15)); +} diff --git a/rust/cuda-coop/src/alignment.rs b/rust/cuda-coop/src/alignment.rs new file mode 100644 index 000000000000..63e8f2987607 --- /dev/null +++ b/rust/cuda-coop/src/alignment.rs @@ -0,0 +1,65 @@ +// Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +/// Type-level byte alignment for [`crate::ThreadData`]. +/// +/// `BYTES` must be a power of two from 1 through 536,870,912 (2²⁹), inclusive. +/// These are Rust's supported explicit alignments. The payload's actual +/// alignment is at least this value and the element type's natural alignment. +/// Larger alignments can increase the payload's size through padding. +pub struct Alignment; + +/// A supported [`Alignment`] selector. +/// +/// This trait is sealed: implementations are supplied for all powers of two +/// from 1 through 2²⁹. Use this public bound when the alignment is generic: +/// +/// ``` +/// use cuda_coop as coop; +/// +/// fn payload(value: u32) +/// -> coop::ThreadData +/// where +/// coop::Alignment: coop::SupportedAlignment, +/// { +/// let mut items = coop::ThreadData::<_, ITEMS_PER_THREAD, ALIGN>::new(); +/// items.write(0, value); +/// items +/// } +/// +/// assert_eq!(payload::<4, 32>(11).get(0), Some(11)); +/// ``` +pub trait SupportedAlignment: sealed::Alignment {} + +pub(crate) mod sealed { + pub trait Alignment { + type Marker: Copy; + } + + macro_rules! alignments { + ($($marker:ident: $bytes:literal),* $(,)?) => { + $( + #[repr(align($bytes))] + #[derive(Clone, Copy)] + pub struct $marker; + + impl Alignment for super::Alignment<$bytes> { + type Marker = $marker; + } + + impl super::SupportedAlignment for super::Alignment<$bytes> {} + )* + }; + } + + alignments! { + A1: 1, A2: 2, A4: 4, A8: 8, A16: 16, A32: 32, + A64: 64, A128: 128, A256: 256, A512: 512, A1024: 1024, + A2048: 2048, A4096: 4096, A8192: 8192, A16384: 16384, + A32768: 32768, A65536: 65536, A131072: 131072, A262144: 262144, + A524288: 524288, A1048576: 1048576, A2097152: 2097152, + A4194304: 4194304, A8388608: 8388608, A16777216: 16777216, + A33554432: 33554432, A67108864: 67108864, A134217728: 134217728, + A268435456: 268435456, A536870912: 536870912, + } +} diff --git a/rust/cuda-coop/src/lib.rs b/rust/cuda-coop/src/lib.rs new file mode 100644 index 000000000000..ec97f008886b --- /dev/null +++ b/rust/cuda-coop/src/lib.rs @@ -0,0 +1,40 @@ +// Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +//! Checked per-thread payloads for Rust `cuda.coop`. +//! +//! This unpublished, feature-free `no_std` crate provides [`ThreadData`] and +//! its sealed [`ThreadDataLike`] interface. Each slot starts unwritten. +//! [`ThreadData::write`] initializes one slot; checked access exposes only +//! initialized values. Copies and moves preserve each slot's validity. +//! +//! ``` +//! use cuda_coop as coop; +//! +//! fn payload(value: i16) +//! -> coop::ThreadData +//! { +//! let mut items = coop::ThreadData::<_, ITEMS_PER_THREAD>::new(); +//! items.write(0, value); // Infers i16 without changing its width. +//! items +//! } +//! +//! let mut items = payload::<4>(7); +//! assert_eq!(items.get(0), Some(7)); +//! assert_eq!(items.get(1), None); +//! items[0] += 2; // Indexing can mutate an already initialized slot. +//! assert_eq!(items[0], 9); +//! ``` +//! +//! The crate requires Rust 1.85 or newer and has no host compiler or CUDA +//! dependencies. Oxide is the intended device compiler integration. This +//! payload API supplies no collective operations, group descriptors, scratch +//! storage, or synchronization guarantees. + +#![no_std] + +mod alignment; +mod thread_data; + +pub use alignment::{Alignment, SupportedAlignment}; +pub use thread_data::{Numeric, ThreadData, ThreadDataLike}; diff --git a/rust/cuda-coop/src/thread_data.rs b/rust/cuda-coop/src/thread_data.rs new file mode 100644 index 000000000000..4d955cd2cbb2 --- /dev/null +++ b/rust/cuda-coop/src/thread_data.rs @@ -0,0 +1,301 @@ +// Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +use core::mem::MaybeUninit; +use core::ops::{Index, IndexMut}; + +use crate::{Alignment, SupportedAlignment}; + +/// A numeric element supported by [`ThreadData`]. +/// +/// This sealed trait is implemented for exactly `i8`, `u8`, `i16`, `u16`, `i32`, +/// `u32`, `i64`, `u64`, `f32`, and `f64`. Values retain their Rust type and width. +pub trait Numeric: sealed::Numeric + Copy {} + +/// An owned, fixed-size payload whose slots are initialized independently. +/// +/// The complete signature is `ThreadData`, where `T: Numeric` and +/// `Alignment: SupportedAlignment`. `ITEMS` must be positive. `T` can be +/// inferred from typed writes; [`new`](Self::new) does not initialize values. +/// `ALIGN` specifies the minimum byte alignment of the actual element storage, +/// in addition to `T`'s natural alignment. The default of 1 adds no alignment +/// requirement. See [`Alignment`] for supported values. +/// +/// [`write`](Self::write) initializes or replaces a single slot. +/// [`get`](Self::get) and [`get_mut`](Self::get_mut) return `None` for an unwritten +/// or out-of-bounds slot. Indexing requires an initialized, in-bounds slot and +/// panics otherwise; use `write`, not indexed assignment, for initialization. +/// [`len`](Self::len) counts all slots, including unwritten ones. Copies, moves, +/// and selections carry the values and each slot's validity together. +/// +/// The values are contiguous, with private validity metadata. There is no +/// unrestricted initialized slice or public foreign-function layout contract. +/// This per-thread value performs no allocation or synchronization. +/// +/// ``` +/// use cuda_coop as coop; +/// let mut items = coop::ThreadData::<_, 4>::new(); +/// items.write(3, 42_u64); +/// assert_eq!(items.get(0), None); +/// assert_eq!(items.get(3), Some(42_u64)); +/// assert_eq!(items.len(), 4); +/// let copy = items; +/// assert_eq!(copy.get(3), Some(42_u64)); +/// ``` +/// +/// Unsupported element types and alignments fail to compile: +/// +/// ```compile_fail,E0277 +/// let _ = cuda_coop::ThreadData::::new(); +/// ``` +/// +/// ```compile_fail,E0277 +/// let _ = cuda_coop::ThreadData::::new(); +/// ``` +/// +/// ```compile_fail,E0277 +/// let _ = cuda_coop::ThreadData::::new(); +/// ``` +/// +/// ```compile_fail,E0277 +/// let _ = cuda_coop::ThreadData::::new(); +/// ``` +#[repr(C)] +#[derive(Clone, Copy)] +pub struct ThreadData +where + Alignment: SupportedAlignment, +{ + // In repr(C), this zero-length field aligns the following values array + // without moving it away from offset zero. + alignment: [ as crate::alignment::sealed::Alignment>::Marker; 0], + values: [MaybeUninit; ITEMS], + initialized: [u8; ITEMS], +} + +impl ThreadData +where + Alignment: SupportedAlignment, +{ + const POSITIVE_EXTENT: () = assert!(ITEMS > 0, "ThreadData requires at least one item"); + + /// Construct a payload with every slot unwritten. + /// + /// Instantiating this constructor with zero items fails constant evaluation: + /// + /// ```compile_fail,E0080 + /// const EMPTY: cuda_coop::ThreadData = cuda_coop::ThreadData::new(); + /// ``` + #[inline] + pub const fn new() -> Self { + let () = Self::POSITIVE_EXTENT; + Self { + alignment: [], + values: [MaybeUninit::uninit(); ITEMS], + initialized: [0; ITEMS], + } + } + + /// Return the static slot count, including unwritten slots. + #[inline] + pub const fn len(&self) -> usize { + ITEMS + } + + /// Return whether the payload has zero slots. + /// + /// This is always `false` for a successfully constructed payload. + #[inline] + pub const fn is_empty(&self) -> bool { + ITEMS == 0 + } + + /// Initialize or replace slot `index` with `value`, leaving other slots alone. + /// + /// # Panics + /// + /// Panics if `index >= self.len()`. + #[inline] + pub fn write(&mut self, index: usize, value: T) { + self.values[index].write(value); + self.initialized[index] = 1; + } + + /// Copy an initialized value, or return `None` if unwritten or out of bounds. + #[inline] + pub fn get(&self, index: usize) -> Option { + if index < ITEMS && self.initialized[index] == 1 { + // SAFETY: Only write sets this flag, after initializing the value. + // Copies preserve both arrays. Numeric types are Copy. + Some(unsafe { self.values[index].assume_init() }) + } else { + None + } + } + + /// Borrow an initialized value mutably, or return `None` if unwritten or + /// out of bounds. The exclusive borrow prevents simultaneous payload access. + #[inline] + pub fn get_mut(&mut self, index: usize) -> Option<&mut T> { + if index < ITEMS && self.initialized[index] == 1 { + // SAFETY: The flag proves initialization, and &mut self is exclusive. + Some(unsafe { self.values[index].assume_init_mut() }) + } else { + None + } + } +} + +impl Default for ThreadData +where + Alignment: SupportedAlignment, +{ + fn default() -> Self { + Self::new() + } +} + +impl Index + for ThreadData +where + Alignment: SupportedAlignment, +{ + type Output = T; + + #[inline] + fn index(&self, index: usize) -> &T { + assert!(index < ITEMS, "ThreadData index out of bounds"); + assert!(self.initialized[index] == 1, "ThreadData slot is unwritten"); + // SAFETY: The bounds and exact initialization flag were checked above. + unsafe { self.values[index].assume_init_ref() } + } +} + +impl IndexMut + for ThreadData +where + Alignment: SupportedAlignment, +{ + #[inline] + fn index_mut(&mut self, index: usize) -> &mut T { + assert!(index < ITEMS, "ThreadData index out of bounds"); + assert!(self.initialized[index] == 1, "ThreadData slot is unwritten"); + // SAFETY: The checks prove initialization; &mut self is exclusive. + unsafe { self.values[index].assume_init_mut() } + } +} + +/// A sealed interface for generic access to fixed-size per-thread payloads. +/// +/// Implemented by [`ThreadData`] for every supported element type, extent and +/// alignment. Consumers can write generic helpers without assuming every slot +/// is initialized. User implementations are not supported. +/// +/// ``` +/// use cuda_coop as coop; +/// fn copy_first(items: &mut D) { +/// if let Some(value) = items.get(0) { +/// items.write(D::EXTENT - 1, value); +/// } +/// } +/// let mut items = coop::ThreadData::<_, 4>::new(); +/// items.write(0, 3_i8); +/// copy_first(&mut items); +/// assert_eq!(items.get(3), Some(3_i8)); +/// assert_eq!(items.get(1), None); +/// ``` +pub trait ThreadDataLike: sealed::ThreadDataLike { + /// The numeric element type. + type Item: Numeric; + + /// The compile-time slot count, including unwritten slots. + const EXTENT: usize; + + /// Return the static slot count, including unwritten slots. + #[inline] + fn len(&self) -> usize { + Self::EXTENT + } + + /// Return whether the payload has zero slots. + #[inline] + fn is_empty(&self) -> bool { + Self::EXTENT == 0 + } + + /// Copy an initialized value, or return `None` if unwritten or out of bounds. + fn get(&self, index: usize) -> Option; + + /// Initialize or replace a single slot, preserving all other slots. + /// + /// # Panics + /// + /// Panics if `index >= self.len()`. + fn write(&mut self, index: usize, value: Self::Item); +} + +impl ThreadDataLike + for ThreadData +where + Alignment: SupportedAlignment, +{ + type Item = T; + const EXTENT: usize = ITEMS; + + #[inline] + fn get(&self, index: usize) -> Option { + self.get(index) + } + + #[inline] + fn write(&mut self, index: usize, value: T) { + self.write(index, value); + } +} + +mod sealed { + pub trait Numeric {} + pub trait ThreadDataLike {} +} + +macro_rules! numeric { + ($($ty:ty),* $(,)?) => { + $( + impl sealed::Numeric for $ty {} + impl Numeric for $ty {} + )* + }; +} + +numeric!(i8, u8, i16, u16, i32, u32, i64, u64, f32, f64); + +impl sealed::ThreadDataLike + for ThreadData +where + Alignment: SupportedAlignment, +{ +} + +#[cfg(test)] +mod tests { + extern crate std; + + use super::ThreadData; + use std::panic::{AssertUnwindSafe, catch_unwind}; + + #[test] + fn only_exactly_one_marks_a_slot_initialized() { + let mut items = ThreadData::<_, 1>::new(); + items.write(0, 42_i32); + for flag in [0, 2, 255] { + items.initialized[0] = flag; + assert_eq!(items.get(0), None); + assert!(items.get_mut(0).is_none()); + assert!(catch_unwind(|| items[0]).is_err()); + assert!(catch_unwind(AssertUnwindSafe(|| items[0] = 9)).is_err()); + } + items.write(0, 7); + assert_eq!(items.get(0), Some(7)); + } +} diff --git a/rust/cuda-coop/tests/payload.rs b/rust/cuda-coop/tests/payload.rs new file mode 100644 index 000000000000..8b24d6a7c2c8 --- /dev/null +++ b/rust/cuda-coop/tests/payload.rs @@ -0,0 +1,148 @@ +// Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +//! Public payload behavior, numeric inference, validity and element alignment. + +use core::mem::{align_of, size_of}; +use std::panic::{AssertUnwindSafe, catch_unwind}; + +use cuda_coop as coop; + +#[test] +fn typed_writes_infer_every_numeric_width_at_one_and_four_items() { + macro_rules! check { + ($($ty:ty),* $(,)?) => { + $( + for value in [<$ty>::MIN, <$ty>::MAX] { + let mut one = coop::ThreadData::<_, 1>::new(); + one.write(0, value); + assert_eq!(one.get(0), Some(value)); + assert_eq!(size_of_val(&one[0]), size_of::<$ty>()); + + let mut four = coop::ThreadData::<_, 4>::new(); + four.write(3, value); + assert_eq!(four.get(3), Some(value)); + assert_eq!(four.get(0), None); + assert_eq!(four.get(1), None); + assert_eq!(four.get(2), None); + } + )* + }; + } + check!(i8, u8, i16, u16, i32, u32, i64, u64, f32, f64); + + let mut single = coop::ThreadData::<_, 1>::new(); + single.write(0, f32::from_bits(0x7fc0_1234)); + assert_eq!(single[0].to_bits(), 0x7fc0_1234); + let mut double = coop::ThreadData::<_, 1>::new(); + double.write(0, f64::from_bits(0xfff8_0000_0000_1234)); + assert_eq!(double[0].to_bits(), 0xfff8_0000_0000_1234); +} + +#[test] +fn arbitrary_writes_and_checked_mutation_preserve_unwritten_slots() { + let mut items = coop::ThreadData::<_, 4>::new(); + items.write(2, 8_i32); + items.write(0, 3_i32); + items[2] *= 2; + *items.get_mut(0).unwrap() += 1; + assert_eq!( + [items.get(0), items.get(1), items.get(2), items.get(3)], + [Some(4), None, Some(16), None] + ); + assert_eq!(items.len(), 4); + assert!(!items.is_empty()); + assert!(items.get_mut(1).is_none()); + for index in [4, usize::MAX] { + assert_eq!(items.get(index), None); + assert!(items.get_mut(index).is_none()); + } + items.write(2, -1); + assert_eq!(items.get(2), Some(-1)); +} + +#[test] +fn invalid_indexing_and_writes_panic_without_initializing_other_slots() { + let mut items = coop::ThreadData::<_, 4>::new(); + items.write(2, 9_i32); + for index in [0, 4, usize::MAX] { + assert!(catch_unwind(|| items[index]).is_err()); + assert!(catch_unwind(AssertUnwindSafe(|| items[index] = 17)).is_err()); + } + for index in [4, usize::MAX] { + assert!(catch_unwind(AssertUnwindSafe(|| items.write(index, 17))).is_err()); + } + assert_eq!( + [items.get(0), items.get(1), items.get(2), items.get(3)], + [None, None, Some(9), None] + ); +} + +#[test] +fn copies_moves_swaps_and_selection_keep_exact_masks() { + let mut first = coop::ThreadData::<_, 4>::new(); + first.write(0, 1_u16); + first.write(3, 4_u16); + let original = first; + let mut second = coop::ThreadData::<_, 4>::new(); + second.write(1, 2_u16); + core::mem::swap(&mut first, &mut second); + assert_eq!( + (first.get(0), first.get(1), first.get(3)), + (None, Some(2), None) + ); + assert_eq!( + (second.get(0), second.get(1), second.get(3)), + (Some(1), None, Some(4)) + ); + for choose_first in [true, false] { + let mut selected = if choose_first { first } else { second }; + assert_eq!(selected.get(1), if choose_first { Some(2) } else { None }); + selected.write(2, 99); + assert_eq!(first.get(2), None); + assert_eq!(second.get(2), None); + } + assert_eq!( + (original.get(0), original.get(1), original.get(3)), + (Some(1), None, Some(4)) + ); +} + +#[test] +fn sealed_generic_interface_preserves_inference_and_extent() { + fn fill_last(items: &mut D, value: D::Item) { + assert_eq!(items.len(), D::EXTENT); + assert!(!items.is_empty()); + assert!(items.get(D::EXTENT - 1).is_none()); + items.write(D::EXTENT - 1, value); + } + let mut items = coop::ThreadData::<_, 4, 16>::new(); + fill_last(&mut items, 255_u8); + assert_eq!(items.get(3), Some(255_u8)); + assert_eq!(items.get(0), None); +} + +#[test] +fn alignment_applies_to_the_values_array() { + let mut items = coop::ThreadData::<_, 4, 64>::new(); + items.write(0, 13_u16); + items.write(1, 27_u16); + let values = &items[0] as *const u16 as usize; + assert_eq!(values, &items as *const _ as usize); + assert_eq!(values % 64, 0); + assert_eq!(&items[1] as *const u16 as usize - values, size_of::()); + assert_eq!(align_of::>(), align_of::()); + + // Query, rather than instantiate, the largest layouts: an over-aligned + // object can have a large padded size even with one element. + macro_rules! check { + ($($align:literal),* $(,)?) => { + $(assert_eq!(align_of::>(), $align);)* + }; + } + check!( + 1, 2, 4, 8, 16, 32, 64, 128, 256, 512, 1024, 2048, 4096, 8192, 16384, 32768, 65536, 131072, + 262144, 524288, 1048576, 2097152, 4194304, 8388608, 16777216, 33554432, 67108864, + 134217728, 268435456, 536870912 + ); +} From 8faf14b900143fabe717e72924c70d0edb26db4a Mon Sep 17 00:00:00 2001 From: Trent Nelson Date: Mon, 19 Oct 2026 22:59:30 -0700 Subject: [PATCH 06/33] Qualify checked native block calls and rejection cases Signed-off-by: Trent Nelson --- .../tests/check_collective_rejections.py | 254 ++++++++++++++++++ .../tests/oxide-fixture/README.md | 19 ++ .../tests/oxide-fixture/api/Cargo.toml | 3 + .../tests/oxide-fixture/api/src/lib.rs | 88 ++++++ .../compile_cases/native_mutation.rs | 43 +++ .../compile_cases/nested_reference.rs | 39 +++ .../tests/oxide-fixture/consumer/Cargo.toml | 3 + .../tests/oxide-fixture/consumer/src/main.rs | 10 +- .../tests/oxide-fixture/helper/Cargo.toml | 3 + .../tests/oxide-fixture/helper/src/checked.rs | 82 ++++++ .../tests/oxide-fixture/helper/src/lib.rs | 6 + .../tests/run_oxide_fixture.py | 6 + 12 files changed, 551 insertions(+), 5 deletions(-) create mode 100644 rust/cuda-coop-codegen/tests/check_collective_rejections.py create mode 100644 rust/cuda-coop-codegen/tests/oxide-fixture/compile_cases/native_mutation.rs create mode 100644 rust/cuda-coop-codegen/tests/oxide-fixture/compile_cases/nested_reference.rs create mode 100644 rust/cuda-coop-codegen/tests/oxide-fixture/helper/src/checked.rs diff --git a/rust/cuda-coop-codegen/tests/check_collective_rejections.py b/rust/cuda-coop-codegen/tests/check_collective_rejections.py new file mode 100644 index 000000000000..1d3f23f503c6 --- /dev/null +++ b/rust/cuda-coop-codegen/tests/check_collective_rejections.py @@ -0,0 +1,254 @@ +# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. +# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +"""Check rejected source programs in a copy made by run_oxide_fixture.py --checked. + +All edits are confined to that generated fixture and restored after each case. +Negative cases fail in the checked collective diagnostic, before launch. A +separate-site geometry case must compile with two shapes and one marker Instance. +""" + +import argparse +import json +import os +import subprocess +from pathlib import Path + +parser = argparse.ArgumentParser(description=__doc__) +parser.add_argument("output", type=Path) +parser.add_argument("--case", action="append", default=[]) +args = parser.parse_args() +output = args.output.resolve() +record = json.loads((output / "command.json").read_text()) +fixture = output / "fixture" +paths = { + "helper": fixture / "helper/src/checked.rs", + "api": fixture / "api/src/lib.rs", + "root": fixture / "consumer/src/main.rs", +} +original = {key: path.read_text() for key, path in paths.items()} +command = record["argv"].copy() +assert "checked" in command, "first run the fixture with --checked" +command[command.index("run")] = "build" +env = os.environ.copy() +env.update( + CCCL_SOURCE=str(Path(__file__).resolve().parents[3]), + CUDA_OXIDE_BACKEND=record["backend"], +) +if "cuda_root" in record: + env.update(CUDA_HOME=record["cuda_root"], CUDA_PATH=record["cuda_root"]) +env.pop("CUDA_OXIDE_UNCHECKED_INDEXING", None) +anchor = " let mut first = [0_i32; ITEMS_PER_THREAD];" + + +def in_helper(code): + return {"helper": original["helper"].replace(anchor, code + "\n" + anchor)} + + +shape_root = original["root"].replace("block = (128, 1, 1)", "block = (64, 1, 1)", 1) +for before, after in [ + ( + "roundtrip::<1, false>(input, output, post, valid, mode)", + "shape_probe::<0>(input)", + ), + ( + "roundtrip::<4, false>(input, output, post, valid, mode)", + "shape_probe::<1>(input)", + ), +]: + shape_root = shape_root.replace(before, after) +shape_api = ( + original["api"] + + """ +#[inline(never)] +pub unsafe fn shape_probe(input: *const i32, values: *mut i32) { + unsafe { cuda_device::library::checked_block_call_v1::(direct_load_full, (input,), (values,)); } +} +""" +) +shape_helper = ( + original["helper"] + + """ +#[inline(never)] +pub unsafe fn shape_probe(input: *const i32) { + let mut value = 0; + unsafe { coop_marker_fixture_api::shape_probe::(input, &mut value); } +} +""" +) +# Expose the additional helper only in the generated fixture copy. +helper_lib = fixture / "helper/src/lib.rs" +helper_lib_original = helper_lib.read_text() +cases = { + "distinct_shapes": ( + {"root": shape_root, "api": shape_api, "helper": shape_helper}, + None, + ), + "count": ( + in_helper(" let valid = valid + (lane & 1) as i32;"), + "not proved block-uniform", + ), + "carrier": ( + in_helper( + " let input = if lane & 1 == 0 { input } else { output as *const i32 };" + ), + "not proved block-uniform", + ), + "offset": ( + in_helper(" let input = unsafe { input.add(lane) };"), + "not proved block-uniform", + ), + "early_return": (in_helper(" if lane & 1 == 0 { return; }"), "lane-varying"), + "cycle": ( + in_helper( + " let mut remaining = lane;\n while remaining > 0 { remaining -= 1; }" + ), + "lane-varying cycle", + ), + "default": ( + { + "api": original["api"].replace( + "(input, valid, -7_i32),", + "(input, valid, cuda_device::thread::threadIdx_x() as i32),", + ) + }, + "not proved block-uniform", + ), + # The same helper Instance first receives uniform inputs, then varying inputs. + "call_context": ( + { + "root": original["root"].replace( + "roundtrip::<4, true>(input, output, post, valid, mode);", + "roundtrip::<4, true>(input, output, post, valid, mode);\n coop_marker_fixture_helper::roundtrip::<4, true>(input, output, post, valid + cuda_device::thread::threadIdx_x() as i32, mode);", + 1, + ) + }, + "not proved block-uniform", + ), + "mixed_raw": ( + in_helper( + " unsafe { coop_marker_fixture_api::load::(input, post_store, valid, true); }" + ), + "raw and checked library calls", + ), + "missing_shape": ( + {"root": original["root"].replace("block = (128, 1, 1), ", "")}, + "requires #[launch_contract(block", + ), + "shared_shape": ( + { + "root": "block = (64, 1, 1)".join( + original["root"].rsplit("block = (128, 1, 1)", 1) + ) + }, + "conflicting exact block shapes", + ), + "mode": (in_helper(" let mode = mode + (lane & 1) as u32;"), "lane-varying"), + "borrow_mutation": ( + { + "helper": original["helper"].replace( + anchor, + " let mut changed = valid;\n vary(&mut changed);\n let valid = changed;\n" + + anchor, + ) + + "\n#[inline(never)]\nfn vary(value: &mut i32) { *value += cuda_device::thread::threadIdx_x() as i32; }\n" + }, + "not proved block-uniform", + ), + "destructor": ( + { + "helper": original["helper"].replace( + anchor, " let _guard = Guard;\n" + anchor + ) + + "\nstruct Guard;\nimpl Drop for Guard { fn drop(&mut self) { cuda_device::thread::sync_threads(); } }\n" + }, + "destructor", + ), + "recursion": ( + { + "helper": original["helper"].replace( + anchor, " recurse(mode);\n" + anchor + ) + + "\n#[inline(never)]\nfn recurse(n: u32) { if n > 0 { recurse(n - 1); } }\n" + }, + "recursive helper", + ), + "unknown_call": ( + { + "helper": original["helper"].replace( + anchor, " unsafe { external_effect(); }\n" + anchor + ) + + '\nunsafe extern "C" { fn external_effect(); }\n' + }, + "without MIR", + ), + "panic": ( + in_helper(' if lane & 1 == 0 { panic!("lane exit"); }'), + "lane-varying", + ), + "unchecked_bounds": ({}, "unchecked indexing"), +} +probe_api = original["api"].replace("fn $name()", "pub fn $name()") +for name in ["native_mutation", "nested_reference", "nested_uniform"]: + source_name = "nested_reference" if name == "nested_uniform" else name + probe = ( + Path(__file__).parent / "oxide-fixture/compile_cases" / f"{source_name}.rs" + ).read_text() + if name == "nested_uniform": + probe = probe.replace("cuda_device::thread::threadIdx_x() as i32", "128_i32") + cases[name] = ( + {"root": probe, "api": probe_api}, + None if name == "nested_uniform" else "not proved block-uniform", + ) +unknown = set(args.case) - cases.keys() +if unknown: + parser.error(f"unknown cases: {', '.join(sorted(unknown))}") +results = {} +try: + for name, (changes, diagnostic) in cases.items(): + if args.case and name not in args.case: + continue + for key, path in paths.items(): + path.write_text(changes.get(key, original[key])) + helper_lib.write_text( + helper_lib_original + + ( + '\n#[cfg(feature = "checked")]\npub use checked::shape_probe;\n' + if name == "distinct_shapes" + else "" + ) + ) + if name in ["native_mutation", "nested_reference", "nested_uniform"]: + helper_lib.write_text( + helper_lib_original + + "\npub use coop_marker_fixture_api::{direct_load_full, direct_load_filled};\n" + ) + case_env = env.copy() + if name == "unchecked_bounds": + case_env["CUDA_OXIDE_UNCHECKED_INDEXING"] = "1" + log_path = output / f"reject-{name}.log" + with log_path.open("w") as log: + result = subprocess.run( + command, + cwd=fixture / "consumer", + env=case_env, + stdout=log, + stderr=subprocess.STDOUT, + ) + text = log_path.read_text() + if diagnostic is None: + assert result.returncode == 0, ( + f"{name}: expected successful compilation; see {log_path}" + ) + else: + assert result.returncode != 0, f"{name} unexpectedly compiled" + assert diagnostic in text, ( + f"{name}: expected {diagnostic!r}; see {log_path}" + ) + results[name] = {"returncode": result.returncode, "diagnostic": diagnostic} + print(f"PASS {name}", flush=True) +finally: + helper_lib.write_text(helper_lib_original) + for key, path in paths.items(): + path.write_text(original[key]) + (output / "rejections.json").write_text(json.dumps(results, indent=2) + "\n") diff --git a/rust/cuda-coop-codegen/tests/oxide-fixture/README.md b/rust/cuda-coop-codegen/tests/oxide-fixture/README.md index fb76581bf752..07f3b0f18834 100644 --- a/rust/cuda-coop-codegen/tests/oxide-fixture/README.md +++ b/rust/cuda-coop-codegen/tests/oxide-fixture/README.md @@ -36,3 +36,22 @@ 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. + +Add `--checked` to exercise the internal checked full-block marker. This route +gets block geometry from each root contract, separates common controls from +lane-local payload pointers, selects either of two payload arrays per lane after +reconvergence, and repeats Load in a uniform loop with explicit synchronization. +The raw unsafe route remains the default. Both routes retain caller obligations +for bounds, initialized values, unique output ownership and scratch reuse. + +After a checked run, exercise rejected source programs in the generated copy: + +```sh +python3 rust/cuda-coop-codegen/tests/check_collective_rejections.py \ + /path/to/new-evidence-directory +``` + +The cases cover varying controls/carriers, different contexts of one helper, +participation and unsupported bounds-elision modes. Sources are restored after +each run. The initial checker diagnoses a shared helper site reached under two +different exact block shapes; separate sites remain supported. diff --git a/rust/cuda-coop-codegen/tests/oxide-fixture/api/Cargo.toml b/rust/cuda-coop-codegen/tests/oxide-fixture/api/Cargo.toml index 0d5afeb1bb00..450d33d5678e 100644 --- a/rust/cuda-coop-codegen/tests/oxide-fixture/api/Cargo.toml +++ b/rust/cuda-coop-codegen/tests/oxide-fixture/api/Cargo.toml @@ -12,3 +12,6 @@ cuda-device.workspace = true [build-dependencies] cuda-coop-codegen = { path = "../../.." } + +[features] +checked = [] 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 index 5f555bfd6953..f0f464a2c760 100644 --- a/rust/cuda-coop-codegen/tests/oxide-fixture/api/src/lib.rs +++ b/rust/cuda-coop-codegen/tests/oxide-fixture/api/src/lib.rs @@ -93,3 +93,91 @@ pub unsafe fn store( } } } + +#[cfg(feature = "checked")] +mod checked { + use super::*; + /// # 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 checked_load( + input: *const i32, + values: *mut i32, + valid: i32, + filled: bool, + ) { + unsafe { + if filled { + if TRANSPOSE { + cuda_device::library::checked_block_call_v1::( + transpose_load_filled, + (input, valid, -7_i32), + (values,), + ); + } else { + cuda_device::library::checked_block_call_v1::( + direct_load_filled, + (input, valid, -7_i32), + (values,), + ); + } + } else if TRANSPOSE { + cuda_device::library::checked_block_call_v1::( + transpose_load_full, + (input,), + (values,), + ); + } else { + cuda_device::library::checked_block_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 checked_store( + output: *mut i32, + values: *mut i32, + valid: i32, + partial: bool, + ) { + unsafe { + if partial { + if TRANSPOSE { + cuda_device::library::checked_block_call_v1::( + transpose_store_partial, + (output, valid), + (values,), + ); + } else { + cuda_device::library::checked_block_call_v1::( + direct_store_partial, + (output, valid), + (values,), + ); + } + } else if TRANSPOSE { + cuda_device::library::checked_block_call_v1::( + transpose_store_full, + (output,), + (values,), + ); + } else { + cuda_device::library::checked_block_call_v1::( + direct_store_full, + (output,), + (values,), + ); + } + } + } +} +#[cfg(feature = "checked")] +pub use checked::{checked_load, checked_store}; diff --git a/rust/cuda-coop-codegen/tests/oxide-fixture/compile_cases/native_mutation.rs b/rust/cuda-coop-codegen/tests/oxide-fixture/compile_cases/native_mutation.rs new file mode 100644 index 000000000000..a0e57f30b2b4 --- /dev/null +++ b/rust/cuda-coop-codegen/tests/oxide-fixture/compile_cases/native_mutation.rs @@ -0,0 +1,43 @@ +// Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception +use coop_marker_fixture_helper::{direct_load_filled, direct_load_full}; +use cuda_device::{kernel, launch_contract}; +use cuda_host::cuda_module; +#[inline(never)] +fn twice( + value: &&i32, + _a: i32, + _b: i32, + _c: i32, + _d: i32, + _e: i32, + _f: i32, + _g: i32, + _h: i32, +) -> i32 { + **value +} +#[cuda_module] +mod kernels { + use super::*; + #[kernel] + #[launch_contract(block = (128,1,1), domain = 3, dynamic_shared = inferred)] + pub unsafe fn probe(input: *const i32) { + let mut count = 1_i32; + let mut payload = 0_i32; + unsafe { + cuda_device::library::checked_block_call_v1::( + direct_load_full, + (input,), + (&mut count as *mut i32,), + ); + cuda_device::thread::sync_threads(); + cuda_device::library::checked_block_call_v1::( + direct_load_filled, + (input, count, -7_i32), + (&mut payload as *mut i32,), + ); + } + } +} +fn main() {} diff --git a/rust/cuda-coop-codegen/tests/oxide-fixture/compile_cases/nested_reference.rs b/rust/cuda-coop-codegen/tests/oxide-fixture/compile_cases/nested_reference.rs new file mode 100644 index 000000000000..c2df4f7d7382 --- /dev/null +++ b/rust/cuda-coop-codegen/tests/oxide-fixture/compile_cases/nested_reference.rs @@ -0,0 +1,39 @@ +// Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception +use coop_marker_fixture_helper::{direct_load_filled, direct_load_full}; +use cuda_device::{kernel, launch_contract}; +use cuda_host::cuda_module; +#[inline(never)] +fn twice( + value: &&i32, + _a: i32, + _b: i32, + _c: i32, + _d: i32, + _e: i32, + _f: i32, + _g: i32, + _h: i32, +) -> i32 { + **value +} +#[cuda_module] +mod kernels { + use super::*; + #[kernel] + #[launch_contract(block = (128,1,1), domain = 3, dynamic_shared = inferred)] + pub unsafe fn probe(input: *const i32) { + let value = cuda_device::thread::threadIdx_x() as i32; + let reference = &value; + let count = twice(&reference, 0, 0, 0, 0, 0, 0, 0, 0); + let mut payload = 0_i32; + unsafe { + cuda_device::library::checked_block_call_v1::( + direct_load_filled, + (input, count, -7_i32), + (&mut payload as *mut i32,), + ); + } + } +} +fn main() {} diff --git a/rust/cuda-coop-codegen/tests/oxide-fixture/consumer/Cargo.toml b/rust/cuda-coop-codegen/tests/oxide-fixture/consumer/Cargo.toml index e2393ccec577..eb5841743397 100644 --- a/rust/cuda-coop-codegen/tests/oxide-fixture/consumer/Cargo.toml +++ b/rust/cuda-coop-codegen/tests/oxide-fixture/consumer/Cargo.toml @@ -13,3 +13,6 @@ cuda-host.workspace = true cuda-core.workspace = true oxide-artifacts.workspace = true coop-marker-fixture-helper = { path = "../helper" } + +[features] +checked = ["coop-marker-fixture-helper/checked"] 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 index eba29e10753e..716937b82d51 100644 --- a/rust/cuda-coop-codegen/tests/oxide-fixture/consumer/src/main.rs +++ b/rust/cuda-coop-codegen/tests/oxide-fixture/consumer/src/main.rs @@ -9,7 +9,7 @@ mod kernels { use super::*; #[kernel] #[launch_bounds(256)] - #[launch_contract(domain = 3, dynamic_shared = inferred)] + #[launch_contract(block = (128, 1, 1), domain = 3, dynamic_shared = inferred)] pub unsafe fn direct_1( input: *const i32, output: *mut i32, @@ -23,7 +23,7 @@ mod kernels { } #[kernel] #[launch_bounds(256)] - #[launch_contract(domain = 3, dynamic_shared = inferred)] + #[launch_contract(block = (128, 1, 1), domain = 3, dynamic_shared = inferred)] pub unsafe fn direct_4( input: *const i32, output: *mut i32, @@ -37,7 +37,7 @@ mod kernels { } #[kernel] #[launch_bounds(256)] - #[launch_contract(domain = 3, dynamic_shared = inferred)] + #[launch_contract(block = (128, 1, 1), domain = 3, dynamic_shared = inferred)] pub unsafe fn transpose_1( input: *const i32, output: *mut i32, @@ -51,7 +51,7 @@ mod kernels { } #[kernel] #[launch_bounds(256)] - #[launch_contract(domain = 3, dynamic_shared = inferred)] + #[launch_contract(block = (128, 1, 1), domain = 3, dynamic_shared = inferred)] pub unsafe fn transpose_4( input: *const i32, output: *mut i32, @@ -65,7 +65,7 @@ mod kernels { } #[kernel] #[launch_bounds(256)] - #[launch_contract(domain = 3, dynamic_shared = inferred)] + #[launch_contract(block = (128, 1, 1), domain = 3, dynamic_shared = inferred)] pub unsafe fn shared_helper_root( input: *const i32, output: *mut i32, diff --git a/rust/cuda-coop-codegen/tests/oxide-fixture/helper/Cargo.toml b/rust/cuda-coop-codegen/tests/oxide-fixture/helper/Cargo.toml index a997768dc07c..abb1aba763ac 100644 --- a/rust/cuda-coop-codegen/tests/oxide-fixture/helper/Cargo.toml +++ b/rust/cuda-coop-codegen/tests/oxide-fixture/helper/Cargo.toml @@ -10,3 +10,6 @@ publish = false [dependencies] cuda-device.workspace = true coop-marker-fixture-api = { path = "../api" } + +[features] +checked = ["coop-marker-fixture-api/checked"] diff --git a/rust/cuda-coop-codegen/tests/oxide-fixture/helper/src/checked.rs b/rust/cuda-coop-codegen/tests/oxide-fixture/helper/src/checked.rs new file mode 100644 index 000000000000..f425c21babda --- /dev/null +++ b/rust/cuda-coop-codegen/tests/oxide-fixture/helper/src/checked.rs @@ -0,0 +1,82 @@ +// Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +/// 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 first = [0_i32; ITEMS_PER_THREAD]; + let mut second = [0_i32; ITEMS_PER_THREAD]; + // Lane-local choice reconverges before every collective. + let payload = if lane & 1 == 0 { + &mut first + } else { + &mut second + }; + unsafe { + // A uniform loop retains full participation. Scratch reuse still needs + // the explicit caller-owned barrier; the checker inserts none. + let mut iteration = 0; + while iteration < 2 { + coop_marker_fixture_api::checked_load::( + input.add(origin), + payload.as_mut_ptr(), + valid, + mode & 1 != 0, + ); + cuda_device::thread::sync_threads(); + iteration += 1; + } + let mut i = 0; + while i < ITEMS_PER_THREAD { + payload[i] += 1; + i += 1; + } + coop_marker_fixture_api::checked_store::( + output.add(origin), + payload.as_mut_ptr(), + valid, + mode & 2 != 0, + ); + let mut i = 0; + while i < ITEMS_PER_THREAD { + post_store + .add(origin + lane * ITEMS_PER_THREAD + i) + .write(payload[i]); + i += 1; + } + cuda_device::thread::sync_threads(); + if mode & 4 != 0 { + coop_marker_fixture_api::checked_load::( + output.add(origin), + payload.as_mut_ptr(), + valid, + false, + ); + let mut i = 0; + while i < ITEMS_PER_THREAD { + post_store + .add(origin + lane * ITEMS_PER_THREAD + i) + .write(payload[i]); + i += 1; + } + } + } +} 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 index 025d76c5fb6d..019ca58eefcc 100644 --- a/rust/cuda-coop-codegen/tests/oxide-fixture/helper/src/lib.rs +++ b/rust/cuda-coop-codegen/tests/oxide-fixture/helper/src/lib.rs @@ -11,6 +11,7 @@ /// 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. +#[cfg(not(feature = "checked"))] #[inline(never)] pub unsafe fn roundtrip( input: *const i32, @@ -63,3 +64,8 @@ pub unsafe fn roundtrip( } } } + +#[cfg(feature = "checked")] +mod checked; +#[cfg(feature = "checked")] +pub use checked::roundtrip; diff --git a/rust/cuda-coop-codegen/tests/run_oxide_fixture.py b/rust/cuda-coop-codegen/tests/run_oxide_fixture.py index 6b357575759d..a6644ccc769c 100644 --- a/rust/cuda-coop-codegen/tests/run_oxide_fixture.py +++ b/rust/cuda-coop-codegen/tests/run_oxide_fixture.py @@ -17,6 +17,9 @@ 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") +parser.add_argument( + "--checked", action="store_true", help="exercise the checked full-block marker" +) args = parser.parse_args() source = Path(__file__).resolve().parent checkout = source.parents[2] @@ -63,6 +66,8 @@ args.arch, "--materialize-cubin", ] +if args.checked: + command.extend(["--features", "checked"]) with (output / "build-run.log").open("w") as log: result = subprocess.run( command, cwd=fixture / "consumer", env=env, stdout=log, stderr=subprocess.STDOUT @@ -74,6 +79,7 @@ "cwd": str(fixture / "consumer"), "oxide_root": str(oxide), "backend": env["CUDA_OXIDE_BACKEND"], + "cuda_root": env["CUDA_HOME"], "returncode": result.returncode, }, indent=2, From 184a1c9f3e3bf0bf33e8b340d7aec81d39c06724 Mon Sep 17 00:00:00 2001 From: Trent Nelson Date: Mon, 19 Oct 2026 23:19:30 -0700 Subject: [PATCH 07/33] Add dimension-free CUDA hierarchy and group queries Signed-off-by: Trent Nelson --- docs/rust/groups.rst | 94 +++ docs/rust/index.rst | 12 +- docs/rust/payload.rst | 8 +- rust/Cargo.lock | 165 ++++ rust/cuda-coop/Cargo.toml | 4 +- rust/cuda-coop/src/group.rs | 582 +++++++++++++ rust/cuda-coop/src/lib.rs | 19 +- .../cuda-coop/tests/device_queries/.gitignore | 7 + .../cuda-coop/tests/device_queries/Cargo.lock | 785 ++++++++++++++++++ .../cuda-coop/tests/device_queries/Cargo.toml | 16 + .../tests/device_queries/src/main.rs | 335 ++++++++ 11 files changed, 2014 insertions(+), 13 deletions(-) create mode 100644 docs/rust/groups.rst create mode 100644 rust/cuda-coop/src/group.rs create mode 100644 rust/cuda-coop/tests/device_queries/.gitignore create mode 100644 rust/cuda-coop/tests/device_queries/Cargo.lock create mode 100644 rust/cuda-coop/tests/device_queries/Cargo.toml create mode 100644 rust/cuda-coop/tests/device_queries/src/main.rs diff --git a/docs/rust/groups.rst b/docs/rust/groups.rst new file mode 100644 index 000000000000..8e411a01c500 --- /dev/null +++ b/docs/rust/groups.rst @@ -0,0 +1,94 @@ +.. Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. +.. SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +Hierarchy and group queries +========================== + +``cuda-coop`` describes the current CUDA launch without requiring dimensions in +the Rust source. ``this_thread()``, ``this_warp()``, ``this_block()``, +``this_cluster()`` and ``this_grid()`` return typed group descriptors. +``ThreadHierarchy::new()`` and ``ThreadHierarchy::current()`` provide the same +physical groups through ``thread()``, ``warp()``, ``block()``, ``cluster()`` and +``grid()``. ``Hierarchy`` aliases ``ThreadHierarchy``. + +Queries run in an Oxide-compiled kernel with the workspace's pinned +``nightly-2026-10-11`` toolchain. The device dependency is pinned in Cargo; there +is no backend-selection feature. Host payload tests require no CUDA installation. +The host provider remains independently usable on Rust 1.85. + +Rank and count +-------------- + +``group.rank()`` and ``group.count()`` default to the Thread level. Rust selects +an explicit level with ``rank_at::()`` or +``count_at::()``. An inner level queries constituents within the +group; an outer level queries that group within its parent. Equal levels have +rank zero and count one. Physical queries use x-fastest linearization. + +For example, ``this_block().rank()`` is the current thread's linear rank in the +block, while ``this_block().rank_at::()`` is the block's linear rank +in the grid. The result is ``u64`` when either scope or query level is Grid, +and ``u32`` otherwise. Grid arithmetic promotes coordinates before multiplying. +``rank_as::()`` and ``count_as::()`` select one of the eight fixed-width +integer types; ``rank_at_as::()`` also selects a level. +Conversions follow integer ``as`` semantics, so callers must choose a sufficient +width to avoid truncation. + +A physical warp always counts 32 lanes. Block warp counts round up when the +last warp is partially occupied. Cluster queries currently require sm_90 or +newer; an ordinary, unclustered launch has one block per cluster. No query +allocates storage or synchronizes threads. Copying a descriptor preserves the +scope, and subsequent queries describe the thread executing them. + +Mapped groups +------------- + +``this_warp().group_by::<8>()`` partitions a physical warp into groups of eight +lanes. ``this_block().group_by::<2>()`` partitions a block into groups of two +**physical warps**, each containing 64 threads. The default is exhaustive: +partition size must divide the parent's unit count. Both forms require complete +physical warps in the enclosing block. + +``group_by_non_exhaustive::()`` excludes leftover units. Its ``is_member()`` +query returns ``1u8`` for members and ``0u8`` for excluded threads. Counts remain +defined for nonmembers; rank queries on a nonmember trap. Lane groups accept +only Thread and Warp levels; warp groups accept Thread, Warp and Block levels. +Mapped groups cannot be partitioned again. + +Sizes must be positive compile-time constants. Invalid static sizes fail during +compilation. Checks that depend on the actual block dimensions trap on device +when violated. These descriptors and checks provide no proof of collective +participation, uniform control flow, or scratch ownership. Synchronization, +temporary storage and collective primitives are separate work. + +Run the native query checks +--------------------------- + +The standalone ``rust/cuda-coop/tests/device_queries`` package contains the +complete native program. It verifies all physical scope/level pairs with +multidimensional blocks and grids, a partial final warp, a clustered launch, +exhaustive and non-exhaustive mappings, and integer conversions. It uses the +same pinned public Oxide dependencies as the library. Keeping this native +package separate lets ordinary workspace tests run without CUDA host libraries. + +With the matching ``cargo-oxide`` compiler installed, a CUDA toolkit, and an +sm_120 GPU, run from the CCCL checkout: + +.. code-block:: shell + + cd rust/cuda-coop/tests/device_queries + CUDA_HOME=/usr/local/cuda-13.3 cargo oxide run --arch sm_120 --materialize-cubin + compute-sanitizer --tool memcheck --error-exitcode 1 target/release/device_queries + +The four rejection modes run in separate processes because a device trap makes +its CUDA context unusable: + +.. code-block:: shell + + for case in 0 1 2 3; do + target/release/device_queries "$case" + done + +They check partial physical warps, an indivisible exhaustive partition, a +partition larger than its parent, and a rank query by an excluded lane. The +program reports success only when CUDA returns the expected device-trap error. diff --git a/docs/rust/index.rst b/docs/rust/index.rst index 72be476dcb99..6250a0f5cf3a 100644 --- a/docs/rust/index.rst +++ b/docs/rust/index.rst @@ -4,14 +4,19 @@ Rust development ================ -The Rust workspace contains per-thread payload types and the native host compiler -foundation for ``cuda.coop``. The ``cuda-coop`` crate supplies checked -``ThreadData`` values; see :doc:`payload`. Its private ``cuda-coop-codegen`` crate +The Rust workspace contains per-thread payloads, CUDA group queries, and the +native host compiler foundation for ``cuda.coop``. The ``cuda-coop`` crate supplies checked +``ThreadData`` values and physical/mapped hierarchy queries; see :doc:`payload` +and :doc:`groups`. 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. +Device code uses the pinned Oxide dependency and ``nightly-2026-10-11``. +The host provider supports Rust 1.85 independently; validate it with +``cargo +1.85.0 test -p cuda-coop-codegen --locked`` from ``rust/``. + 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. @@ -114,3 +119,4 @@ invalidation of external compiler inputs belong to the compiler integration. :maxdepth: 1 payload + groups diff --git a/docs/rust/payload.rst b/docs/rust/payload.rst index ae2d0f55d26c..be5aebab8b98 100644 --- a/docs/rust/payload.rst +++ b/docs/rust/payload.rst @@ -5,10 +5,10 @@ Per-thread payloads =================== The unpublished ``cuda-coop`` crate provides ``ThreadData`` and its sealed -``ThreadDataLike`` interface. It is a feature-free ``no_std`` crate with no -dependencies, and supports Rust 1.85 or newer. Oxide is the intended device -compiler integration. This crate currently supplies the payload vocabulary; -collectives, groups and scratch storage are still under development. +``ThreadDataLike`` interface. It is a feature-free ``no_std`` crate using the +pinned Oxide device dependency and ``nightly-2026-10-11`` toolchain. See +:doc:`groups` for physical and mapped group queries. Collective operations and +scratch storage are still under development. Construct and access a payload ------------------------------ diff --git a/rust/Cargo.lock b/rust/Cargo.lock index 650a0f45e4ce..c0d1de3a583d 100644 --- a/rust/Cargo.lock +++ b/rust/Cargo.lock @@ -39,6 +39,9 @@ dependencies = [ [[package]] name = "cuda-coop" version = "0.0.0" +dependencies = [ + "cuda-device", +] [[package]] name = "cuda-coop-codegen" @@ -48,6 +51,26 @@ dependencies = [ "sha2", ] +[[package]] +name = "cuda-device" +version = "0.2.1" +source = "git+https://github.com/tpn/cuda-rust.git?rev=8c4a42e9c95b2a6bdc14a7531904d8469f0e7ca6#8c4a42e9c95b2a6bdc14a7531904d8469f0e7ca6" +dependencies = [ + "cuda-macros", +] + +[[package]] +name = "cuda-macros" +version = "0.2.1" +source = "git+https://github.com/tpn/cuda-rust.git?rev=8c4a42e9c95b2a6bdc14a7531904d8469f0e7ca6#8c4a42e9c95b2a6bdc14a7531904d8469f0e7ca6" +dependencies = [ + "oxide-tile-export", + "proc-macro2", + "quote", + "reserved-oxide-symbols", + "syn 2.0.119", +] + [[package]] name = "digest" version = "0.10.7" @@ -68,6 +91,12 @@ dependencies = [ "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" @@ -84,6 +113,88 @@ dependencies = [ "windows-link", ] +[[package]] +name = "memchr" +version = "2.8.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cf8baf1c55e62ffcace7a9f06f4bd9cd3f0c4beb022d3b367256b91b87513d98" + +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +source = "git+https://github.com/tpn/cuda-rust.git?rev=8c4a42e9c95b2a6bdc14a7531904d8469f0e7ca6#8c4a42e9c95b2a6bdc14a7531904d8469f0e7ca6" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + +[[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 = "reserved-oxide-symbols" +version = "0.2.1" +source = "git+https://github.com/tpn/cuda-rust.git?rev=8c4a42e9c95b2a6bdc14a7531904d8469f0e7ca6#8c4a42e9c95b2a6bdc14a7531904d8469f0e7ca6" + +[[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 3.0.9", +] + +[[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" @@ -95,12 +206,60 @@ dependencies = [ "digest", ] +[[package]] +name = "syn" +version = "2.0.119" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "872831b642d1a07999a962a351ed35b955ea2cfc8f3862091e2a240a84f17297" +dependencies = [ + "proc-macro2", + "quote", + "unicode-ident", +] + +[[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 = "thiserror" +version = "2.0.21" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "09e52cb86a36cede5cb101bf8908837b3e4c6e5e59fe7fd85c23fb56200d189e" +dependencies = [ + "thiserror-impl", +] + +[[package]] +name = "thiserror-impl" +version = "2.0.21" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fe5197923287db20a58125f0bc85c062f7f2c892de97b18c356f9efb14b28524" +dependencies = [ + "proc-macro2", + "quote", + "syn 3.0.9", +] + [[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" @@ -112,3 +271,9 @@ 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/cuda-coop/Cargo.toml b/rust/cuda-coop/Cargo.toml index 1ca43ea36394..03e46ac0c204 100644 --- a/rust/cuda-coop/Cargo.toml +++ b/rust/cuda-coop/Cargo.toml @@ -6,8 +6,10 @@ name = "cuda-coop" version = "0.0.0" edition.workspace = true license.workspace = true -rust-version.workspace = true publish = false +[dependencies] +cuda-device = { git = "https://github.com/tpn/cuda-rust.git", rev = "8c4a42e9c95b2a6bdc14a7531904d8469f0e7ca6" } + [lints] workspace = true diff --git a/rust/cuda-coop/src/group.rs b/rust/cuda-coop/src/group.rs new file mode 100644 index 000000000000..d3ea2e900b23 --- /dev/null +++ b/rust/cuda-coop/src/group.rs @@ -0,0 +1,582 @@ +// Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +use core::marker::PhantomData; +use cuda_device::{cluster, thread}; + +mod sealed { + pub trait Scope { + fn member() -> u8; + } + pub trait Level {} + pub trait Integer {} + pub trait Group {} +} + +/// A supported physical or mapped group scope. +/// +/// Sealed to this crate's scope types. A scope grants no collective +/// participation or synchronization guarantee. +pub trait GroupScope: sealed::Scope {} + +/// A physical CUDA hierarchy level: Thread, Warp, Block, Cluster, or Grid. +/// +/// This sealed trait supplies the level argument to [`ThreadGroup::rank_at`] +/// and [`ThreadGroup::count_at`]. Mapped scopes are not query levels. +pub trait Level: GroupScope + sealed::Level {} + +macro_rules! levels { + ($(($name:ident, $new:ident, $doc:literal)),+ $(,)?) => {$ ( + #[doc = $doc] + #[derive(Clone, Copy, Debug)] + pub enum $name {} + impl sealed::Scope for $name { + #[inline(always)] fn member() -> u8 { 1 } + } + impl GroupScope for $name {} + impl sealed::Level for $name {} + impl Level for $name {} + #[doc = concat!("Describe the current ", $doc, "\n\nNo dimensions are supplied. This owns no resources and does not synchronize.")] + #[inline(always)] + pub const fn $new() -> ThreadGroup<$name> { ThreadGroup(PhantomData) } + )+}; +} +levels!( + (Thread, this_thread, "CUDA thread."), + (Warp, this_warp, "physical CUDA warp (32 lanes)."), + (Block, this_block, "CUDA thread block."), + (Cluster, this_cluster, "CUDA thread-block cluster."), + (Grid, this_grid, "CUDA grid."), +); + +/// Integer types accepted by rank/count conversions. +/// +/// Sealed to `i8`, `u8`, `i16`, `u16`, `i32`, `u32`, `i64`, and `u64`. +/// Conversions follow integer `as` casts, including truncation and signed +/// reinterpretation. Choose a type wide enough for the result. +pub trait QueryInteger: sealed::Integer + Copy { + #[doc(hidden)] + fn from_u64(value: u64) -> Self; + #[doc(hidden)] + fn to_u64(self) -> u64; +} +macro_rules! integers { + ($($t:ty),+) => {$ ( + impl sealed::Integer for $t {} + impl QueryInteger for $t { + #[inline(always)] fn from_u64(value: u64) -> Self { value as Self } + #[inline(always)] fn to_u64(self) -> u64 { self as u64 } + } + )+}; +} +integers!(i8, u8, i16, u16, i32, u32, i64, u64); + +/// A supported scope/level query and its natural integer width. +/// +/// Sealed to [`ThreadGroup`] implementations. Physical groups accept every +/// physical level. [`ThreadsWithinWarp`] accepts Thread/Warp; +/// [`WarpsWithinBlock`] accepts Thread/Warp/Block. +/// +/// ```no_run +/// use cuda_coop::{GroupScope, QueryAt, Thread, ThreadGroup}; +/// fn threads(group: ThreadGroup) +/// -> as QueryAt>::Output +/// where ThreadGroup: QueryAt +/// { group.count() } +/// ``` +pub trait QueryAt: sealed::Group { + /// `u64` when either scope or query level is Grid, otherwise `u32`. + type Output: QueryInteger; + #[doc(hidden)] + fn query_rank(&self) -> Self::Output; + #[doc(hidden)] + fn query_count(&self) -> Self::Output; +} + +/// A dimension-free descriptor of the executing thread's current group. +/// +/// Construct with [`this_thread`], [`this_warp`], [`this_block`], +/// [`this_cluster`], or [`this_grid`]. `S` fixes the available queries. +/// Copies describe the thread executing each query, not the original thread. +/// +/// Native queries require an Oxide-compiled kernel. They read launch state +/// without allocating, moving payloads, or synchronizing. Descriptors carry +/// no uniformity or collective safety proof. Cluster queries currently need +/// sm_90 or newer; an ordinary launch has one block per cluster. +#[derive(Debug)] +pub struct ThreadGroup(PhantomData); +impl Copy for ThreadGroup {} +impl Clone for ThreadGroup { + fn clone(&self) -> Self { + *self + } +} +impl sealed::Group for ThreadGroup {} +impl ThreadGroup { + /// Return the thread's rank in this group; equivalent to + /// `rank_at::()`. Panics (traps on device) for mapped nonmembers. + #[inline(always)] + pub fn rank(self) -> >::Output + where + Self: QueryAt, + { + self.rank_at::() + } + + /// Count threads in this group; equivalent to `count_at::()`. + /// A physical warp has 32 lanes, including a partially occupied last warp. + #[inline(always)] + pub fn count(self) -> >::Output + where + Self: QueryAt, + { + self.count_at::() + } + + /// Query rank at `L`, with x-fastest linearization. + /// + /// An inner level selects a constituent's rank in this group; an outer + /// level selects this group's rank in that parent. Same-level rank is + /// zero. The result is `u64` if Grid is involved, otherwise `u32`. + /// Panics (traps on device) for an excluded mapped-group member. + /// + /// ```no_run + /// use cuda_coop as coop; + /// let block = coop::this_block(); + /// let thread: u32 = block.rank(); + /// let warp: u32 = block.rank_at::(); + /// let block_in_grid: u64 = block.rank_at::(); + /// ``` + #[inline(always)] + pub fn rank_at(self) -> >::Output + where + Self: QueryAt, + { + self.query_rank() + } + + /// Count constituents or parents at `L`. + /// + /// An inner level counts constituents in this group; an outer level + /// counts these groups in that parent. Same-level count is one. The + /// result is `u64` if Grid is involved, otherwise `u32`. Mapped counts + /// remain defined for excluded trailing threads. + #[inline(always)] + pub fn count_at(self) -> >::Output + where + Self: QueryAt, + { + self.query_count() + } + + /// Convert [`Self::rank`] to integer `I` with integer `as` semantics. + /// Mapped membership checks still apply. + #[inline(always)] + pub fn rank_as(self) -> I + where + Self: QueryAt, + { + self.rank_at_as::() + } + + /// Convert [`Self::count`] to integer `I` with integer `as` semantics. + #[inline(always)] + pub fn count_as(self) -> I + where + Self: QueryAt, + { + self.count_at_as::() + } + + /// Convert [`Self::rank_at`] at `L` to integer `I` with integer `as` + /// semantics. Mapped membership checks still apply. + #[inline(always)] + pub fn rank_at_as(self) -> I + where + Self: QueryAt, + { + I::from_u64(self.rank_at::().to_u64()) + } + + /// Convert [`Self::count_at`] at `L` to integer `I` with integer `as` + /// semantics. + #[inline(always)] + pub fn count_at_as(self) -> I + where + Self: QueryAt, + { + I::from_u64(self.count_at::().to_u64()) + } + + /// Return `1u8` for a member, `0u8` for an excluded trailing unit. + /// Physical groups return one. This query grants no participation proof. + #[inline(always)] + pub fn is_member(self) -> u8 { + S::member() + } +} + +/// A dimension-free view of the current physical CUDA hierarchy. +/// +/// [`Self::new`] and [`Self::current`] are equivalent. The hierarchy stores +/// no launch dimensions or resources and grants no synchronization capability. +#[derive(Clone, Copy, Debug, Default)] +pub struct ThreadHierarchy(()); +/// Alias for [`ThreadHierarchy`]. +pub type Hierarchy = ThreadHierarchy; +impl ThreadHierarchy { + /// Describe the current launch without supplying dimensions. + pub const fn new() -> Self { + Self(()) + } + /// Describe the current launch; equivalent to [`Self::new`]. + pub const fn current() -> Self { + Self::new() + } + /// Describe the current thread. + pub const fn thread(self) -> ThreadGroup { + this_thread() + } + /// Describe the current physical warp. + pub const fn warp(self) -> ThreadGroup { + this_warp() + } + /// Describe the current thread block. + pub const fn block(self) -> ThreadGroup { + this_block() + } + /// Describe the current thread-block cluster. + pub const fn cluster(self) -> ThreadGroup { + this_cluster() + } + /// Describe the current grid. + pub const fn grid(self) -> ThreadGroup { + this_grid() + } +} + +// Read only the registers required by the monomorphized query. +#[inline(always)] +fn thread_rank_block() -> u32 { + (thread::threadIdx_z() * thread::blockDim_y() + thread::threadIdx_y()) * thread::blockDim_x() + + thread::threadIdx_x() +} +#[inline(always)] +fn threads_block() -> u32 { + thread::blockDim_x() * thread::blockDim_y() * thread::blockDim_z() +} +#[inline(always)] +fn warps_block() -> u32 { + threads_block().div_ceil(32) +} +#[inline(always)] +fn linear_u64(x: u32, y: u32, z: u32, nx: u32, ny: u32) -> u64 { + (u64::from(z) * u64::from(ny) + u64::from(y)) * u64::from(nx) + u64::from(x) +} +#[inline(always)] +fn block_rank_grid() -> u64 { + linear_u64( + thread::blockIdx_x(), + thread::blockIdx_y(), + thread::blockIdx_z(), + thread::gridDim_x(), + thread::gridDim_y(), + ) +} +#[inline(always)] +fn blocks_grid() -> u64 { + u64::from(thread::gridDim_x()) * u64::from(thread::gridDim_y()) * u64::from(thread::gridDim_z()) +} +#[inline(always)] +fn cluster_rank_grid() -> u64 { + linear_u64( + thread::blockIdx_x() / cluster::cluster_nctaidX(), + thread::blockIdx_y() / cluster::cluster_nctaidY(), + thread::blockIdx_z() / cluster::cluster_nctaidZ(), + thread::gridDim_x() / cluster::cluster_nctaidX(), + thread::gridDim_y() / cluster::cluster_nctaidY(), + ) +} +#[inline(always)] +fn clusters_grid() -> u64 { + u64::from(thread::gridDim_x() / cluster::cluster_nctaidX()) + * u64::from(thread::gridDim_y() / cluster::cluster_nctaidY()) + * u64::from(thread::gridDim_z() / cluster::cluster_nctaidZ()) +} +macro_rules! query { + ($scope:ty, $level:ty, $out:ty, $rank:expr, $count:expr) => { + impl QueryAt<$level> for ThreadGroup<$scope> { + type Output = $out; + #[inline(always)] + fn query_rank(&self) -> $out { + $rank + } + #[inline(always)] + fn query_count(&self) -> $out { + $count + } + } + }; +} +macro_rules! pair { + ($inner:ty, $outer:ty, $out:ty, $rank:expr, $count:expr) => { + query!($inner, $outer, $out, $rank, $count); + query!($outer, $inner, $out, $rank, $count); + }; +} +query!(Thread, Thread, u32, 0, 1); +query!(Warp, Warp, u32, 0, 1); +query!(Block, Block, u32, 0, 1); +query!(Cluster, Cluster, u32, 0, 1); +query!(Grid, Grid, u64, 0, 1); +pair!(Thread, Warp, u32, thread_rank_block() % 32, 32); +pair!(Thread, Block, u32, thread_rank_block(), threads_block()); +pair!(Warp, Block, u32, thread_rank_block() / 32, warps_block()); +pair!( + Thread, + Cluster, + u32, + cluster::block_rank() * threads_block() + thread_rank_block(), + cluster::cluster_size() * threads_block() +); +pair!( + Warp, + Cluster, + u32, + cluster::block_rank() * warps_block() + thread_rank_block() / 32, + cluster::cluster_size() * warps_block() +); +pair!( + Block, + Cluster, + u32, + cluster::block_rank(), + cluster::cluster_size() +); +pair!( + Thread, + Grid, + u64, + block_rank_grid() + .wrapping_mul(u64::from(threads_block())) + .wrapping_add(u64::from(thread_rank_block())), + blocks_grid().wrapping_mul(u64::from(threads_block())) +); +pair!( + Warp, + Grid, + u64, + block_rank_grid() + .wrapping_mul(u64::from(warps_block())) + .wrapping_add(u64::from(thread_rank_block() / 32)), + blocks_grid().wrapping_mul(u64::from(warps_block())) +); +pair!(Block, Grid, u64, block_rank_grid(), blocks_grid()); +pair!(Cluster, Grid, u64, cluster_rank_grid(), clusters_grid()); + +/// Consecutive groups of `N` lanes within a physical warp. +/// +/// Construct with [`ThreadGroup::group_by`] or its non-exhaustive +/// variant. `N` is in 1..=32. `EXHAUSTIVE = true` requires divisibility; +/// otherwise trailing lanes are excluded. The block must have complete +/// warps. Only Thread/Warp queries are available; groups cannot be nested. +/// +/// ```compile_fail +/// let group = cuda_coop::this_warp().group_by::<8>(); +/// group.rank_at::(); +/// ``` +/// ```compile_fail +/// cuda_coop::this_warp().group_by::<8>().group_by::<2>(); +/// ``` +#[derive(Clone, Copy, Debug)] +pub enum ThreadsWithinWarp {} +/// Consecutive groups of `N` physical warps within a block. +/// +/// Construct with [`ThreadGroup::group_by`] or its non-exhaustive +/// variant. The block must have complete warps and at least `N` warps. +/// `EXHAUSTIVE = true` requires divisibility; otherwise trailing warps are +/// excluded. Only Thread/Warp/Block queries are available; no nesting. +/// +/// ```compile_fail +/// cuda_coop::this_block().group_by::<2>().count_at::(); +/// ``` +#[derive(Clone, Copy, Debug)] +pub enum WarpsWithinBlock {} +impl ThreadsWithinWarp { + const VALID: () = { + assert!(N > 0 && N <= 32, "lane partition size must be in 1..=32"); + assert!( + !E || 32 % N == 0, + "exhaustive lane partition must divide 32" + ); + }; +} +impl WarpsWithinBlock { + const VALID: () = { + assert!(N > 0 && N <= 32, "warp partition size must be in 1..=32"); + }; +} +impl sealed::Scope for ThreadsWithinWarp { + #[inline(always)] + fn member() -> u8 { + ((thread_rank_block() % 32) < (32 / N as u32) * N as u32) as u8 + } +} +impl GroupScope for ThreadsWithinWarp {} +impl sealed::Scope for WarpsWithinBlock { + #[inline(always)] + fn member() -> u8 { + ((thread_rank_block() / 32) < (warps_block() / N as u32) * N as u32) as u8 + } +} +impl GroupScope for WarpsWithinBlock {} +impl ThreadGroup { + /// Partition into consecutive groups of `N` lanes. + /// + /// `N` must be positive and divide 32 (checked at compilation). Panics + /// (traps on device) if the block has a partial physical warp. + /// + /// ```compile_fail + /// cuda_coop::this_warp().group_by::<0>(); + /// ``` + /// ```compile_fail + /// cuda_coop::this_warp().group_by::<6>(); + /// ``` + #[inline(always)] + pub fn group_by(self) -> ThreadGroup> { + self.partition::() + } + /// Partition into `N`-lane groups, excluding leftover lanes. + /// + /// `N` must be in 1..=32. Panics (traps on device) for a partial physical + /// warp. Check [`ThreadGroup::is_member`] before querying a member rank. + #[inline(always)] + pub fn group_by_non_exhaustive( + self, + ) -> ThreadGroup> { + self.partition::() + } + #[inline(always)] + fn partition(self) -> ThreadGroup> { + let () = ThreadsWithinWarp::::VALID; + assert!( + threads_block().is_multiple_of(32), + "lane partitions require complete physical warps" + ); + ThreadGroup(PhantomData) + } +} +impl ThreadGroup { + /// Partition into consecutive groups of `N` physical warps. + /// + /// `N` must be in 1..=32 (checked at compilation). Panics (traps on + /// device) unless the block has complete warps and its warp count is + /// divisible by `N`. + /// + /// ```compile_fail + /// cuda_coop::this_block().group_by::<33>(); + /// ``` + #[inline(always)] + pub fn group_by(self) -> ThreadGroup> { + self.partition::() + } + /// Partition into `N`-warp groups, excluding leftover warps. + /// + /// `N` must be in 1..=32. Panics (traps on device) unless the block has + /// complete warps and at least `N` warps. Check [`ThreadGroup::is_member`] + /// before querying a possibly excluded thread's rank. + #[inline(always)] + pub fn group_by_non_exhaustive( + self, + ) -> ThreadGroup> { + self.partition::() + } + #[inline(always)] + fn partition(self) -> ThreadGroup> { + let () = WarpsWithinBlock::::VALID; + let warps = warps_block(); + assert!( + threads_block().is_multiple_of(32), + "warp partitions require complete physical warps" + ); + assert!( + N as u32 <= warps, + "partition exceeds the block's warp count" + ); + assert!( + !E || warps.is_multiple_of(N as u32), + "exhaustive partition must divide the block's warp count" + ); + ThreadGroup(PhantomData) + } +} +macro_rules! mapped_query { + ($scope:ident, $level:ty, $rank:expr, $count:expr) => { + impl QueryAt<$level> for ThreadGroup<$scope> { + type Output = u32; + #[inline(always)] + fn query_rank(&self) -> u32 { + assert!( + self.is_member() == 1, + "excluded thread has no mapped-group rank" + ); + $rank + } + #[inline(always)] + fn query_count(&self) -> u32 { + $count + } + } + }; +} +mapped_query!( + ThreadsWithinWarp, + Thread, + thread_rank_block() % 32 % N as u32, + N as u32 +); +mapped_query!( + ThreadsWithinWarp, + Warp, + thread_rank_block() % 32 / N as u32, + 32 / N as u32 +); +mapped_query!( + WarpsWithinBlock, + Thread, + (thread_rank_block() / 32 % N as u32) * 32 + thread_rank_block() % 32, + N as u32 * 32 +); +mapped_query!( + WarpsWithinBlock, + Warp, + thread_rank_block() / 32 % N as u32, + N as u32 +); +mapped_query!( + WarpsWithinBlock, + Block, + thread_rank_block() / 32 / N as u32, + warps_block() / N as u32 +); + +#[cfg(test)] +mod tests { + use super::*; + #[test] + fn grid_linearization_promotes_before_multiplication() { + assert_eq!( + linear_u64(7, 65534, 2, 2_147_483_647, 65535), + 422_203_874_934_795 + ); + } + #[test] + fn query_casts_preserve_integer_cast_semantics() { + assert_eq!(::from_u64(255), -1); + assert_eq!(::from_u64(65537), 1); + assert_eq!( + ::from_u64(u32::MAX as u64 + 1), + 4_294_967_296 + ); + } +} diff --git a/rust/cuda-coop/src/lib.rs b/rust/cuda-coop/src/lib.rs index ec97f008886b..3d575be24c8a 100644 --- a/rust/cuda-coop/src/lib.rs +++ b/rust/cuda-coop/src/lib.rs @@ -1,7 +1,7 @@ // Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception -//! Checked per-thread payloads for Rust `cuda.coop`. +//! Per-thread payloads and CUDA hierarchy queries for Rust `cuda.coop`. //! //! This unpublished, feature-free `no_std` crate provides [`ThreadData`] and //! its sealed [`ThreadDataLike`] interface. Each slot starts unwritten. @@ -26,15 +26,24 @@ //! assert_eq!(items[0], 9); //! ``` //! -//! The crate requires Rust 1.85 or newer and has no host compiler or CUDA -//! dependencies. Oxide is the intended device compiler integration. This -//! payload API supplies no collective operations, group descriptors, scratch -//! storage, or synchronization guarantees. +//! Device queries use Oxide and the pinned `nightly-2026-10-11` toolchain. +//! [`this_block`] and the other dimension-free group constructors describe +//! the executing thread's current launch. Queries grant no collective +//! participation or synchronization guarantee. This crate supplies no +//! collective operations or scratch storage. Host payload tests need no GPU; +//! executing hardware queries requires an Oxide-compiled CUDA kernel. #![no_std] mod alignment; +mod group; mod thread_data; pub use alignment::{Alignment, SupportedAlignment}; pub use thread_data::{Numeric, ThreadData, ThreadDataLike}; + +pub use group::{ + Block, Cluster, Grid, GroupScope, Hierarchy, Level, QueryAt, QueryInteger, Thread, ThreadGroup, + ThreadHierarchy, ThreadsWithinWarp, Warp, WarpsWithinBlock, this_block, this_cluster, + this_grid, this_thread, this_warp, +}; diff --git a/rust/cuda-coop/tests/device_queries/.gitignore b/rust/cuda-coop/tests/device_queries/.gitignore new file mode 100644 index 000000000000..74bffa8af4a4 --- /dev/null +++ b/rust/cuda-coop/tests/device_queries/.gitignore @@ -0,0 +1,7 @@ +# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. +# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception +/target/ +/.oxide-artifacts/ +/*.ll +/*.options +/*.target diff --git a/rust/cuda-coop/tests/device_queries/Cargo.lock b/rust/cuda-coop/tests/device_queries/Cargo.lock new file mode 100644 index 000000000000..2026e4b3bedf --- /dev/null +++ b/rust/cuda-coop/tests/device_queries/Cargo.lock @@ -0,0 +1,785 @@ +# This file is automatically @generated by Cargo. +# It is not intended for manual editing. +version = 4 + +[[package]] +name = "aho-corasick" +version = "1.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c982642fa9e8606056828ee9a8505737230110bb1099153c79efe865c59d12ba" +dependencies = [ + "memchr", +] + +[[package]] +name = "anyhow" +version = "1.0.104" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "330a5ed07fa54e4702c9d6c4174f74427fc0ef6e214bbd677ae50a5099946470" + +[[package]] +name = "bindgen" +version = "0.69.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "271383c67ccabffb7381723dea0672a673f292304fcb45c01cc648c7a8d58088" +dependencies = [ + "bitflags", + "cexpr", + "clang-sys", + "itertools", + "lazy_static", + "lazycell", + "log", + "prettyplease", + "proc-macro2", + "quote", + "regex", + "rustc-hash", + "shlex", + "syn 2.0.119", + "which", +] + +[[package]] +name = "bitflags" +version = "2.13.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3ded4057c258ba199e2d26386d3af3780957ecaee6c4ef4041c6b4b8b97c0b06" + +[[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 = "cexpr" +version = "0.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6fac387a98bb7c37292057cffc56d62ecb629900026402633ae9160df93a8766" +dependencies = [ + "nom", +] + +[[package]] +name = "cfg-if" +version = "1.0.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4e7648175b45a9a48536d676f68d918270699102aa8dab5496df06904c914600" + +[[package]] +name = "clang-sys" +version = "1.9.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "157a8ba7b480713b56f4c09fd13fc3e0a22a5dfab8097ba61cbc5feef950788a" +dependencies = [ + "glob", + "libc", + "libloading", +] + +[[package]] +name = "cpufeatures" +version = "0.2.17" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "59ed5838eebb26a2bb2e58f6d5b5316989ae9d08bab10e0e6d103e656d1b0280" +dependencies = [ + "libc", +] + +[[package]] +name = "crc32fast" +version = "1.5.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "01a7799fd6b852db0e61728dde9a204c423b44d689dbd432522543614b490e78" +dependencies = [ + "cfg-if", +] + +[[package]] +name = "crunchy" +version = "0.2.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "460fbee9c2c2f33933d720630a6a0bac33ba7053db5344fac858d4b8952d77d5" + +[[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-artifact-finalizer" +version = "0.2.1" +source = "git+https://github.com/tpn/cuda-rust.git?rev=8c4a42e9c95b2a6bdc14a7531904d8469f0e7ca6#8c4a42e9c95b2a6bdc14a7531904d8469f0e7ca6" +dependencies = [ + "libnvvm-sys", + "nvjitlink-sys", + "serde", + "sha2", + "thiserror", +] + +[[package]] +name = "cuda-bindings" +version = "0.4.0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=8c4a42e9c95b2a6bdc14a7531904d8469f0e7ca6#8c4a42e9c95b2a6bdc14a7531904d8469f0e7ca6" +dependencies = [ + "bindgen", + "libloading", + "prettyplease", + "proc-macro2", + "quote", + "syn 2.0.119", +] + +[[package]] +name = "cuda-coop" +version = "0.0.0" +dependencies = [ + "cuda-device", +] + +[[package]] +name = "cuda-core" +version = "0.4.0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=8c4a42e9c95b2a6bdc14a7531904d8469f0e7ca6#8c4a42e9c95b2a6bdc14a7531904d8469f0e7ca6" +dependencies = [ + "anyhow", + "cuda-bindings", + "cuda-core-derive", + "half", + "oxide-artifacts", +] + +[[package]] +name = "cuda-core-derive" +version = "0.4.0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=8c4a42e9c95b2a6bdc14a7531904d8469f0e7ca6#8c4a42e9c95b2a6bdc14a7531904d8469f0e7ca6" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", +] + +[[package]] +name = "cuda-device" +version = "0.2.1" +source = "git+https://github.com/tpn/cuda-rust.git?rev=8c4a42e9c95b2a6bdc14a7531904d8469f0e7ca6#8c4a42e9c95b2a6bdc14a7531904d8469f0e7ca6" +dependencies = [ + "cuda-macros", +] + +[[package]] +name = "cuda-host" +version = "0.2.1" +source = "git+https://github.com/tpn/cuda-rust.git?rev=8c4a42e9c95b2a6bdc14a7531904d8469f0e7ca6#8c4a42e9c95b2a6bdc14a7531904d8469f0e7ca6" +dependencies = [ + "cuda-artifact-finalizer", + "cuda-core", + "cuda-macros", + "half", + "libc", + "oxide-artifacts", + "ptx-parse", + "sha2", + "thiserror", +] + +[[package]] +name = "cuda-macros" +version = "0.2.1" +source = "git+https://github.com/tpn/cuda-rust.git?rev=8c4a42e9c95b2a6bdc14a7531904d8469f0e7ca6#8c4a42e9c95b2a6bdc14a7531904d8469f0e7ca6" +dependencies = [ + "oxide-tile-export", + "proc-macro2", + "quote", + "reserved-oxide-symbols", + "syn 2.0.119", +] + +[[package]] +name = "cuda-target-spec" +version = "0.1.0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=8c4a42e9c95b2a6bdc14a7531904d8469f0e7ca6#8c4a42e9c95b2a6bdc14a7531904d8469f0e7ca6" + +[[package]] +name = "device_queries" +version = "0.0.0" +dependencies = [ + "cuda-coop", + "cuda-core", + "cuda-device", + "cuda-host", +] + +[[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 = "either" +version = "1.19.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0e9c71c2167ca323c882b99918929403426e2373ea17242ff5653e0d5e1058be" + +[[package]] +name = "equivalent" +version = "1.0.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "00d174d5400e5e8fd687ad1049e2f578285fa914201b1af7e8b112a4546bd826" + +[[package]] +name = "errno" +version = "0.3.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "39cab71617ae0d63f51a36d69f866391735b51691dbda63cf6f96d042b63efeb" +dependencies = [ + "libc", + "windows-sys 0.61.2", +] + +[[package]] +name = "foldhash" +version = "0.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d9c4f5dac5e15c24eb999c26181a6ca40b39fe946cbe4c263c7209467bc83af2" + +[[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 = "glob" +version = "0.3.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e4eba85ea1d0a966a983acd07deee566e67395d2d96b6fb39e62b5a833f1eb0b" + +[[package]] +name = "half" +version = "2.7.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" +dependencies = [ + "cfg-if", + "crunchy", + "zerocopy", +] + +[[package]] +name = "hashbrown" +version = "0.15.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9229cfe53dfd69f0609a49f65461bd93001ea1ef889cd5529dd176593f5338a1" +dependencies = [ + "foldhash", +] + +[[package]] +name = "hashbrown" +version = "0.17.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ed5909b6e89a2db4456e54cd5f673791d7eca6732202bbf2a9cc504fe2f9b84a" + +[[package]] +name = "home" +version = "0.5.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cc627f471c528ff0c4a49e1d5e60450c8f6461dd6d10ba9dcd3a61d3dff7728d" +dependencies = [ + "windows-sys 0.61.2", +] + +[[package]] +name = "indexmap" +version = "2.14.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cc4e190f5d26ca7051642629da2c52fc03bde85a03197c99408dcd291734c855" +dependencies = [ + "equivalent", + "hashbrown 0.17.1", +] + +[[package]] +name = "itertools" +version = "0.12.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ba291022dbbd398a455acf126c1e341954079855bc60dfdda641363bd6922569" +dependencies = [ + "either", +] + +[[package]] +name = "itoa" +version = "1.0.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8f42a60cbdf9a97f5d2305f08a87dc4e09308d1276d28c869c684d7777685682" + +[[package]] +name = "lazy_static" +version = "1.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "20870f649af7073d53e38067b2a84312175d56ea15217e1b15bc83506ec50afb" + +[[package]] +name = "lazycell" +version = "1.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "830d08ce1d1d941e6b30645f1a0eb5643013d835ce3779a5fc208261dbe10f55" + +[[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 = "libnvvm-sys" +version = "0.2.1" +source = "git+https://github.com/tpn/cuda-rust.git?rev=8c4a42e9c95b2a6bdc14a7531904d8469f0e7ca6#8c4a42e9c95b2a6bdc14a7531904d8469f0e7ca6" +dependencies = [ + "cuda-target-spec", + "libloading", + "thiserror", +] + +[[package]] +name = "linux-raw-sys" +version = "0.4.15" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d26c52dbd32dccf2d10cac7725f8eae5296885fb5703b261f7d0a0739ec807ab" + +[[package]] +name = "log" +version = "0.4.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9f8bd3e56ce4dfc153cf470fffbfa98c7620958b312ca5c3a4b8d5181fd13c6" + +[[package]] +name = "memchr" +version = "2.8.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cf8baf1c55e62ffcace7a9f06f4bd9cd3f0c4beb022d3b367256b91b87513d98" + +[[package]] +name = "minimal-lexical" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "68354c5c6bd36d73ff3feceb05efa59b6acb7626617f4962be322a825e61f79a" + +[[package]] +name = "nom" +version = "7.1.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d273983c5a657a70a3e8f2a01329822f3b8c8172b73826411a55751e404a0a4a" +dependencies = [ + "memchr", + "minimal-lexical", +] + +[[package]] +name = "nvjitlink-sys" +version = "0.2.1" +source = "git+https://github.com/tpn/cuda-rust.git?rev=8c4a42e9c95b2a6bdc14a7531904d8469f0e7ca6#8c4a42e9c95b2a6bdc14a7531904d8469f0e7ca6" +dependencies = [ + "libloading", + "thiserror", +] + +[[package]] +name = "object" +version = "0.36.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "62948e14d923ea95ea2c7c86c71013138b66525b86bdc08d2dcc262bdb497b87" +dependencies = [ + "crc32fast", + "hashbrown 0.15.5", + "indexmap", + "memchr", +] + +[[package]] +name = "once_cell" +version = "1.21.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9f7c3e4beb33f85d45ae3e3a1792185706c8e16d043238c593331cc7cd313b50" + +[[package]] +name = "oxide-artifacts" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "74d9a941c0c43d2dceaf25f5c9ca4c9f9cceb28dfa5e4e3cfd14e476f003e7c9" +dependencies = [ + "object", +] + +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +source = "git+https://github.com/tpn/cuda-rust.git?rev=8c4a42e9c95b2a6bdc14a7531904d8469f0e7ca6#8c4a42e9c95b2a6bdc14a7531904d8469f0e7ca6" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + +[[package]] +name = "prettyplease" +version = "0.2.37" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "479ca8adacdd7ce8f1fb39ce9ecccbfe93a3f1344b3d0d97f20bc0196208f62b" +dependencies = [ + "proc-macro2", + "syn 2.0.119", +] + +[[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 = "ptx-parse" +version = "0.1.0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=8c4a42e9c95b2a6bdc14a7531904d8469f0e7ca6#8c4a42e9c95b2a6bdc14a7531904d8469f0e7ca6" + +[[package]] +name = "quote" +version = "1.0.47" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1fbf4db142a473a8d80c26bbf18454ed458bf8d26c8219c331daecfdbd079001" +dependencies = [ + "proc-macro2", +] + +[[package]] +name = "regex" +version = "1.13.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f020237b6c8eed93db2e2cb53c00c60a8e1bc73da7d073199a1180401450218d" +dependencies = [ + "aho-corasick", + "memchr", + "regex-automata", + "regex-syntax", +] + +[[package]] +name = "regex-automata" +version = "0.4.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ad8553b9b26413251cbf30e620595c7a41b3887f03da04579c0e6b0d6a06b4b2" +dependencies = [ + "aho-corasick", + "memchr", + "regex-syntax", +] + +[[package]] +name = "regex-syntax" +version = "0.8.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d6f6ff9a378485b298a5286656da665ba74413d36db0979633275d2e708145d4" + +[[package]] +name = "reserved-oxide-symbols" +version = "0.2.1" +source = "git+https://github.com/tpn/cuda-rust.git?rev=8c4a42e9c95b2a6bdc14a7531904d8469f0e7ca6#8c4a42e9c95b2a6bdc14a7531904d8469f0e7ca6" + +[[package]] +name = "rustc-hash" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "08d43f7aa6b08d49f382cde6a7982047c3426db949b1424bc4b7ec9ae12c6ce2" + +[[package]] +name = "rustix" +version = "0.38.44" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fdb5bc1ae2baa591800df16c9ca78619bf65c0488b41b96ccec5d11220d8c154" +dependencies = [ + "bitflags", + "errno", + "libc", + "linux-raw-sys", + "windows-sys 0.59.0", +] + +[[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 3.0.9", +] + +[[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 = "shlex" +version = "1.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0fda2ff0d084019ba4d7c6f371c95d8fd75ce3524c3cb8fb653a3023f6323e64" + +[[package]] +name = "syn" +version = "2.0.119" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "872831b642d1a07999a962a351ed35b955ea2cfc8f3862091e2a240a84f17297" +dependencies = [ + "proc-macro2", + "quote", + "unicode-ident", +] + +[[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 = "thiserror" +version = "2.0.21" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "09e52cb86a36cede5cb101bf8908837b3e4c6e5e59fe7fd85c23fb56200d189e" +dependencies = [ + "thiserror-impl", +] + +[[package]] +name = "thiserror-impl" +version = "2.0.21" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fe5197923287db20a58125f0bc85c062f7f2c892de97b18c356f9efb14b28524" +dependencies = [ + "proc-macro2", + "quote", + "syn 3.0.9", +] + +[[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 = "which" +version = "4.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "87ba24419a2078cd2b0f2ede2691b6c66d8e47836da3b6db8265ebad47afbfc7" +dependencies = [ + "either", + "home", + "once_cell", + "rustix", +] + +[[package]] +name = "windows-link" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f0805222e57f7521d6a62e36fa9163bc891acd422f971defe97d64e70d0a4fe5" + +[[package]] +name = "windows-sys" +version = "0.59.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1e38bc4d79ed67fd075bcc251a1c39b32a1776bbe92e5bef1f0bf1f8c531853b" +dependencies = [ + "windows-targets", +] + +[[package]] +name = "windows-sys" +version = "0.61.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ae137229bcbd6cdf0f7b80a31df61766145077ddf49416a728b02cb3921ff3fc" +dependencies = [ + "windows-link", +] + +[[package]] +name = "windows-targets" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9b724f72796e036ab90c1021d4780d4d3d648aca59e491e6b98e725b84e99973" +dependencies = [ + "windows_aarch64_gnullvm", + "windows_aarch64_msvc", + "windows_i686_gnu", + "windows_i686_gnullvm", + "windows_i686_msvc", + "windows_x86_64_gnu", + "windows_x86_64_gnullvm", + "windows_x86_64_msvc", +] + +[[package]] +name = "windows_aarch64_gnullvm" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "32a4622180e7a0ec044bb555404c800bc9fd9ec262ec147edd5989ccd0c02cd3" + +[[package]] +name = "windows_aarch64_msvc" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "09ec2a7bb152e2252b53fa7803150007879548bc709c039df7627cabbd05d469" + +[[package]] +name = "windows_i686_gnu" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8e9b5ad5ab802e97eb8e295ac6720e509ee4c243f69d781394014ebfe8bbfa0b" + +[[package]] +name = "windows_i686_gnullvm" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0eee52d38c090b3caa76c563b86c3a4bd71ef1a819287c19d586d7334ae8ed66" + +[[package]] +name = "windows_i686_msvc" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "240948bc05c5e7c6dabba28bf89d89ffce3e303022809e73deaefe4f6ec56c66" + +[[package]] +name = "windows_x86_64_gnu" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "147a5c80aabfbf0c7d901cb5895d1de30ef2907eb21fbbab29ca94c5b08b1a78" + +[[package]] +name = "windows_x86_64_gnullvm" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "24d5b23dc417412679681396f2b49f3de8c1473deb516bd34410872eff51ed0d" + +[[package]] +name = "windows_x86_64_msvc" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "589f6da84c646204747d1270a2a5661ea66ed1cced2631d546fdfb155959f9ec" + +[[package]] +name = "zerocopy" +version = "0.8.63" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e5fe1f8f1b06191a00962174c61aa5005e0bb391a6d80d07e24d115c01a92ed8" +dependencies = [ + "zerocopy-derive", +] + +[[package]] +name = "zerocopy-derive" +version = "0.8.63" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "863ad3ac83293fb4d740aedbfdc9240dd8d1a50c1099acd76ce80ce7c7230c7f" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", +] + +[[package]] +name = "zmij" +version = "1.0.23" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "29666d0abbfad1e3dc4dcf6144730dd3a3ab225bbbdac83319345b1b44ccfc1b" diff --git a/rust/cuda-coop/tests/device_queries/Cargo.toml b/rust/cuda-coop/tests/device_queries/Cargo.toml new file mode 100644 index 000000000000..7e29090be630 --- /dev/null +++ b/rust/cuda-coop/tests/device_queries/Cargo.toml @@ -0,0 +1,16 @@ +# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. +# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +[package] +name = "device_queries" +version = "0.0.0" +edition = "2024" +publish = false + +[workspace] + +[dependencies] +cuda-coop = { path = "../.." } +cuda-device = { git = "https://github.com/tpn/cuda-rust.git", rev = "8c4a42e9c95b2a6bdc14a7531904d8469f0e7ca6" } +cuda-host = { git = "https://github.com/tpn/cuda-rust.git", rev = "8c4a42e9c95b2a6bdc14a7531904d8469f0e7ca6" } +cuda-core = { git = "https://github.com/tpn/cuda-rust.git", rev = "8c4a42e9c95b2a6bdc14a7531904d8469f0e7ca6" } diff --git a/rust/cuda-coop/tests/device_queries/src/main.rs b/rust/cuda-coop/tests/device_queries/src/main.rs new file mode 100644 index 000000000000..12ce62674ea4 --- /dev/null +++ b/rust/cuda-coop/tests/device_queries/src/main.rs @@ -0,0 +1,335 @@ +// Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception +//! Query physical and mapped groups in an Oxide-compiled CUDA kernel. +use cuda_coop as coop; +use cuda_device::{cluster_launch, device, kernel}; +use cuda_host::cuda_module; +const PHYSICAL: usize = 55; +const MAPPED: usize = 20; + +// Check each local query's public return width while converting test output. +macro_rules! local { + ($query:expr) => {{ + let value: u32 = $query; + u64::from(value) + }}; +} + +#[device] +unsafe fn physical(output: *mut u64) { + let hierarchy = coop::Hierarchy::current(); + let thread = hierarchy.thread(); + let warp = hierarchy.warp(); + let block = hierarchy.block(); + let cluster = hierarchy.cluster(); + let grid = hierarchy.grid(); + let values = [ + local!(thread.rank_at::()), + local!(thread.count_at::()), + local!(thread.rank_at::()), + local!(thread.count_at::()), + local!(thread.rank_at::()), + local!(thread.count_at::()), + local!(thread.rank_at::()), + local!(thread.count_at::()), + thread.rank_at::(), + thread.count_at::(), + local!(warp.rank_at::()), + local!(warp.count_at::()), + local!(warp.rank_at::()), + local!(warp.count_at::()), + local!(warp.rank_at::()), + local!(warp.count_at::()), + local!(warp.rank_at::()), + local!(warp.count_at::()), + warp.rank_at::(), + warp.count_at::(), + local!(block.rank_at::()), + local!(block.count_at::()), + local!(block.rank_at::()), + local!(block.count_at::()), + local!(block.rank_at::()), + local!(block.count_at::()), + local!(block.rank_at::()), + local!(block.count_at::()), + block.rank_at::(), + block.count_at::(), + local!(cluster.rank_at::()), + local!(cluster.count_at::()), + local!(cluster.rank_at::()), + local!(cluster.count_at::()), + local!(cluster.rank_at::()), + local!(cluster.count_at::()), + local!(cluster.rank_at::()), + local!(cluster.count_at::()), + cluster.rank_at::(), + cluster.count_at::(), + grid.rank_at::(), + grid.count_at::(), + grid.rank_at::(), + grid.count_at::(), + grid.rank_at::(), + grid.count_at::(), + grid.rank_at::(), + grid.count_at::(), + grid.rank_at::(), + grid.count_at::(), + thread.is_member() as u64, + warp.is_member() as u64, + block.is_member() as u64, + cluster.is_member() as u64, + grid.is_member() as u64, + ]; + let base = grid.rank() as usize * PHYSICAL; + for (i, value) in values.into_iter().enumerate() { + unsafe { + output.add(base + i).write(value); + } + } +} + +#[cuda_module] +mod kernels { + use super::*; + #[kernel] + pub unsafe fn queries(output: *mut u64) { + unsafe { + physical(output); + } + } + #[kernel] + #[cluster_launch(2, 2, 1)] + pub unsafe fn clustered(output: *mut u64) { + unsafe { + physical(output); + } + } + #[kernel] + pub unsafe fn mapped(output: *mut u64) { + let lanes = coop::this_warp().group_by_non_exhaustive::<6>(); + let warps = coop::this_block().group_by_non_exhaustive::<2>(); + let exact = coop::this_warp().group_by::<8>(); + let single = coop::this_block().group_by::<1>(); + let lane_member: u8 = lanes.is_member(); + let warp_member: u8 = warps.is_member(); + let values = [ + lane_member as u64, + lanes.count() as u64, + local!(lanes.count_at::()), + if lane_member == 1 { + lanes.rank() as u64 + } else { + u64::MAX + }, + if lane_member == 1 { + local!(lanes.rank_at::()) + } else { + u64::MAX + }, + warp_member as u64, + warps.count() as u64, + local!(warps.count_at::()), + local!(warps.count_at::()), + if warp_member == 1 { + warps.rank() as u64 + } else { + u64::MAX + }, + if warp_member == 1 { + local!(warps.rank_at::()) + } else { + u64::MAX + }, + if warp_member == 1 { + local!(warps.rank_at::()) + } else { + u64::MAX + }, + exact.rank_as::() as u64, + exact.count_as::() as u64, + exact.rank_at_as::() as u64, + exact.count_at_as::() as u64, + single.rank_as::() as u64, + single.count_as::() as u64, + single.rank_at_as::() as u64, + single.count_at_as::(), + ]; + let base = coop::this_grid().rank() as usize * MAPPED; + for (i, value) in values.into_iter().enumerate() { + unsafe { + output.add(base + i).write(value); + } + } + } + // Each negative case runs in a separate process because a device trap + // makes the CUDA context unusable for subsequent launches. + #[kernel] + pub unsafe fn invalid(output: *mut u32, mode: u32) { + let value = match mode { + 0 => coop::this_warp().group_by::<8>().count(), + 1 => coop::this_block().group_by::<2>().count(), + 2 => coop::this_block().group_by_non_exhaustive::<4>().count(), + _ => coop::this_warp().group_by_non_exhaustive::<6>().rank(), + }; + unsafe { + output.add(coop::this_block().rank() as usize).write(value); + } + } +} + +fn main() -> Result<(), Box> { + use cuda_core::simt::LaunchConfig; + use cuda_core::{CudaContext, DeviceBuffer}; + let ctx = CudaContext::new(0)?; + let stream = ctx.default_stream(); + let module = kernels::load(&ctx)?; + if let Some(mode) = std::env::args().nth(1) { + let mode: u32 = mode.parse()?; + assert!(mode <= 3); + let threads = if mode == 0 { 48 } else { 96 }; + let out = DeviceBuffer::::zeroed(&stream, threads)?; + unsafe { + module.invalid( + &stream, + LaunchConfig { + grid_dim: (1, 1, 1), + block_dim: (threads as u32, 1, 1), + shared_mem_bytes: 0, + }, + out.cu_deviceptr() as *mut u32, + mode, + )?; + } + let err = out + .to_host_vec(&stream) + .expect_err("invalid group query must trap"); + assert!(err.to_string().contains("719"), "unexpected error: {err}"); + println!("PASS rejected invalid group case {mode}: {err}"); + return Ok(()); + } + let grid = (4u32, 2u32, 2u32); + for (block, clustered) in [((8, 3, 4), false), ((8, 3, 2), false), ((8, 3, 4), true)] { + let nt = block.0 * block.1 * block.2; + let nb = grid.0 * grid.1 * grid.2; + let out = DeviceBuffer::::zeroed(&stream, (nt * nb) as usize * PHYSICAL)?; + let cfg = LaunchConfig { + grid_dim: grid, + block_dim: block, + shared_mem_bytes: 0, + }; + unsafe { + if clustered { + module.clustered(&stream, cfg, out.cu_deviceptr() as *mut u64)?; + } else { + module.queries(&stream, cfg, out.cu_deviceptr() as *mut u64)?; + } + } + let values = out.to_host_vec(&stream)?; + let c = if clustered { (2, 2, 1) } else { (1, 1, 1) }; + let cb = c.0 * c.1 * c.2; + for bz in 0..grid.2 { + for by in 0..grid.1 { + for bx in 0..grid.0 { + let b = (bz * grid.1 + by) * grid.0 + bx; + let br = ((bz % c.2) * c.1 + by % c.1) * c.0 + bx % c.0; + let cr = ((bz / c.2) * (grid.1 / c.1) + by / c.1) * (grid.0 / c.0) + bx / c.0; + for t in 0..nt { + let nw = nt.div_ceil(32); + let mut expected = [1u64; PHYSICAL]; + let pairs = [ + (0, 1, t % 32, 32), + (0, 2, t, nt), + (1, 2, t / 32, nw), + (0, 3, br * nt + t, cb * nt), + (1, 3, br * nw + t / 32, cb * nw), + (2, 3, br, cb), + (0, 4, b * nt + t, nb * nt), + (1, 4, b * nw + t / 32, nb * nw), + (2, 4, b, nb), + (3, 4, cr, nb / cb), + ]; + for i in 0..5 { + expected[(i * 5 + i) * 2] = 0; + } + for (inner, outer, rank, count) in pairs { + for (s, l) in [(inner, outer), (outer, inner)] { + expected[(s * 5 + l) * 2] = u64::from(rank); + expected[(s * 5 + l) * 2 + 1] = u64::from(count); + } + } + let start = (b * nt + t) as usize * PHYSICAL; + assert_eq!( + &values[start..start + PHYSICAL], + expected, + "block={b}, thread={t}" + ); + } + } + } + } + println!( + "PASS physical grid={grid:?} block={block:?} cluster={c:?}: {} values", + values.len() + ); + } + for block in [(8, 3, 4), (8, 5, 4)] { + let nt = block.0 * block.1 * block.2; + let nb = grid.0 * grid.1 * grid.2; + let out = DeviceBuffer::::zeroed(&stream, (nt * nb) as usize * MAPPED)?; + unsafe { + module.mapped( + &stream, + LaunchConfig { + grid_dim: grid, + block_dim: block, + shared_mem_bytes: 0, + }, + out.cu_deviceptr() as *mut u64, + )?; + } + let values = out.to_host_vec(&stream)?; + for b in 0..nb { + for t in 0..nt { + let lane = t % 32; + let warp = t / 32; + let nw = nt / 32; + let lm = lane < 30; + let wm = warp < nw / 2 * 2; + let expected = [ + lm as u64, + 6, + 5, + if lm { (lane % 6) as u64 } else { u64::MAX }, + if lm { (lane / 6) as u64 } else { u64::MAX }, + wm as u64, + 64, + 2, + (nw / 2) as u64, + if wm { + (warp % 2 * 32 + lane) as u64 + } else { + u64::MAX + }, + if wm { (warp % 2) as u64 } else { u64::MAX }, + if wm { (warp / 2) as u64 } else { u64::MAX }, + (lane % 8) as u64, + 8, + (lane / 8) as u64, + 4, + lane as u64, + 32, + warp as u64, + nw as u64, + ]; + let start = (b * nt + t) as usize * MAPPED; + assert_eq!( + &values[start..start + MAPPED], + expected, + "mapped block={b}, thread={t}" + ); + } + } + println!("PASS mapped block={block:?}: {} values", values.len()); + } + Ok(()) +} From e8709f555d9035fa268366e2297cda285acef71b Mon Sep 17 00:00:00 2001 From: Trent Nelson Date: Mon, 19 Oct 2026 23:29:44 -0700 Subject: [PATCH 08/33] Exercise checked helpers shared by distinct kernel roots Reuse the same helper sites with 64- and 128-thread kernels and retain one device static across their emitted bodies. Check each launch plan and preserve fixed external-symbol diagnostics. Signed-off-by: Trent Nelson --- .../tests/check_collective_rejections.py | 22 +++- .../tests/oxide-fixture/README.md | 13 +- .../tests/oxide-fixture/consumer/Cargo.toml | 1 + .../tests/oxide-fixture/consumer/src/main.rs | 113 ++++++++++++++---- .../tests/oxide-fixture/helper/src/checked.rs | 25 +++- .../tests/run_oxide_fixture.py | 9 +- 6 files changed, 152 insertions(+), 31 deletions(-) diff --git a/rust/cuda-coop-codegen/tests/check_collective_rejections.py b/rust/cuda-coop-codegen/tests/check_collective_rejections.py index 1d3f23f503c6..166ac4f71621 100644 --- a/rust/cuda-coop-codegen/tests/check_collective_rejections.py +++ b/rust/cuda-coop-codegen/tests/check_collective_rejections.py @@ -28,7 +28,9 @@ } original = {key: path.read_text() for key, path in paths.items()} command = record["argv"].copy() -assert "checked" in command, "first run the fixture with --checked" +assert any(feature in command for feature in ["checked", "root-context"]), ( + "first run the fixture with --checked or --root-context" +) command[command.index("run")] = "build" env = os.environ.copy() env.update( @@ -141,7 +143,23 @@ def in_helper(code): original["root"].rsplit("block = (128, 1, 1)", 1) ) }, - "conflicting exact block shapes", + None, + ), + "fixed_export": ( + { + "root": original["root"].replace( + "coop_marker_fixture_helper::roundtrip::<1, false>(input, output, post, valid, mode);", + "exported_helper(input, output, post, valid, mode);", + 1, + ) + + """ +#[unsafe(no_mangle)] +unsafe fn exported_helper(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); } +} +""" + }, + "fixed external symbol", ), "mode": (in_helper(" let mode = mode + (lane & 1) as u32;"), "lane-varying"), "borrow_mutation": ( diff --git a/rust/cuda-coop-codegen/tests/oxide-fixture/README.md b/rust/cuda-coop-codegen/tests/oxide-fixture/README.md index 07f3b0f18834..3d95e7af549b 100644 --- a/rust/cuda-coop-codegen/tests/oxide-fixture/README.md +++ b/rust/cuda-coop-codegen/tests/oxide-fixture/README.md @@ -53,5 +53,14 @@ python3 rust/cuda-coop-codegen/tests/check_collective_rejections.py \ The cases cover varying controls/carriers, different contexts of one helper, participation and unsupported bounds-elision modes. Sources are restored after -each run. The initial checker diagnoses a shared helper site reached under two -different exact block shapes; separate sites remain supported. +each run. Shared helper sites support different exact root shapes; the compiler +checks each root and argument context independently. Fixed external helper symbols +without a root binding remain diagnosed. + +Add `--root-context` instead of `--checked` for 672 native cases. Two additional +64-thread roots call the same item-count 1/4 helpers and marker sites used by +128-thread roots. An intermediate helper retains a second call boundary, and +all copies read one device static. The fixture checks each root's exact shape +and scratch, including 272/1,040 bytes at 64 threads versus 528/2,064 at 128 +threads. Two equal-shape roots also share the original helper Instance. This +route requires the coordinated contextual-helper compiler support. diff --git a/rust/cuda-coop-codegen/tests/oxide-fixture/consumer/Cargo.toml b/rust/cuda-coop-codegen/tests/oxide-fixture/consumer/Cargo.toml index eb5841743397..ff096fc7ba97 100644 --- a/rust/cuda-coop-codegen/tests/oxide-fixture/consumer/Cargo.toml +++ b/rust/cuda-coop-codegen/tests/oxide-fixture/consumer/Cargo.toml @@ -16,3 +16,4 @@ coop-marker-fixture-helper = { path = "../helper" } [features] checked = ["coop-marker-fixture-helper/checked"] +root-context = ["checked"] 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 index 716937b82d51..7b4b06eb74d8 100644 --- a/rust/cuda-coop-codegen/tests/oxide-fixture/consumer/src/main.rs +++ b/rust/cuda-coop-codegen/tests/oxide-fixture/consumer/src/main.rs @@ -77,6 +77,36 @@ mod kernels { coop_marker_fixture_helper::roundtrip::<4, true>(input, output, post, valid, mode); } } + #[cfg(feature = "root-context")] + #[kernel] + #[launch_bounds(256)] + #[launch_contract(block = (64, 1, 1), domain = 3, dynamic_shared = inferred)] + pub unsafe fn smaller_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); + } + } + #[cfg(feature = "root-context")] + #[kernel] + #[launch_bounds(256)] + #[launch_contract(block = (64, 1, 1), domain = 3, dynamic_shared = inferred)] + pub unsafe fn smaller_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); + } + } } fn main() -> Result<(), Box> { @@ -96,15 +126,16 @@ fn main() -> Result<(), Box> { // 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)))?; + ($prepare:ident, $bytes:expr, $alignment:expr, $threads:expr) => {{ + let prepared = + module.$prepare(LaunchConfig3D::auto_shared((1, 1, 1), ($threads, 1, 1)))?; let facts = prepared.function().launch_facts().unwrap(); - assert_eq!(facts.exact_block, Some((128, 1, 1))); + assert_eq!(facts.exact_block, Some(($threads, 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)] { + for shape in [($threads / 2, 2, 1), ($threads * 2, 1, 1)] { assert!(matches!( module.$prepare(LaunchConfig3D::auto_shared((1, 1, 1), shape)), Err(LaunchContractError::BlockShapeMismatch { .. }) @@ -114,27 +145,36 @@ fn main() -> Result<(), Box> { 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)), + module.$prepare(LaunchConfig3D::new((1, 1, 1), ($threads, 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); + check_plan!(prepare_direct_1, 0, 1, 128); + check_plan!(prepare_direct_4, 0, 1, 128); + check_plan!(prepare_transpose_1, 528, 16, 128); + check_plan!(prepare_transpose_4, 2064, 16, 128); + check_plan!(prepare_shared_helper_root, 2064, 16, 128); + #[cfg(feature = "root-context")] + { + check_plan!(prepare_smaller_1, 272, 16, 64); + check_plan!(prepare_smaller_4, 1040, 16, 64); + } let mut cases = 0; - for (kernel, items, transpose) in [ - (0, 1, false), - (1, 4, false), - (2, 1, true), - (3, 4, true), - (4, 4, true), + for (kernel, items, transpose, threads) in [ + (0, 1, false, 128), + (1, 4, false, 128), + (2, 1, true, 128), + (3, 4, true, 128), + (4, 4, true, 128), + #[cfg(feature = "root-context")] + (5, 1, true, 64), + #[cfg(feature = "root-context")] + (6, 4, true, 64), ] { - let tile = 128 * items; + let tile = threads * items; for blocks in [1, 3] { let extent = tile * blocks; let input: Vec = (0..extent).map(|i| i as i32 * 3 - 11).collect(); @@ -143,7 +183,8 @@ fn main() -> Result<(), Box> { 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 config = + LaunchConfig3D::auto_shared((blocks as u32, 1, 1), (threads as u32, 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; @@ -185,7 +226,7 @@ fn main() -> Result<(), Box> { valid as i32, mode, )?, - _ => module.shared_helper_root( + 4 => module.shared_helper_root( &stream, &module.prepare_shared_helper_root(config)?, input_ptr, @@ -194,6 +235,27 @@ fn main() -> Result<(), Box> { valid as i32, mode, )?, + #[cfg(feature = "root-context")] + 5 => module.smaller_1( + &stream, + &module.prepare_smaller_1(config)?, + input_ptr, + output_ptr, + post_ptr, + valid as i32, + mode, + )?, + #[cfg(feature = "root-context")] + 6 => module.smaller_4( + &stream, + &module.prepare_smaller_4(config)?, + input_ptr, + output_ptr, + post_ptr, + valid as i32, + mode, + )?, + _ => unreachable!(), } } let loaded: Vec = input @@ -227,7 +289,7 @@ fn main() -> Result<(), Box> { let lane = (index - origin) / items; let item = index % items; loaded[if transpose { - origin + item * 128 + lane + origin + item * threads + lane } else { index }] @@ -248,9 +310,16 @@ fn main() -> Result<(), Box> { } } } - assert_eq!(cases, 480); + assert_eq!( + cases, + if cfg!(feature = "root-context") { + 672 + } else { + 480 + } + ); println!( - "PASS {cases} native marker GPU cases with inferred scratch through five roots and sixteen provider specializations; {} CUBIN bytes", + "PASS {cases} native marker GPU cases with inferred scratch through shared generic helpers; {} CUBIN bytes", cubin.len() ); Ok(()) diff --git a/rust/cuda-coop-codegen/tests/oxide-fixture/helper/src/checked.rs b/rust/cuda-coop-codegen/tests/oxide-fixture/helper/src/checked.rs index f425c21babda..9899bed7be6a 100644 --- a/rust/cuda-coop-codegen/tests/oxide-fixture/helper/src/checked.rs +++ b/rust/cuda-coop-codegen/tests/oxide-fixture/helper/src/checked.rs @@ -4,7 +4,7 @@ /// Roundtrip through one shared generic helper reached from several kernels. /// /// # Safety -/// Each block owns a disjoint tile. Launch exact 128-thread one-dimensional +/// Each block owns a disjoint tile. Launch exact root-declared 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 @@ -17,11 +17,28 @@ pub unsafe fn roundtrip( valid: i32, mode: u32, ) { + unsafe { + roundtrip_body::(input, output, post_store, valid, mode); + } +} + +// One original static identity must survive emission of this helper per root. +static mut CONTEXT_INCREMENT: i32 = 1; + +#[inline(never)] +unsafe fn roundtrip_body( + input: *const i32, + output: *mut i32, + post_store: *mut i32, + valid: i32, + mode: u32, +) { + let block_threads = cuda_device::thread::blockDim_x() as usize; let lane = cuda_device::thread::threadIdx_x() as usize; - let origin = cuda_device::thread::blockIdx_x() as usize * 128 * ITEMS_PER_THREAD; + let origin = cuda_device::thread::blockIdx_x() as usize * block_threads * ITEMS_PER_THREAD; let valid = (valid - origin as i32) .max(0) - .min((128 * ITEMS_PER_THREAD) as i32); + .min((block_threads * ITEMS_PER_THREAD) as i32); let mut first = [0_i32; ITEMS_PER_THREAD]; let mut second = [0_i32; ITEMS_PER_THREAD]; // Lane-local choice reconverges before every collective. @@ -46,7 +63,7 @@ pub unsafe fn roundtrip( } let mut i = 0; while i < ITEMS_PER_THREAD { - payload[i] += 1; + payload[i] += CONTEXT_INCREMENT; i += 1; } coop_marker_fixture_api::checked_store::( diff --git a/rust/cuda-coop-codegen/tests/run_oxide_fixture.py b/rust/cuda-coop-codegen/tests/run_oxide_fixture.py index a6644ccc769c..52bfb966a211 100644 --- a/rust/cuda-coop-codegen/tests/run_oxide_fixture.py +++ b/rust/cuda-coop-codegen/tests/run_oxide_fixture.py @@ -20,6 +20,11 @@ parser.add_argument( "--checked", action="store_true", help="exercise the checked full-block marker" ) +parser.add_argument( + "--root-context", + action="store_true", + help="exercise shared helpers at 64 and 128 threads", +) args = parser.parse_args() source = Path(__file__).resolve().parent checkout = source.parents[2] @@ -66,7 +71,9 @@ args.arch, "--materialize-cubin", ] -if args.checked: +if args.root_context: + command.extend(["--features", "root-context"]) +elif args.checked: command.extend(["--features", "checked"]) with (output / "build-run.log").open("w") as log: result = subprocess.run( From 2a498e93cdbf2ad1deb98a410fb4a78f9f715367 Mon Sep 17 00:00:00 2001 From: Trent Nelson Date: Mon, 19 Oct 2026 23:39:56 -0700 Subject: [PATCH 09/33] Anchor group queries and retain device payload checks Select fixture output locations from raw CUDA coordinates and exercise payload initialization, access and copying at both item counts. Annotate compile-fail examples and clarify hierarchy ordering in the guide. Signed-off-by: Trent Nelson --- docs/rust/groups.rst | 17 +-- rust/cuda-coop/src/group.rs | 12 +-- .../tests/device_queries/src/main.rs | 101 ++++++++++++++++-- 3 files changed, 110 insertions(+), 20 deletions(-) diff --git a/docs/rust/groups.rst b/docs/rust/groups.rst index 8e411a01c500..4c5050e02914 100644 --- a/docs/rust/groups.rst +++ b/docs/rust/groups.rst @@ -2,7 +2,7 @@ .. SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception Hierarchy and group queries -========================== +=========================== ``cuda-coop`` describes the current CUDA launch without requiring dimensions in the Rust source. ``this_thread()``, ``this_warp()``, ``this_block()``, @@ -23,7 +23,8 @@ Rank and count an explicit level with ``rank_at::()`` or ``count_at::()``. An inner level queries constituents within the group; an outer level queries that group within its parent. Equal levels have -rank zero and count one. Physical queries use x-fastest linearization. +rank zero and count one. Physical queries use x-fastest linearization within +each level, with blocks ordered before their constituent threads or warps. For example, ``this_block().rank()`` is the current thread's linear rank in the block, while ``this_block().rank_at::()`` is the block's linear rank @@ -67,7 +68,10 @@ Run the native query checks The standalone ``rust/cuda-coop/tests/device_queries`` package contains the complete native program. It verifies all physical scope/level pairs with multidimensional blocks and grids, a partial final warp, a clustered launch, -exhaustive and non-exhaustive mappings, and integer conversions. It uses the +exhaustive and non-exhaustive mappings, and integer conversions. Raw CUDA +coordinates select output locations independently of the queries being checked. +The same program checks payload initialization, safe access and copying at +item counts 1 and 4. It uses the same pinned public Oxide dependencies as the library. Keeping this native package separate lets ordinary workspace tests run without CUDA host libraries. @@ -80,15 +84,16 @@ sm_120 GPU, run from the CCCL checkout: CUDA_HOME=/usr/local/cuda-13.3 cargo oxide run --arch sm_120 --materialize-cubin compute-sanitizer --tool memcheck --error-exitcode 1 target/release/device_queries -The four rejection modes run in separate processes because a device trap makes +The five rejection modes run in separate processes because a device trap makes its CUDA context unusable: .. code-block:: shell - for case in 0 1 2 3; do + for case in 0 1 2 3 4; do target/release/device_queries "$case" done They check partial physical warps, an indivisible exhaustive partition, a partition larger than its parent, and a rank query by an excluded lane. The -program reports success only when CUDA returns the expected device-trap error. +last mode indexes an unwritten payload slot. The program reports success only +when CUDA returns the expected device-trap error. diff --git a/rust/cuda-coop/src/group.rs b/rust/cuda-coop/src/group.rs index d3ea2e900b23..f1a1cdf3f5b1 100644 --- a/rust/cuda-coop/src/group.rs +++ b/rust/cuda-coop/src/group.rs @@ -381,11 +381,11 @@ pair!(Cluster, Grid, u64, cluster_rank_grid(), clusters_grid()); /// otherwise trailing lanes are excluded. The block must have complete /// warps. Only Thread/Warp queries are available; groups cannot be nested. /// -/// ```compile_fail +/// ```compile_fail,E0277 /// let group = cuda_coop::this_warp().group_by::<8>(); /// group.rank_at::(); /// ``` -/// ```compile_fail +/// ```compile_fail,E0599 /// cuda_coop::this_warp().group_by::<8>().group_by::<2>(); /// ``` #[derive(Clone, Copy, Debug)] @@ -397,7 +397,7 @@ pub enum ThreadsWithinWarp {} /// `EXHAUSTIVE = true` requires divisibility; otherwise trailing warps are /// excluded. Only Thread/Warp/Block queries are available; no nesting. /// -/// ```compile_fail +/// ```compile_fail,E0277 /// cuda_coop::this_block().group_by::<2>().count_at::(); /// ``` #[derive(Clone, Copy, Debug)] @@ -436,10 +436,10 @@ impl ThreadGroup { /// `N` must be positive and divide 32 (checked at compilation). Panics /// (traps on device) if the block has a partial physical warp. /// - /// ```compile_fail + /// ```compile_fail,E0080 /// cuda_coop::this_warp().group_by::<0>(); /// ``` - /// ```compile_fail + /// ```compile_fail,E0080 /// cuda_coop::this_warp().group_by::<6>(); /// ``` #[inline(always)] @@ -473,7 +473,7 @@ impl ThreadGroup { /// device) unless the block has complete warps and its warp count is /// divisible by `N`. /// - /// ```compile_fail + /// ```compile_fail,E0080 /// cuda_coop::this_block().group_by::<33>(); /// ``` #[inline(always)] diff --git a/rust/cuda-coop/tests/device_queries/src/main.rs b/rust/cuda-coop/tests/device_queries/src/main.rs index 12ce62674ea4..b7188d1def73 100644 --- a/rust/cuda-coop/tests/device_queries/src/main.rs +++ b/rust/cuda-coop/tests/device_queries/src/main.rs @@ -2,11 +2,52 @@ // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception //! Query physical and mapped groups in an Oxide-compiled CUDA kernel. use cuda_coop as coop; -use cuda_device::{cluster_launch, device, kernel}; +use cuda_device::{cluster_launch, device, kernel, thread}; use cuda_host::cuda_module; const PHYSICAL: usize = 55; const MAPPED: usize = 20; +// Output addresses use raw CUDA coordinates, independently of group queries. +#[device] +fn output_index() -> usize { + let block = (u64::from(thread::blockIdx_z()) * u64::from(thread::gridDim_y()) + + u64::from(thread::blockIdx_y())) + * u64::from(thread::gridDim_x()) + + u64::from(thread::blockIdx_x()); + let lane = (thread::threadIdx_z() * thread::blockDim_y() + thread::threadIdx_y()) + * thread::blockDim_x() + + thread::threadIdx_x(); + let threads = thread::blockDim_x() * thread::blockDim_y() * thread::blockDim_z(); + (block * u64::from(threads) + u64::from(lane)) as usize +} + +#[device] +unsafe fn payload_values(output: *mut i32) { + let lane = output_index(); + let selected = lane % ITEMS; + let seed = lane as i32 * 3 - 11; + let mut values = coop::ThreadData::<_, ITEMS>::new(); + for i in 0..ITEMS { + assert!(values.get(i).is_none()); + assert!(values.get_mut(i).is_none()); + } + values.write(selected, seed); + let mut copy = values; + *copy.get_mut(selected).unwrap() += 7; + copy[selected] += 5; + assert_eq!(values[selected], seed); + assert!(copy.get(ITEMS).is_none()); + for i in 0..ITEMS { + values.write(i, seed + i as i32); + unsafe { + output.add(lane * ITEMS * 2 + i).write(values[i]); + output + .add(lane * ITEMS * 2 + ITEMS + i) + .write(copy.get(i).unwrap_or(-999)); + } + } +} + // Check each local query's public return width while converting test output. macro_rules! local { ($query:expr) => {{ @@ -80,7 +121,7 @@ unsafe fn physical(output: *mut u64) { cluster.is_member() as u64, grid.is_member() as u64, ]; - let base = grid.rank() as usize * PHYSICAL; + let base = output_index() * PHYSICAL; for (i, value) in values.into_iter().enumerate() { unsafe { output.add(base + i).write(value); @@ -154,13 +195,23 @@ mod kernels { single.rank_at_as::() as u64, single.count_at_as::(), ]; - let base = coop::this_grid().rank() as usize * MAPPED; + let base = output_index() * MAPPED; for (i, value) in values.into_iter().enumerate() { unsafe { output.add(base + i).write(value); } } } + #[kernel] + pub unsafe fn payload(output: *mut i32, items: u32) { + unsafe { + if items == 1 { + payload_values::<1>(output); + } else { + payload_values::<4>(output); + } + } + } // Each negative case runs in a separate process because a device trap // makes the CUDA context unusable for subsequent launches. #[kernel] @@ -169,10 +220,15 @@ mod kernels { 0 => coop::this_warp().group_by::<8>().count(), 1 => coop::this_block().group_by::<2>().count(), 2 => coop::this_block().group_by_non_exhaustive::<4>().count(), - _ => coop::this_warp().group_by_non_exhaustive::<6>().rank(), + 3 => coop::this_warp().group_by_non_exhaustive::<6>().rank(), + _ => { + let mut data = coop::ThreadData::<_, 4>::new(); + data.write(0, 1_u32); + data[1] + } }; unsafe { - output.add(coop::this_block().rank() as usize).write(value); + output.add(output_index()).write(value); } } } @@ -185,7 +241,7 @@ fn main() -> Result<(), Box> { let module = kernels::load(&ctx)?; if let Some(mode) = std::env::args().nth(1) { let mode: u32 = mode.parse()?; - assert!(mode <= 3); + assert!(mode <= 4); let threads = if mode == 0 { 48 } else { 96 }; let out = DeviceBuffer::::zeroed(&stream, threads)?; unsafe { @@ -202,9 +258,9 @@ fn main() -> Result<(), Box> { } let err = out .to_host_vec(&stream) - .expect_err("invalid group query must trap"); + .expect_err("invalid group query or payload access must trap"); assert!(err.to_string().contains("719"), "unexpected error: {err}"); - println!("PASS rejected invalid group case {mode}: {err}"); + println!("PASS rejected invalid query/payload case {mode}: {err}"); return Ok(()); } let grid = (4u32, 2u32, 2u32); @@ -331,5 +387,34 @@ fn main() -> Result<(), Box> { } println!("PASS mapped block={block:?}: {} values", values.len()); } + for items in [1, 4] { + let block = (8, 3, 4); + let threads = block.0 * block.1 * block.2 * grid.0 * grid.1 * grid.2; + let out = DeviceBuffer::::zeroed(&stream, threads as usize * items * 2)?; + unsafe { + module.payload( + &stream, + LaunchConfig { + grid_dim: grid, + block_dim: block, + shared_mem_bytes: 0, + }, + out.cu_deviceptr() as *mut i32, + items as u32, + )?; + } + let actual = out.to_host_vec(&stream)?; + for lane in 0..threads as usize { + let seed = lane as i32 * 3 - 11; + for i in 0..items { + assert_eq!(actual[lane * items * 2 + i], seed + i as i32); + assert_eq!( + actual[lane * items * 2 + items + i], + if i == lane % items { seed + 12 } else { -999 } + ); + } + } + println!("PASS payload items={items}: {} values", actual.len()); + } Ok(()) } From a9a2c200ca34f9151f6d8159b81b6113a4ecfe55 Mon Sep 17 00:00:00 2001 From: Trent Nelson Date: Mon, 19 Oct 2026 23:59:07 -0700 Subject: [PATCH 10/33] Declare native provider inputs for Cargo rebuilds Signed-off-by: Trent Nelson --- docs/rust/index.rst | 10 ++ rust/cuda-coop-codegen/src/headers.rs | 23 ++-- rust/cuda-coop-codegen/src/nvrtc.rs | 23 ++-- rust/cuda-coop-codegen/src/oxide.rs | 104 +++++++++++++++++- .../tests/oxide-fixture/api/build.rs | 18 ++- .../tests/oxide-fixture/api/src/lib.rs | 2 +- 6 files changed, 156 insertions(+), 24 deletions(-) diff --git a/docs/rust/index.rst b/docs/rust/index.rst index 23776f9573fc..06b7c5d546d2 100644 --- a/docs/rust/index.rst +++ b/docs/rust/index.rst @@ -107,3 +107,13 @@ mutation, and collective requirements. Open ``rust/target/doc/cuda_coop_codegen/ 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; changes that preserve timestamps may +require an explicit rebuild. diff --git a/rust/cuda-coop-codegen/src/headers.rs b/rust/cuda-coop-codegen/src/headers.rs index 7b9105133221..dd55da52337e 100644 --- a/rust/cuda-coop-codegen/src/headers.rs +++ b/rust/cuda-coop-codegen/src/headers.rs @@ -5,7 +5,12 @@ use crate::{Error, Result}; use sha2::{Digest, Sha256}; -use std::{collections::BTreeMap, ffi::CString, fs, path::Path}; +use std::{ + collections::BTreeMap, + ffi::CString, + fs, + path::{Path, PathBuf}, +}; pub(crate) struct Headers { pub(crate) entries: Vec<(CString, CString)>, @@ -13,16 +18,20 @@ pub(crate) struct Headers { } impl Headers { - 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 [ + 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(); diff --git a/rust/cuda-coop-codegen/src/nvrtc.rs b/rust/cuda-coop-codegen/src/nvrtc.rs index 52daf38a400d..fe0ca34e7658 100644 --- a/rust/cuda-coop-codegen/src/nvrtc.rs +++ b/rust/cuda-coop-codegen/src/nvrtc.rs @@ -12,7 +12,7 @@ use sha2::{Digest, Sha256}; use std::{ ffi::{CStr, CString, c_char, c_int, c_void}, fs, - path::Path, + path::{Path, PathBuf}, ptr, }; @@ -56,6 +56,10 @@ pub(crate) struct Output { } 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( @@ -63,12 +67,12 @@ impl Compiler { )); } let mut hash = Sha256::new(); - let mut load = |name: &str| -> Result { - let path = toolkit - .join("lib64") - .join(name) - .canonicalize() - .map_err(|e| Error::input(&toolkit.join("lib64").join(name), e))?; + 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, @@ -79,8 +83,9 @@ impl Compiler { unsafe { Library::new(&path) } .map_err(|e| Error::message(format!("cannot load {}: {e}", path.display()))) }; - let builtins = load("libnvrtc-builtins.so")?; - let library = load("libnvrtc.so")?; + 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 { diff --git a/rust/cuda-coop-codegen/src/oxide.rs b/rust/cuda-coop-codegen/src/oxide.rs index 09d57e186b6b..001a4b55d099 100644 --- a/rust/cuda-coop-codegen/src/oxide.rs +++ b/rust/cuda-coop-codegen/src/oxide.rs @@ -5,7 +5,9 @@ //! //! 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-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`]. @@ -200,6 +202,78 @@ pub fn registration( .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; preserving file timestamps can +/// prevent it from detecting a content change. +/// +/// # 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. @@ -288,6 +362,34 @@ fn write_response(requests: &[Request], artifact: Artifact, output: &mut impl Wr 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( diff --git a/rust/cuda-coop-codegen/tests/oxide-fixture/api/build.rs b/rust/cuda-coop-codegen/tests/oxide-fixture/api/build.rs index 76b2e828a01b..6c1abfb85736 100644 --- a/rust/cuda-coop-codegen/tests/oxide-fixture/api/build.rs +++ b/rust/cuda-coop-codegen/tests/oxide-fixture/api/build.rs @@ -4,12 +4,18 @@ use cuda_coop_codegen::{Algorithm, Operation, oxide}; fn main() { - if std::env::args().nth(1).as_deref() == Some("--oxide-library-request-v1") { - if let Err(error) = oxide::serve( - std::path::Path::new(env!("CCCL_SOURCE")), - std::io::stdin().lock(), - std::io::stdout().lock(), - ) { + 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); } 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 index 5f555bfd6953..5a6d130add0c 100644 --- a/rust/cuda-coop-codegen/tests/oxide-fixture/api/src/lib.rs +++ b/rust/cuda-coop-codegen/tests/oxide-fixture/api/src/lib.rs @@ -7,7 +7,7 @@ macro_rules! registration { ($name:ident, $variable:literal) => { fn $name() -> &'static str { - env!($variable) + cuda_device::library::registration!($variable) } }; } From ed1e7f07cfd2f9dc2fde43f1825bd6ce5ab82035 Mon Sep 17 00:00:00 2001 From: Trent Nelson Date: Tue, 20 Oct 2026 00:10:25 -0700 Subject: [PATCH 11/33] Clarify the scope of Rust payload and group APIs Signed-off-by: Trent Nelson --- docs/rust/index.rst | 4 ++-- rust/cuda-coop/src/group.rs | 2 +- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/docs/rust/index.rst b/docs/rust/index.rst index 6250a0f5cf3a..98d37567acbf 100644 --- a/docs/rust/index.rst +++ b/docs/rust/index.rst @@ -10,8 +10,8 @@ native host compiler foundation for ``cuda.coop``. The ``cuda-coop`` crate suppl and :doc:`groups`. 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. +compiler integration development; collective operations and scratch storage are +still under development. The crate is unpublished and its interface can change. Device code uses the pinned Oxide dependency and ``nightly-2026-10-11``. The host provider supports Rust 1.85 independently; validate it with diff --git a/rust/cuda-coop/src/group.rs b/rust/cuda-coop/src/group.rs index f1a1cdf3f5b1..97004caff679 100644 --- a/rust/cuda-coop/src/group.rs +++ b/rust/cuda-coop/src/group.rs @@ -138,7 +138,7 @@ impl ThreadGroup { /// An inner level selects a constituent's rank in this group; an outer /// level selects this group's rank in that parent. Same-level rank is /// zero. The result is `u64` if Grid is involved, otherwise `u32`. - /// Panics (traps on device) for an excluded mapped-group member. + /// Panics (traps on device) for a thread excluded from a mapped group. /// /// ```no_run /// use cuda_coop as coop; From ee77ca248584a6480ef4e9238b02e7cc79ee3346 Mon Sep 17 00:00:00 2001 From: Trent Nelson Date: Tue, 20 Oct 2026 00:24:03 -0700 Subject: [PATCH 12/33] Clarify native provider timestamp freshness limits Signed-off-by: Trent Nelson --- docs/rust/index.rst | 4 ++-- rust/cuda-coop-codegen/src/oxide.rs | 5 +++-- 2 files changed, 5 insertions(+), 4 deletions(-) diff --git a/docs/rust/index.rst b/docs/rust/index.rst index 06b7c5d546d2..2846d81e9413 100644 --- a/docs/rust/index.rst +++ b/docs/rust/index.rst @@ -115,5 +115,5 @@ 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; changes that preserve timestamps may -require an explicit rebuild. +Cargo's timestamp freshness rules apply; replacements with unchanged or older +timestamps may require an explicit rebuild. diff --git a/rust/cuda-coop-codegen/src/oxide.rs b/rust/cuda-coop-codegen/src/oxide.rs index 001a4b55d099..291e3f90004f 100644 --- a/rust/cuda-coop-codegen/src/oxide.rs +++ b/rust/cuda-coop-codegen/src/oxide.rs @@ -224,8 +224,9 @@ struct InputResponse { /// 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; preserving file timestamps can -/// prevent it from detecting a content change. +/// 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 /// From 2e14aef59caebb81b6c2cf37c985138d9b190d4d Mon Sep 17 00:00:00 2001 From: Trent Nelson Date: Tue, 20 Oct 2026 00:35:55 -0700 Subject: [PATCH 13/33] Qualify collective carrier ownership and full-block reuse Signed-off-by: Trent Nelson --- .../tests/check_carrier_rejections.py | 204 ++++++++++++++++++ .../tests/oxide-fixture/Cargo.toml | 3 +- .../tests/oxide-fixture/README.md | 28 +++ .../tests/oxide-fixture/api/Cargo.toml | 1 + .../tests/oxide-fixture/api/src/carrier.rs | 105 +++++++++ .../tests/oxide-fixture/api/src/lib.rs | 3 + .../tests/oxide-fixture/carrier/Cargo.toml | 15 ++ .../tests/oxide-fixture/carrier/src/main.rs | 152 +++++++++++++ .../tests/run_oxide_fixture.py | 10 +- 9 files changed, 517 insertions(+), 4 deletions(-) create mode 100644 rust/cuda-coop-codegen/tests/check_carrier_rejections.py create mode 100644 rust/cuda-coop-codegen/tests/oxide-fixture/api/src/carrier.rs create mode 100644 rust/cuda-coop-codegen/tests/oxide-fixture/carrier/Cargo.toml create mode 100644 rust/cuda-coop-codegen/tests/oxide-fixture/carrier/src/main.rs diff --git a/rust/cuda-coop-codegen/tests/check_carrier_rejections.py b/rust/cuda-coop-codegen/tests/check_carrier_rejections.py new file mode 100644 index 000000000000..6e63d79b9ace --- /dev/null +++ b/rust/cuda-coop-codegen/tests/check_carrier_rejections.py @@ -0,0 +1,204 @@ +# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. +# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +"""Compile resource/epoch counterexamples in a generated --carrier fixture.""" + +import argparse +import json +import os +import subprocess +from pathlib import Path + +parser = argparse.ArgumentParser(description=__doc__) +parser.add_argument("output", type=Path) +parser.add_argument("--case", action="append", default=[]) +args = parser.parse_args() +output = args.output.resolve() +record = json.loads((output / "command.json").read_text()) +root = output / "fixture/carrier/src/main.rs" +original = root.read_text() +api = output / "fixture/api/src/carrier.rs" +original_api = api.read_text() +command = record["argv"].copy() +command[command.index("run")] = "build" +env = os.environ.copy() +env.update( + CUDA_OXIDE_BACKEND=record["backend"], + CUDA_HOME=record["cuda_root"], + CUDA_PATH=record["cuda_root"], +) +prefix = """// Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception +use coop_marker_fixture_api::carrier; +use cuda_device::{CollectiveTiles, DisjointSlice, kernel, launch_contract, thread}; +type Buffer<'a> = DisjointSlice<'a, i32, CollectiveTiles<4>>; +#[kernel] +#[launch_contract(domain = 1, block = (128, 1, 1), dynamic_shared = inferred)] +pub fn probe<'a>(input: &'a [i32], mut output: Buffer<'a>, mut other: Buffer<'a>, valid: usize) { + let offset = thread::blockIdx_x() as usize * 512; + let mut values = [0_i32; 4]; +""" +suffix = "\n}\nfn main() {}\n" +cases = { + "missing_barrier": ( + "carrier::store(&mut output, offset, valid, &mut values);\n" + "carrier::reload(&mut output, offset, valid, &mut values);", + "requires a full-block barrier", + ), + "loop_backedge": ( + "let mut remaining = 2;\nwhile remaining > 0 {\n" + "carrier::store(&mut output, offset, valid, &mut values);\n" + "remaining -= 1;\n}", + "requires a full-block barrier", + ), + "conditional_barrier": ( + "carrier::store(&mut output, offset, valid, &mut values);\n" + "if valid == 0 { thread::sync_threads(); }\n" + "carrier::reload(&mut output, offset, valid, &mut values);", + "requires a full-block barrier", + ), + "conditional_access": ( + "if valid > 0 { carrier::store(&mut output, offset, valid, &mut values); }\n" + "carrier::reload(&mut output, offset, valid, &mut values);", + "requires a full-block barrier", + ), + "divergent_barrier": ( + "carrier::store(&mut output, offset, valid, &mut values);\n" + "if thread::threadIdx_x() & 1 == 0 { thread::sync_threads(); }", + "barrier is reached through lane-varying", + ), + "varying_count": ( + "carrier::store(&mut output, offset, thread::threadIdx_x() as usize, &mut values);", + "not proved block-uniform", + ), + "varying_carrier": ( + "let selected = if thread::threadIdx_x() & 1 == 0 { &mut output } else { &mut other };\n" + "carrier::store(selected, offset, valid, &mut values);", + "not proved block-uniform", + ), + "ambiguous_origin": ( + "let selected = if valid == 0 { &mut output } else { &mut other };\n" + "carrier::store(selected, offset, valid, &mut values);", + "not proved block-uniform", + ), + "descriptor_mutation": ( + "replace(&mut output, other);\n" + "carrier::store(&mut output, offset, valid, &mut values);", + "mutation through a reference to a resource descriptor", + ), + "readonly_descriptor_mutation": ( + "let mut selected = input;\n" + "redirect_read(&mut selected, input);\n" + "carrier::load(selected, offset, valid, &mut values);", + "mutation through a reference to a resource descriptor", + ), + "descriptor_swap": ( + "carrier::store(&mut other, offset, valid, &mut values);\n" + "core::mem::swap(&mut output, &mut other);\n" + "carrier::store(&mut output, offset, valid, &mut values);", + "intrinsic is outside the supported checked block subset", + ), + "descriptor_copy": ( + "unsafe { core::ptr::copy_nonoverlapping(&other, &mut output, 1); }\n" + "carrier::store(&mut output, offset, valid, &mut values);", + "bulk memory mutation", + ), + "aggregate_reference_mutation": ( + "carrier::store(&mut other, offset, valid, &mut values);\n" + "let mut refs = (&mut output,);\n" + "redirect(&mut refs, &mut other);\n" + "carrier::store(refs.0, offset, valid, &mut values);", + "mutation through a reference to a resource descriptor", + ), + "projected_reborrow": ( + "carrier::store(&mut other, offset, valid, &mut values);\n" + "let mut refs = (&mut output,);\n" + "let target = &mut refs.0;\n*target = &mut other;\n" + "carrier::store(refs.0, offset, valid, &mut values);", + "requires a full-block barrier", + ), + "mixed_marker": ( + "unsafe { coop_marker_fixture_api::store::<4, false>(output.as_mut_ptr(), values.as_mut_ptr(), valid as i32, true); }\n" + "carrier::reload(&mut output, offset, valid, &mut values);", + "untracked and resource-checked", + ), + "scratch_provider": ( + "carrier::store(&mut output, offset, valid, &mut values);", + "require a scratch-free provider", + ), + "lane_reference": ( + "let _ = output.get_mut(thread::index_1d());", + "FlatIndexSpace", + ), + "distinct_allocations": ( + "carrier::store(&mut output, offset, valid, &mut values);\n" + "carrier::store(&mut other, offset, valid, &mut values);", + None, + ), + "same_origin_join": ( + "let selected = if valid == 0 { &mut output } else { &mut output };\n" + "carrier::store(selected, offset, valid, &mut values);", + None, + ), + "helper_barrier": ( + "carrier::store(&mut output, offset, valid, &mut values);\n" + "complete();\ncarrier::reload(&mut output, offset, valid, &mut values);", + None, + ), + "loop_barrier": ( + "let mut remaining = 2;\nwhile remaining > 0 {\n" + "carrier::store(&mut output, offset, valid, &mut values);\n" + "thread::sync_threads();\nremaining -= 1;\n}", + None, + ), +} +helpers = """ +#[inline(never)] +fn complete() { thread::sync_threads(); } +#[inline(never)] +fn replace<'a>(target: &mut Buffer<'a>, value: Buffer<'a>) { *target = value; } +#[inline(never)] +fn redirect<'a, 'b>(target: &mut (&'a mut Buffer<'b>,), value: &'a mut Buffer<'b>) { + target.0 = value; +} +#[inline(never)] +fn redirect_read<'a>(target: &mut &'a [i32], value: &'a [i32]) { *target = value; } +""" +results = {} +unknown = set(args.case) - cases.keys() +if unknown: + parser.error(f"unknown cases: {', '.join(sorted(unknown))}") +try: + for name, (body, diagnostic) in cases.items(): + if args.case and name not in args.case: + continue + root.write_text(prefix + body + suffix + helpers) + api.write_text( + original_api.replace( + "super::direct_store_partial", "super::transpose_store_partial" + ) + if name == "scratch_provider" + else original_api + ) + log_path = output / f"carrier-{name}.log" + with log_path.open("w") as log: + result = subprocess.run( + command, + cwd=record["cwd"], + env=env, + stdout=log, + stderr=subprocess.STDOUT, + ) + log = log_path.read_text() + passed = ( + result.returncode == 0 + if diagnostic is None + else result.returncode != 0 and diagnostic in log + ) + results[name] = {"passed": passed, "returncode": result.returncode} + print(f"{name}: {'PASS' if passed else 'FAIL'}", flush=True) +finally: + root.write_text(original) + api.write_text(original_api) + (output / "carrier-checks.json").write_text(json.dumps(results, indent=2) + "\n") +raise SystemExit(0 if all(item["passed"] for item in results.values()) else 1) diff --git a/rust/cuda-coop-codegen/tests/oxide-fixture/Cargo.toml b/rust/cuda-coop-codegen/tests/oxide-fixture/Cargo.toml index d2b9d7f747af..3278c198b444 100644 --- a/rust/cuda-coop-codegen/tests/oxide-fixture/Cargo.toml +++ b/rust/cuda-coop-codegen/tests/oxide-fixture/Cargo.toml @@ -2,7 +2,7 @@ # SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception [workspace] -members = ["api", "helper", "consumer"] +members = ["api", "helper", "consumer", "carrier"] default-members = ["consumer"] resolver = "3" @@ -10,4 +10,5 @@ resolver = "3" cuda-device = "0.2.1" cuda-host = "0.2.1" cuda-core = "0.4.0" +cuda-async = "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 index 3d95e7af549b..be7a15222acf 100644 --- a/rust/cuda-coop-codegen/tests/oxide-fixture/README.md +++ b/rust/cuda-coop-codegen/tests/oxide-fixture/README.md @@ -64,3 +64,31 @@ all copies read one device static. The fixture checks each root's exact shape and scratch, including 272/1,040 bytes at 64 threads versus 528/2,064 at 128 threads. Two equal-shape roots also share the original helper Instance. This route requires the coordinated contextual-helper compiler support. + +Add `--carrier` to run the separate private scratch-free carrier slice. Its safe +fixture wrappers use `DisjointSlice>` and initialized +lane-local arrays. The compiler tracks the original kernel allocations through +direct helpers, joins and explicit `while` loops, and requires a full-block +barrier before conflicting reuse. The wrapper checks the supplied offset and +extent against a fixed per-block ownership tile; changing an operation's item +count never repartitions that allocation. The carrier exposes no safe lane +reference or conversion to an ordinary slice. + +This bounded route covers only canonical full-block Direct Load/Store. Arbitrary +offsets retain their meaning and must fit the owning tile. The raw unsafe marker +remains available for other ownership mappings. Warp groups, reservations, +explicit scratch, public payload validity, descriptor mutation through references +and unresolved allocation identities are outside this fixture's supported slice. +The numerical checks include shifted offsets, zero/tail/full counts, item counts +1 and 4 in one kernel, Store/barrier/Load cycles and owned async completion. + +```sh +python3 rust/cuda-coop-codegen/tests/check_carrier_rejections.py \ + /path/to/carrier-evidence-directory +``` + +This command checks missing and conditional barriers, loop backedges, divergent +controls, descriptor/reference mutation, ambiguous origins and mixed markers, +with positive cases for independent allocations and barriers in helpers. The +native carrier binary also accepts `overlap` or `bounds`; each invocation must +fail through a device trap, before an invalid native access. diff --git a/rust/cuda-coop-codegen/tests/oxide-fixture/api/Cargo.toml b/rust/cuda-coop-codegen/tests/oxide-fixture/api/Cargo.toml index 450d33d5678e..2dce4d7120fc 100644 --- a/rust/cuda-coop-codegen/tests/oxide-fixture/api/Cargo.toml +++ b/rust/cuda-coop-codegen/tests/oxide-fixture/api/Cargo.toml @@ -15,3 +15,4 @@ cuda-coop-codegen = { path = "../../.." } [features] checked = [] +carrier = [] diff --git a/rust/cuda-coop-codegen/tests/oxide-fixture/api/src/carrier.rs b/rust/cuda-coop-codegen/tests/oxide-fixture/api/src/carrier.rs new file mode 100644 index 000000000000..e46c515f0d61 --- /dev/null +++ b/rust/cuda-coop-codegen/tests/oxide-fixture/api/src/carrier.rs @@ -0,0 +1,105 @@ +// Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +//! Private, scratch-free full-block integration fixture. These wrappers are not +//! the public Load/Store API: the compiler checks participation and resource +//! epochs, while this trusted boundary checks the unchanged offset and extent. +use cuda_device::{CollectiveTiles, DisjointSlice, library, thread}; + +#[inline(never)] +fn extent(len: usize, offset: usize, valid: usize) { + let threads = thread::blockDim_x() as usize + * thread::blockDim_y() as usize + * thread::blockDim_z() as usize; + assert!(ITEMS > 0 && threads <= usize::MAX / ITEMS); + assert!(valid <= threads * ITEMS && valid <= i32::MAX as usize); + assert!(offset <= len && valid <= len - offset); +} + +#[inline(never)] +fn containment(offset: usize, valid: usize) { + let threads = thread::blockDim_x() as usize + * thread::blockDim_y() as usize + * thread::blockDim_z() as usize; + let groups = + thread::gridDim_x() as usize * thread::gridDim_y() as usize * thread::gridDim_z() as usize; + let rank = thread::blockIdx_x() as usize + + thread::gridDim_x() as usize + * (thread::blockIdx_y() as usize + + thread::gridDim_y() as usize * thread::blockIdx_z() as usize); + assert!(CAPACITY > 0 && threads <= usize::MAX / CAPACITY); + let capacity = threads * CAPACITY; + assert!(groups > 0 && groups <= usize::MAX / capacity); + let period = groups * capacity; + if valid != 0 { + let within = offset % period; + let begin = rank * capacity; + assert!(within >= begin && within - begin <= capacity); + assert!(valid <= capacity - (within - begin)); + } +} + +#[inline(never)] +pub fn load( + source: &[i32], + offset: usize, + valid: usize, + values: &mut [i32; ITEMS], +) { + extent::(source.len(), offset, valid); + // SAFETY: filled Load initializes every payload element. Bounds are checked + // before forming the effective pointer. The witness and pointer describe + // the same allocation; the compiler proves full-block control uniformity. + unsafe { + library::checked_block_resource_call_v1::( + super::direct_load_filled, + source, + (source.as_ptr().add(offset), valid as i32, -7_i32), + (values.as_mut_ptr(),), + ); + } +} + +#[inline(never)] +pub fn store( + target: &mut DisjointSlice>, + offset: usize, + valid: usize, + values: &mut [i32; ITEMS], +) { + extent::(target.len(), offset, valid); + containment::(offset, valid); + let pointer = target.as_mut_ptr(); + // SAFETY: this exact carrier owns the checked range. No lane reference to + // it can escape. The payload is initialized, and the compiler checks epochs. + unsafe { + library::checked_block_resource_call_v1::( + super::direct_store_partial, + target, + (pointer.add(offset), valid as i32), + (values.as_mut_ptr(),), + ); + } +} + +#[inline(never)] +pub fn reload( + source: &mut DisjointSlice>, + offset: usize, + valid: usize, + values: &mut [i32; ITEMS], +) { + extent::(source.len(), offset, valid); + containment::(offset, valid); + let pointer = source.as_mut_ptr(); + // SAFETY: immutable witness and effective pointer have the same origin. + // The compiler requires a full-block barrier after any prior write. + unsafe { + library::checked_block_resource_call_v1::( + super::direct_load_filled, + &*source, + (pointer.add(offset).cast_const(), valid as i32, -7_i32), + (values.as_mut_ptr(),), + ); + } +} 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 index f0f464a2c760..2f51b0c273b5 100644 --- a/rust/cuda-coop-codegen/tests/oxide-fixture/api/src/lib.rs +++ b/rust/cuda-coop-codegen/tests/oxide-fixture/api/src/lib.rs @@ -4,6 +4,9 @@ //! Private unsafe fixture wrappers, with no public collective API contract. #![no_std] +#[cfg(feature = "carrier")] +pub mod carrier; + macro_rules! registration { ($name:ident, $variable:literal) => { fn $name() -> &'static str { diff --git a/rust/cuda-coop-codegen/tests/oxide-fixture/carrier/Cargo.toml b/rust/cuda-coop-codegen/tests/oxide-fixture/carrier/Cargo.toml new file mode 100644 index 000000000000..fe068a69bb79 --- /dev/null +++ b/rust/cuda-coop-codegen/tests/oxide-fixture/carrier/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-carrier-fixture" +version = "0.0.0" +edition = "2024" +publish = false + +[dependencies] +cuda-device.workspace = true +cuda-host = { workspace = true, features = ["async"] } +cuda-core.workspace = true +cuda-async.workspace = true +coop-marker-fixture-api = { path = "../api", features = ["carrier"] } diff --git a/rust/cuda-coop-codegen/tests/oxide-fixture/carrier/src/main.rs b/rust/cuda-coop-codegen/tests/oxide-fixture/carrier/src/main.rs new file mode 100644 index 000000000000..2b39d7138a89 --- /dev/null +++ b/rust/cuda-coop-codegen/tests/oxide-fixture/carrier/src/main.rs @@ -0,0 +1,152 @@ +// Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +use coop_marker_fixture_api::carrier; +use cuda_device::{CollectiveTiles, DisjointSlice, kernel, launch_contract, thread}; +use cuda_host::cuda_module; + +#[cuda_module] +mod kernels { + use super::*; + + #[inline(never)] + fn cycle( + input: &[i32], + output: &mut DisjointSlice>, + offset: usize, + valid: usize, + ) { + let mut values = [0; ITEMS]; + carrier::load(input, offset, valid, &mut values); + carrier::store(output, offset, valid, &mut values); + thread::sync_threads(); + carrier::reload(output, offset, valid, &mut values); + thread::sync_threads(); + carrier::store(output, offset, valid, &mut values); + thread::sync_threads(); + } + + #[inline(never)] + fn repeat( + input: &[i32], + output: &mut DisjointSlice>, + valid: usize, + shift: usize, + rounds: u32, + ) { + let offset = thread::blockIdx_x() as usize * 128 * 4 + shift; + let mut remaining = rounds; + while remaining > 0 { + cycle::(input, output, offset, valid); + remaining -= 1; + } + } + + #[kernel] + #[launch_contract(domain = 1, block = (128, 1, 1), dynamic_shared = inferred)] + pub fn copy_one( + input: &[i32], + mut output: DisjointSlice>, + valid: usize, + shift: usize, + rounds: u32, + ) { + repeat::<1>(input, &mut output, valid, shift, rounds); + } + + #[kernel] + #[launch_contract(domain = 1, block = (128, 1, 1), dynamic_shared = inferred)] + pub fn copy_four( + input: &[i32], + mut output: DisjointSlice>, + valid: usize, + shift: usize, + rounds: u32, + ) { + repeat::<4>(input, &mut output, valid, shift, rounds); + } + #[kernel] + #[launch_contract(domain = 1, block = (128, 1, 1), dynamic_shared = inferred)] + pub fn changed_items(input: &[i32], mut output: DisjointSlice>) { + repeat::<1>(input, &mut output, 128, 0, 1); + let offset = thread::blockIdx_x() as usize * 512; + let mut values = [0; 4]; + carrier::load(input, offset, 512, &mut values); + carrier::store(&mut output, offset, 512, &mut values); + // No final barrier: returning the owned buffers waits for this kernel. + } +} + +fn main() -> Result<(), Box> { + use cuda_async::simt::device_context::{init_device_contexts, with_cuda_context}; + use cuda_async::simt::device_operation::DeviceOperation; + use cuda_core::{DeviceBuffer, LaunchConfig1D}; + + init_device_contexts(0, 1)?; + let context = with_cuda_context(0, |context| context.clone())?; + let stream = context.default_stream(); + // SAFETY: this module contains the declarations above. + let module = unsafe { kernels::load_async(0)? }; + let config = LaunchConfig1D::auto_shared(3, 128); + let one = module.prepare_copy_one(config)?; + let four = module.prepare_copy_four(config)?; + assert_eq!(one.__raw_config().shared_mem_bytes, 0); + let input: Vec<_> = (0..1536).map(|i| i * 3 + 17).collect(); + if let Some(mode) = std::env::args().nth(1) { + // All shifted ranges fit the allocation in the overlap case. Only + // canonical group containment forbids the adjacent-block overlap. + let source = DeviceBuffer::from_host(&stream, &vec![1_i32; 2048])?; + let target = DeviceBuffer::from_host(&stream, &vec![0_i32; 2048])?; + let shift = if mode == "overlap" { 1 } else { 2048 }; + let _ = module + .copy_four_async_owned(&four, source, target, 512, shift, 1) + .sync()?; + panic!("invalid collective range did not trap"); + } + for items in [1, 4] { + for shift in [0, 17] { + for valid in [0, 1, items * 128 - 1, items * 128] { + if valid > 512 - shift { + continue; + } + let source = DeviceBuffer::from_host(&stream, &input)?; + let target = DeviceBuffer::from_host(&stream, &vec![-99; input.len()])?; + let (_source, target) = if items == 1 { + module + .copy_one_async_owned(&one, source, target, valid, shift, 2) + .sync()? + } else { + module + .copy_four_async_owned(&four, source, target, valid, shift, 2) + .sync()? + }; + let mut result = vec![0; input.len()]; + target.copy_to_host(&stream, &mut result)?; + stream.synchronize()?; + for (index, value) in result.into_iter().enumerate() { + let expected = if index % 512 >= shift && index % 512 - shift < valid { + input[index] + } else { + -99 + }; + assert_eq!( + value, expected, + "items={items}, valid={valid}, index={index}" + ); + } + } + } + } + let changed = module.prepare_changed_items(config)?; + let source = DeviceBuffer::from_host(&stream, &input)?; + let target = DeviceBuffer::from_host(&stream, &vec![-99; input.len()])?; + let (_source, target) = module + .changed_items_async_owned(&changed, source, target) + .sync()?; + let mut result = vec![0; input.len()]; + target.copy_to_host(&stream, &mut result)?; + stream.synchronize()?; + assert_eq!(result, input); + println!("PASS: carrier bounds, item counts, repeated Store/barrier/Load, owned async"); + Ok(()) +} diff --git a/rust/cuda-coop-codegen/tests/run_oxide_fixture.py b/rust/cuda-coop-codegen/tests/run_oxide_fixture.py index 52bfb966a211..c2899825874b 100644 --- a/rust/cuda-coop-codegen/tests/run_oxide_fixture.py +++ b/rust/cuda-coop-codegen/tests/run_oxide_fixture.py @@ -25,6 +25,9 @@ action="store_true", help="exercise shared helpers at 64 and 128 threads", ) +parser.add_argument( + "--carrier", action="store_true", help="exercise private collective-only carriers" +) args = parser.parse_args() source = Path(__file__).resolve().parent checkout = source.parents[2] @@ -48,7 +51,7 @@ stream.write( f"{name} = {{ path = {json.dumps(str(oxide / 'cuda-oxide/crates' / name))} }}\n" ) - for name in ["cuda-core", "cuda-core-derive", "cuda-bindings"]: + for name in ["cuda-core", "cuda-core-derive", "cuda-bindings", "cuda-async"]: 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. @@ -75,15 +78,16 @@ command.extend(["--features", "root-context"]) elif args.checked: command.extend(["--features", "checked"]) +consumer = fixture / ("carrier" if args.carrier else "consumer") with (output / "build-run.log").open("w") as log: result = subprocess.run( - command, cwd=fixture / "consumer", env=env, stdout=log, stderr=subprocess.STDOUT + command, cwd=consumer, env=env, stdout=log, stderr=subprocess.STDOUT ) (output / "command.json").write_text( json.dumps( { "argv": command, - "cwd": str(fixture / "consumer"), + "cwd": str(consumer), "oxide_root": str(oxide), "backend": env["CUDA_OXIDE_BACKEND"], "cuda_root": env["CUDA_HOME"], From e3ba38952c8a2925617f31b4a4df5feb9b427e1a Mon Sep 17 00:00:00 2001 From: Trent Nelson Date: Tue, 20 Oct 2026 00:48:09 -0700 Subject: [PATCH 14/33] Cover checked helper projections and rooted contexts Signed-off-by: Trent Nelson --- .../tests/check_collective_rejections.py | 78 ++++++++++++++++++- .../tests/oxide-fixture/README.md | 21 ++++- .../compile_cases/atomic_control.rs | 34 ++++++++ .../tests/oxide-fixture/consumer/src/main.rs | 24 +++++- .../tests/oxide-fixture/helper/src/checked.rs | 11 ++- .../tests/oxide-fixture/helper/src/lib.rs | 15 ++++ 6 files changed, 172 insertions(+), 11 deletions(-) create mode 100644 rust/cuda-coop-codegen/tests/oxide-fixture/compile_cases/atomic_control.rs diff --git a/rust/cuda-coop-codegen/tests/check_collective_rejections.py b/rust/cuda-coop-codegen/tests/check_collective_rejections.py index 166ac4f71621..d67393be7fe2 100644 --- a/rust/cuda-coop-codegen/tests/check_collective_rejections.py +++ b/rust/cuda-coop-codegen/tests/check_collective_rejections.py @@ -137,6 +137,14 @@ def in_helper(code): {"root": original["root"].replace("block = (128, 1, 1), ", "")}, "requires #[launch_contract(block", ), + "two_dimensional_shape": ( + { + "root": original["root"].replace( + "block = (128, 1, 1)", "block = (64, 2, 1)", 1 + ) + }, + "exact block shape [B,1,1]", + ), "shared_shape": ( { "root": "block = (64, 1, 1)".join( @@ -173,6 +181,48 @@ def in_helper(code): }, "not proved block-uniform", ), + "aggregate_reference": ( + in_helper( + " let mut changed = valid;\n" + " let holder = (&mut changed,);\n" + " *holder.0 = lane as i32;\n" + " let valid = changed;" + ), + "not proved block-uniform", + ), + "projected_reborrow": ( + in_helper( + " let mut pair = (lane as i32, valid);\n" + " let mut pointer = &mut pair;\n" + " if mode & 8 != 0 { pointer = &mut pair; }\n" + " let other = &mut pointer.1;\n" + " *other = 0;\n" + " let valid = pair.0;" + ), + "not proved block-uniform", + ), + "reconvergent_loop": ( + { + "helper": original["helper"].replace( + "payload[i] += CONTEXT_INCREMENT;", + "if lane & 1 == 0 { payload[i] += CONTEXT_INCREMENT; }" + " else { payload[i] += 2; }", + ) + }, + None, + ), + "exhaustive_match": ( + { + "helper": in_helper( + " let choice = if mode & 8 == 0 { Some(1) } else { None };\n" + " let adjustment = match choice { Some(value) => value, None => 0 };" + )["helper"].replace( + "payload[i] += CONTEXT_INCREMENT;", + "payload[i] += CONTEXT_INCREMENT + adjustment;", + ) + }, + None, + ), "destructor": ( { "helper": original["helper"].replace( @@ -207,7 +257,7 @@ def in_helper(code): "unchecked_bounds": ({}, "unchecked indexing"), } probe_api = original["api"].replace("fn $name()", "pub fn $name()") -for name in ["native_mutation", "nested_reference", "nested_uniform"]: +for name in ["native_mutation", "nested_reference", "nested_uniform", "atomic_control"]: source_name = "nested_reference" if name == "nested_uniform" else name probe = ( Path(__file__).parent / "oxide-fixture/compile_cases" / f"{source_name}.rs" @@ -218,6 +268,14 @@ def in_helper(code): {"root": probe, "api": probe_api}, None if name == "nested_uniform" else "not proved block-uniform", ) +cases["reborrow_uniform"] = ( + { + "helper": cases["projected_reborrow"][0]["helper"].replace( + "let valid = pair.0;", "let valid = pair.1;" + ) + }, + None, +) unknown = set(args.case) - cases.keys() if unknown: parser.error(f"unknown cases: {', '.join(sorted(unknown))}") @@ -236,7 +294,12 @@ def in_helper(code): else "" ) ) - if name in ["native_mutation", "nested_reference", "nested_uniform"]: + if name in [ + "native_mutation", + "nested_reference", + "nested_uniform", + "atomic_control", + ]: helper_lib.write_text( helper_lib_original + "\npub use coop_marker_fixture_api::{direct_load_full, direct_load_filled};\n" @@ -244,6 +307,17 @@ def in_helper(code): case_env = env.copy() if name == "unchecked_bounds": case_env["CUDA_OXIDE_UNCHECKED_INDEXING"] = "1" + if name in [ + "aggregate_reference", + "projected_reborrow", + "reborrow_uniform", + "exhaustive_match", + ]: + # Preserve reference projections and exhaustive-match branches + # before the checker sees them. + case_env["RUSTFLAGS"] = ( + case_env.get("RUSTFLAGS", "") + " -Zmir-opt-level=0" + ).strip() log_path = output / f"reject-{name}.log" with log_path.open("w") as log: result = subprocess.run( diff --git a/rust/cuda-coop-codegen/tests/oxide-fixture/README.md b/rust/cuda-coop-codegen/tests/oxide-fixture/README.md index 3d95e7af549b..c47b65d48f90 100644 --- a/rust/cuda-coop-codegen/tests/oxide-fixture/README.md +++ b/rust/cuda-coop-codegen/tests/oxide-fixture/README.md @@ -4,7 +4,7 @@ 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. +concrete kernel roots (seven with `--root-context`). No consumer supplies an LTO artifact list. From the CCCL checkout, with the agreed Oxide source and backend already built: @@ -52,15 +52,28 @@ python3 rust/cuda-coop-codegen/tests/check_collective_rejections.py \ ``` The cases cover varying controls/carriers, different contexts of one helper, -participation and unsupported bounds-elision modes. Sources are restored after +projected references, atomic results, participation, and unsupported +bounds-elision modes. Positive contrasts retain reconvergent lane work in a +uniform loop and an exhaustive enum match. Sources are restored after each run. Shared helper sites support different exact root shapes; the compiler checks each root and argument context independently. Fixed external helper symbols -without a root binding remain diagnosed. +(`#[no_mangle]` or `#[export_name]`) are unsupported on checked paths, including +single-root paths. Remove those attributes to permit root-specific naming. Add `--root-context` instead of `--checked` for 672 native cases. Two additional 64-thread roots call the same item-count 1/4 helpers and marker sites used by -128-thread roots. An intermediate helper retains a second call boundary, and +128-thread roots. The outer generic helper uses `#[device]`; a nongeneric +dependency bridge also reaches it at both block sizes. Checked helpers defer +device emission until a kernel supplies its context. An intermediate +helper retains another call boundary, and all copies read one device static. The fixture checks each root's exact shape and scratch, including 272/1,040 bytes at 64 threads versus 528/2,064 at 128 threads. Two equal-shape roots also share the original helper Instance. This route requires the coordinated contextual-helper compiler support. + +The checked executable also accepts `--trap-before` and `--trap-after-load`. +Run each in a separate process: lane zero deliberately indexes past its payload, +before the first Load or between that Load and its scratch-reuse barrier. +Both runs must terminate with a CUDA execution error, without the success +message. These negative modes check the compiler's bounds-failure abort path; +the normal numerical run uses neither mode. diff --git a/rust/cuda-coop-codegen/tests/oxide-fixture/compile_cases/atomic_control.rs b/rust/cuda-coop-codegen/tests/oxide-fixture/compile_cases/atomic_control.rs new file mode 100644 index 000000000000..9646c1c74a30 --- /dev/null +++ b/rust/cuda-coop-codegen/tests/oxide-fixture/compile_cases/atomic_control.rs @@ -0,0 +1,34 @@ +// Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +use coop_marker_fixture_helper::direct_load_filled; +use cuda_device::atomic::{AtomicOrdering, DeviceAtomicU32}; +use cuda_device::{kernel, launch_contract}; +use cuda_host::cuda_module; + +#[cuda_module] +mod kernels { + use super::*; + + #[kernel] + #[launch_contract(block = (128, 1, 1), domain = 3, dynamic_shared = inferred)] + pub unsafe fn probe(input: *const i32, counter: *const DeviceAtomicU32, flag: bool) { + // A panic-only host placeholder still returns a varying device value. + // Merging it with a constant branch must not prove uniformity. + let count = if flag { + unsafe { (*counter).fetch_add(1, AtomicOrdering::Relaxed) as i32 } + } else { + 0 + }; + let mut value = 0_i32; + unsafe { + cuda_device::library::checked_block_call_v1::( + direct_load_filled, + (input, count, -7), + (&mut value as *mut i32,), + ); + } + } +} + +fn main() {} 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 index 7b4b06eb74d8..fb2f5eff8a13 100644 --- a/rust/cuda-coop-codegen/tests/oxide-fixture/consumer/src/main.rs +++ b/rust/cuda-coop-codegen/tests/oxide-fixture/consumer/src/main.rs @@ -74,7 +74,7 @@ mod kernels { mode: u32, ) { unsafe { - coop_marker_fixture_helper::roundtrip::<4, true>(input, output, post, valid, mode); + coop_marker_fixture_helper::shared_roundtrip(input, output, post, valid, mode); } } #[cfg(feature = "root-context")] @@ -104,7 +104,7 @@ mod kernels { mode: u32, ) { unsafe { - coop_marker_fixture_helper::roundtrip::<4, true>(input, output, post, valid, mode); + coop_marker_fixture_helper::shared_roundtrip(input, output, post, valid, mode); } } } @@ -117,10 +117,21 @@ fn main() -> Result<(), Box> { let bundles = read_artifact_bundles_from_object_bytes(&executable)?; let bundle = bundles .iter() - .find(|b| b.payload(ArtifactPayloadKind::Cubin).is_some()) + .find(|b| { + b.entry("transpose_4").is_some() && 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 trap_mode = match std::env::args().nth(1).as_deref() { + None => 0, + Some("--trap-before") => 8, + Some("--trap-after-load") => 16, + Some(_) => return Err("expected --trap-before or --trap-after-load".into()), + }; + if trap_mode != 0 && !cfg!(feature = "checked") { + return Err("trap modes require the checked fixture".into()); + } let context = CudaContext::new(0)?; let stream = context.default_stream(); // SAFETY: the embedded artifact is generated from these kernel declarations. @@ -135,7 +146,11 @@ fn main() -> Result<(), Box> { 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 [($threads / 2, 2, 1), ($threads * 2, 1, 1)] { + for shape in [ + ($threads / 2, 1, 1), + ($threads / 2, 2, 1), + ($threads * 2, 1, 1), + ] { assert!(matches!( module.$prepare(LaunchConfig3D::auto_shared((1, 1, 1), shape)), Err(LaunchContractError::BlockShapeMismatch { .. }) @@ -181,6 +196,7 @@ fn main() -> Result<(), Box> { 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 mode = mode | trap_mode; let output = DeviceBuffer::from_host(&stream, &vec![-99_i32; extent])?; let post = DeviceBuffer::from_host(&stream, &vec![-99_i32; extent])?; let config = diff --git a/rust/cuda-coop-codegen/tests/oxide-fixture/helper/src/checked.rs b/rust/cuda-coop-codegen/tests/oxide-fixture/helper/src/checked.rs index 9899bed7be6a..9dcf388dc50d 100644 --- a/rust/cuda-coop-codegen/tests/oxide-fixture/helper/src/checked.rs +++ b/rust/cuda-coop-codegen/tests/oxide-fixture/helper/src/checked.rs @@ -9,7 +9,7 @@ /// 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)] +#[cuda_device::device] pub unsafe fn roundtrip( input: *const i32, output: *mut i32, @@ -41,6 +41,12 @@ unsafe fn roundtrip_body( .min((block_threads * ITEMS_PER_THREAD) as i32); let mut first = [0_i32; ITEMS_PER_THREAD]; let mut second = [0_i32; ITEMS_PER_THREAD]; + // Negative native modes: exactly lane zero fails a checked index, before + // the collective sequence or between its first Load and reuse barrier. + let trap_index = if lane == 0 { ITEMS_PER_THREAD } else { 0 }; + if mode & 8 != 0 { + first[trap_index] = 0; + } // Lane-local choice reconverges before every collective. let payload = if lane & 1 == 0 { &mut first @@ -58,6 +64,9 @@ unsafe fn roundtrip_body( valid, mode & 1 != 0, ); + if mode & 16 != 0 { + payload[trap_index] = 0; + } cuda_device::thread::sync_threads(); iteration += 1; } 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 index 019ca58eefcc..8f0145b72ba3 100644 --- a/rust/cuda-coop-codegen/tests/oxide-fixture/helper/src/lib.rs +++ b/rust/cuda-coop-codegen/tests/oxide-fixture/helper/src/lib.rs @@ -3,6 +3,21 @@ #![no_std] +/// Preserve a nongeneric dependency boundary above the marked generic helper. +/// +/// # Safety +/// The caller must meet `roundtrip::<4, true>`'s complete safety contract. +#[inline(never)] +pub unsafe fn shared_roundtrip( + input: *const i32, + output: *mut i32, + post_store: *mut i32, + valid: i32, + mode: u32, +) { + unsafe { roundtrip::<4, true>(input, output, post_store, valid, mode) } +} + /// Roundtrip through one shared generic helper reached from several kernels. /// /// # Safety From bfadda28c95f6887222736fc8b2b65e5b222a894 Mon Sep 17 00:00:00 2001 From: Trent Nelson Date: Tue, 20 Oct 2026 08:28:29 -0700 Subject: [PATCH 15/33] Cover carrier ownership and aggregate mutation boundaries Signed-off-by: Trent Nelson --- .../tests/check_carrier_rejections.py | 126 +++++++++++++++++- .../tests/oxide-fixture/README.md | 14 +- .../tests/oxide-fixture/carrier/src/main.rs | 70 ++++++++-- 3 files changed, 188 insertions(+), 22 deletions(-) diff --git a/rust/cuda-coop-codegen/tests/check_carrier_rejections.py b/rust/cuda-coop-codegen/tests/check_carrier_rejections.py index 6e63d79b9ace..4a982dc62222 100644 --- a/rust/cuda-coop-codegen/tests/check_carrier_rejections.py +++ b/rust/cuda-coop-codegen/tests/check_carrier_rejections.py @@ -117,6 +117,66 @@ "carrier::store(refs.0, offset, valid, &mut values);", "requires a full-block barrier", ), + "joined_owner_replacement": ( + "let mut first = Owner { buf: output };\n" + "let mut second = Owner { buf: other };\n" + "let mut replacement = Owner { buf: third };\n" + "let mut one = [7_i32; 1];\n" + "carrier::store(&mut replacement.buf, offset, 4, &mut one);\n" + "let selected = if choose { &mut first } else { &mut second };\n" + "*selected = replacement;\n" + "carrier::store(&mut first.buf, offset, 4, &mut values);", + "mutation through a reference to a resource descriptor", + ), + "joined_owner_helper": ( + "let mut first = Owner { buf: output };\n" + "let mut second = Owner { buf: other };\n" + "let mut replacement = Owner { buf: third };\n" + "let mut one = [7_i32; 1];\n" + "carrier::store(&mut replacement.buf, offset, 4, &mut one);\n" + "let selected = if choose { &mut first } else { &mut second };\n" + "replace_owner(selected, replacement);\n" + "carrier::store(&mut first.buf, offset, 4, &mut values);", + "mutation through a reference to a resource descriptor", + ), + "tuple_owner_replacement": ( + "let mut first = Owner { buf: output };\n" + "let mut replacement = Owner { buf: other };\n" + "let mut one = [7_i32; 1];\n" + "carrier::store(&mut replacement.buf, offset, 4, &mut one);\n" + "let references = (&mut first,);\n" + "*references.0 = replacement;\n" + "carrier::store(&mut first.buf, offset, 4, &mut values);", + "mutation through a reference to a resource descriptor", + ), + "warp_barrier": ( + "carrier::store(&mut output, offset, valid, &mut values);\n" + "cuda_device::warp::sync_mask(u32::MAX);\n" + "carrier::reload(&mut output, offset, valid, &mut values);", + "intrinsic is outside the supported checked block subset", + ), + "changed_items_without_barrier": ( + "let mut one = [0_i32; 1];\n" + "carrier::store(&mut output, offset, 128, &mut one);\n" + "carrier::store(&mut output, offset, 512, &mut values);", + "requires a full-block barrier", + ), + "slice_write_witness": ( + "carrier::load(input, offset, valid, &mut values);", + "resource witness requires DisjointSlice>", + ), + "immutable_write_witness": ( + "carrier::store(&mut output, offset, valid, &mut values);", + "resource writes require a mutable carrier witness", + ), + "wrong_element_witness": ( + "carrier::store(&mut output, offset, valid, &mut values);", + "resource witness element type disagrees with the library element type", + ), + "zero_capacity": ( + "carrier::store(&mut output, offset, valid, &mut values);", + "resource carriers require positive capacity", + ), "mixed_marker": ( "unsafe { coop_marker_fixture_api::store::<4, false>(output.as_mut_ptr(), values.as_mut_ptr(), valid as i32, true); }\n" "carrier::reload(&mut output, offset, valid, &mut values);", @@ -135,6 +195,20 @@ "carrier::store(&mut other, offset, valid, &mut values);", None, ), + "aggregate_forward": ( + "let owned = Owner { buf: output };\n" + "let mut moved = owned;\n" + "carrier::store(&mut moved.buf, offset, valid, &mut values);\n" + "thread::sync_threads();\n" + "carrier::reload(&mut moved.buf, offset, valid, &mut values);", + None, + ), + "payload_reference_write": ( + "let references = (&mut values,);\n" + "references.0[0] = 37;\n" + "carrier::store(&mut output, offset, valid, &mut values);", + None, + ), "same_origin_join": ( "let selected = if valid == 0 { &mut output } else { &mut output };\n" "carrier::store(selected, offset, valid, &mut values);", @@ -153,6 +227,9 @@ ), } helpers = """ +struct Owner<'a> { buf: Buffer<'a> } +#[inline(never)] +fn replace_owner<'a>(target: &mut Owner<'a>, value: Owner<'a>) { *target = value; } #[inline(never)] fn complete() { thread::sync_threads(); } #[inline(never)] @@ -172,20 +249,55 @@ for name, (body, diagnostic) in cases.items(): if args.case and name not in args.case: continue - root.write_text(prefix + body + suffix + helpers) - api.write_text( - original_api.replace( + case_prefix = prefix + if name.startswith("joined_owner"): + case_prefix = prefix.replace( + "valid: usize)", "valid: usize, third: Buffer<'a>, choose: bool)" + ) + if name == "zero_capacity": + case_prefix = prefix.replace("CollectiveTiles<4>", "CollectiveTiles<0>") + root.write_text(case_prefix + body + suffix + helpers) + case_api = original_api + if name == "scratch_provider": + case_api = original_api.replace( "super::direct_store_partial", "super::transpose_store_partial" ) - if name == "scratch_provider" - else original_api - ) + elif name == "slice_write_witness": + case_api = original_api.replace("ITEMS, false>", "ITEMS, true>") + elif name == "immutable_write_witness": + case_api = original_api.replace( + " target,", " &*target," + ) + elif name == "wrong_element_witness": + case_api = ( + original_api.replace( + "::", + "::", + ) + .replace( + "(pointer.add(offset), valid as i32)", + "(pointer.add(offset).cast::(), valid as i32)", + ) + .replace( + "(values.as_mut_ptr(),)", "(values.as_mut_ptr().cast::(),)" + ) + ) + elif name == "zero_capacity": + # Keep the marker reachable: this case tests structural witness + # validation, independently of the wrapper's runtime containment. + case_api = original_api.replace( + " containment::(offset, valid);\n", "" + ) + api.write_text(case_api) log_path = output / f"carrier-{name}.log" + # Each program gets its own artifacts, including wrapper dependencies. + case_env = env.copy() + case_env["CARGO_TARGET_DIR"] = str(output / "carrier-targets" / name) with log_path.open("w") as log: result = subprocess.run( command, cwd=record["cwd"], - env=env, + env=case_env, stdout=log, stderr=subprocess.STDOUT, ) diff --git a/rust/cuda-coop-codegen/tests/oxide-fixture/README.md b/rust/cuda-coop-codegen/tests/oxide-fixture/README.md index cb3e919abe50..8ae3934937a1 100644 --- a/rust/cuda-coop-codegen/tests/oxide-fixture/README.md +++ b/rust/cuda-coop-codegen/tests/oxide-fixture/README.md @@ -94,6 +94,8 @@ explicit scratch, public payload validity, descriptor mutation through reference and unresolved allocation identities are outside this fixture's supported slice. The numerical checks include shifted offsets, zero/tail/full counts, item counts 1 and 4 in one kernel, Store/barrier/Load cycles and owned async completion. +The shifted four-item tails reach the fixed capacity boundary. A 2-by-2 grid +also exercises the x-fastest block rank used by canonical containment. ```sh python3 rust/cuda-coop-codegen/tests/check_carrier_rejections.py \ @@ -101,7 +103,11 @@ python3 rust/cuda-coop-codegen/tests/check_carrier_rejections.py \ ``` This command checks missing and conditional barriers, loop backedges, divergent -controls, descriptor/reference mutation, ambiguous origins and mixed markers, -with positive cases for independent allocations and barriers in helpers. The -native carrier binary also accepts `overlap` or `bounds`; each invocation must -fail through a device trap, before an invalid native access. +controls, descriptor/reference mutation (including enclosing aggregates), +ambiguous origins, wrong witnesses and mixed markers, with positive cases for +independent allocations, moved aggregates and barriers in helpers. Warp sync +remains outside this full-block subset and cannot substitute for a block +barrier. The native carrier binary also accepts `overlap`, `bounds` or +`repartition`; each invocation must fail through a device trap. `repartition` +uses an in-bounds offset computed from the smaller operation item count, which +would incorrectly change ownership if the fixed carrier capacity were ignored. diff --git a/rust/cuda-coop-codegen/tests/oxide-fixture/carrier/src/main.rs b/rust/cuda-coop-codegen/tests/oxide-fixture/carrier/src/main.rs index 2b39d7138a89..394c0da136c9 100644 --- a/rust/cuda-coop-codegen/tests/oxide-fixture/carrier/src/main.rs +++ b/rust/cuda-coop-codegen/tests/oxide-fixture/carrier/src/main.rs @@ -34,7 +34,11 @@ mod kernels { shift: usize, rounds: u32, ) { - let offset = thread::blockIdx_x() as usize * 128 * 4 + shift; + let rank = thread::blockIdx_x() as usize + + thread::gridDim_x() as usize + * (thread::blockIdx_y() as usize + + thread::gridDim_y() as usize * thread::blockIdx_z() as usize); + let offset = rank * 128 * 4 + shift; let mut remaining = rounds; while remaining > 0 { cycle::(input, output, offset, valid); @@ -75,12 +79,27 @@ mod kernels { carrier::store(&mut output, offset, 512, &mut values); // No final barrier: returning the owned buffers waits for this kernel. } + + #[kernel] + #[launch_contract(domain = 1, block = (128, 1, 1), dynamic_shared = inferred)] + pub fn repartitioned(input: &[i32], mut output: DisjointSlice>) { + // ITEMS=1 does not change the capacity-4 ownership grain. Block 1's + // offset is in bounds but belongs to block 0. + let offset = thread::blockIdx_x() as usize * 128; + cycle::<1>(input, &mut output, offset, 128); + } + + #[kernel] + #[launch_contract(domain = 2, block = (128, 1, 1), dynamic_shared = inferred)] + pub fn copy_four_2d(input: &[i32], mut output: DisjointSlice>) { + repeat::<4>(input, &mut output, 495, 17, 2); + } } fn main() -> Result<(), Box> { use cuda_async::simt::device_context::{init_device_contexts, with_cuda_context}; use cuda_async::simt::device_operation::DeviceOperation; - use cuda_core::{DeviceBuffer, LaunchConfig1D}; + use cuda_core::{DeviceBuffer, LaunchConfig1D, LaunchConfig2D}; init_device_contexts(0, 1)?; let context = with_cuda_context(0, |context| context.clone())?; @@ -97,18 +116,27 @@ fn main() -> Result<(), Box> { // canonical group containment forbids the adjacent-block overlap. let source = DeviceBuffer::from_host(&stream, &vec![1_i32; 2048])?; let target = DeviceBuffer::from_host(&stream, &vec![0_i32; 2048])?; - let shift = if mode == "overlap" { 1 } else { 2048 }; - let _ = module - .copy_four_async_owned(&four, source, target, 512, shift, 1) - .sync()?; + if mode == "repartition" { + let repartitioned = module.prepare_repartitioned(config)?; + let _ = module + .repartitioned_async_owned(&repartitioned, source, target) + .sync()?; + } else { + let shift = match mode.as_str() { + "overlap" => 1, + "bounds" => 2048, + _ => return Err(format!("unknown trap mode: {mode}").into()), + }; + let _ = module + .copy_four_async_owned(&four, source, target, 512, shift, 1) + .sync()?; + } panic!("invalid collective range did not trap"); } for items in [1, 4] { for shift in [0, 17] { - for valid in [0, 1, items * 128 - 1, items * 128] { - if valid > 512 - shift { - continue; - } + let limit = (items * 128).min(512 - shift); + for valid in [0, 1, limit - 1, limit] { let source = DeviceBuffer::from_host(&stream, &input)?; let target = DeviceBuffer::from_host(&stream, &vec![-99; input.len()])?; let (_source, target) = if items == 1 { @@ -147,6 +175,26 @@ fn main() -> Result<(), Box> { target.copy_to_host(&stream, &mut result)?; stream.synchronize()?; assert_eq!(result, input); - println!("PASS: carrier bounds, item counts, repeated Store/barrier/Load, owned async"); + + let two_dimensional = + module.prepare_copy_four_2d(LaunchConfig2D::auto_shared((2, 2), (128, 1)))?; + let input_2d: Vec<_> = (0..2048).map(|i| 11 * i - 23).collect(); + let source = DeviceBuffer::from_host(&stream, &input_2d)?; + let target = DeviceBuffer::from_host(&stream, &vec![-99; input_2d.len()])?; + let (_source, target) = module + .copy_four_2d_async_owned(&two_dimensional, source, target) + .sync()?; + let mut result = vec![0; input_2d.len()]; + target.copy_to_host(&stream, &mut result)?; + stream.synchronize()?; + for (index, value) in result.into_iter().enumerate() { + let expected = if index % 512 >= 17 { + input_2d[index] + } else { + -99 + }; + assert_eq!(value, expected, "2D grid index={index}"); + } + println!("PASS: 28,160 carrier values, bounds, item counts, 2D grid, reuse, owned async"); Ok(()) } From 7fa10952b28b159f24be394d3d3291b966977ca4 Mon Sep 17 00:00:00 2001 From: Trent Nelson Date: Tue, 20 Oct 2026 08:29:37 -0700 Subject: [PATCH 16/33] Exercise native storage composition and completion Signed-off-by: Trent Nelson --- .../tests/check_storage_rejections.py | 235 ++++++++++++++++++ .../tests/oxide-fixture/README.md | 14 ++ .../tests/oxide-fixture/api/Cargo.toml | 1 + .../tests/oxide-fixture/api/src/lib.rs | 3 + .../tests/oxide-fixture/api/src/storage.rs | 108 ++++++++ .../tests/oxide-fixture/consumer/Cargo.toml | 2 + .../tests/oxide-fixture/consumer/src/main.rs | 11 + .../oxide-fixture/consumer/src/storage.rs | 161 ++++++++++++ .../tests/oxide-fixture/helper/Cargo.toml | 1 + .../tests/oxide-fixture/helper/src/lib.rs | 3 + .../tests/oxide-fixture/helper/src/storage.rs | 121 +++++++++ .../tests/run_oxide_fixture.py | 12 +- 12 files changed, 670 insertions(+), 2 deletions(-) create mode 100644 rust/cuda-coop-codegen/tests/check_storage_rejections.py create mode 100644 rust/cuda-coop-codegen/tests/oxide-fixture/api/src/storage.rs create mode 100644 rust/cuda-coop-codegen/tests/oxide-fixture/consumer/src/storage.rs create mode 100644 rust/cuda-coop-codegen/tests/oxide-fixture/helper/src/storage.rs diff --git a/rust/cuda-coop-codegen/tests/check_storage_rejections.py b/rust/cuda-coop-codegen/tests/check_storage_rejections.py new file mode 100644 index 000000000000..f67f47730ab5 --- /dev/null +++ b/rust/cuda-coop-codegen/tests/check_storage_rejections.py @@ -0,0 +1,235 @@ +# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. +# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +"""Compile storage/layout/epoch contrasts in an isolated --storage fixture.""" + +import argparse +import json +import os +import subprocess +from pathlib import Path + +parser = argparse.ArgumentParser(description=__doc__) +parser.add_argument("output", type=Path) +parser.add_argument("--case", action="append", default=[]) +args = parser.parse_args() +output = args.output.resolve() +record = json.loads((output / "command.json").read_text()) +source = output / "fixture/consumer/src/storage.rs" +original = source.read_text() +command = record["argv"].copy() +if "run" in command: + command[command.index("run")] = "build" +env = os.environ.copy() +env.update( + CUDA_OXIDE_BACKEND=record["backend"], + CUDA_HOME=record["cuda_root"], + CUDA_PATH=record["cuda_root"], +) +prefix = """// Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception +use coop_marker_fixture_helper::storage::{load,store,reload}; +use cuda_device::{CollectiveTiles,DisjointSlice,library,thread,kernel,launch_contract}; +#[kernel] +#[launch_contract(domain=1,block=(128,1,1),dynamic_shared=inferred)] +pub fn probe(input:&[i32],mut output:DisjointSlice>,valid:usize) { + let offset=thread::blockIdx_x() as usize*512; + let mut values=[0_i32;1]; + unsafe { +""" +suffix = "\n}}\npub fn run()->Result<(),Box> {Ok(())}\n" +new = "let mut scratch=library::storage_new_v1::<0,0,false,false>();\n" +load = "load::<1,true,_>(input,offset,valid,&mut values,&mut scratch);\n" +store = "store::<1,true,_>(&mut output,offset,valid,&mut values,&mut scratch);\n" +reserve = "let view=library::storage_reserve_v1::(&mut scratch);\n" +release = "library::storage_release_v1(view);\n" +sync = "thread::sync_threads();\n" +cases = { + "manual_reuse": (new + load + load, "explicit storage reuse requires"), + "manual_sync": (new + load + sync + load, None), + "automatic": (new.replace("false,false", "true,false") + load + load, None), + "exclusive_sites": (new.replace("false,false", "false,true") + load + load, None), + "exclusive_loop": ( + new.replace("false,false", "false,true") + + "let mut i=0;while i<2 {" + + load + + "i+=1;}", + "explicit storage reuse requires", + ), + "managed_loop": ( + "let mut i=0;while i<2 {load::<1,true,_>(input,offset,valid,&mut values,());i+=1;}", + None, + ), + "constructor_loop": ( + "let mut i=0;while i<2 {" + new + load + "i+=1;}", + "explicit storage reuse requires", + ), + "live_reservation": ( + new + reserve + sync + load + release, + "overlaps a live reservation", + ), + "release_without_sync": ( + new + reserve + release + load, + "explicit storage reuse requires", + ), + "sync_before_release": ( + new + reserve + sync + release + load, + "explicit storage reuse requires", + ), + "release_then_sync": (new + reserve + release + sync + load, None), + "forgotten_view": ( + new + reserve + "core::mem::forget(view);" + sync + load, + "overlaps a live reservation", + ), + "exclusive_views": ( + new.replace("false,false", "false,true") + + reserve + + "let other=library::storage_reserve_v1::(&mut scratch);" + + release + + "library::storage_release_v1(other);", + None, + ), + "reserve_auto": ( + new.replace("false,false", "true,false") + reserve + release, + "reservations cannot use auto_sync", + ), + "zero_extent": ( + new + reserve.replace("128,0", "0,0") + release, + "extent must be positive", + ), + "invalid_alignment": ( + new.replace("0,0,false", "0,3,false") + load, + "alignment must be zero or a power of two", + ), + "unsupported_scalar": ( + new + reserve.replace("::();let selected=if valid==0 {&mut scratch}else{&mut other};load::<1,true,_>(input,offset,valid,&mut values,selected);", + "no unambiguous descriptor origin", + ), + "direct_output_epoch": ( + "store::<1,false,_>(&mut output,offset,valid,&mut values,());reload::<1,false,_>(&mut output,offset,valid,&mut values,());", + "conflicting collective resource access", + ), + "managed_output_epoch": ( + "store::<1,true,_>(&mut output,offset,valid,&mut values,());reload::<1,true,_>(&mut output,offset,valid,&mut values,());", + None, + ), + "independent_output_epoch": ( + new + + "let mut other=library::storage_new_v1::<0,0,false,false>();" + + store + + "reload::<1,true,_>(&mut output,offset,valid,&mut values,&mut other);", + "conflicting collective resource access", + ), + "explicit_output_sync": ( + new + + store + + sync + + "reload::<1,true,_>(&mut output,offset,valid,&mut values,&mut scratch);", + None, + ), + "native_fixed_prefix": ( + "let _prefix=cuda_device::DynamicSharedArray::<[u64;8],32>::get_fixed();" + + new + + load + + sync + + store, + None, + ), + "native_unconstrained": ( + "let _prefix=cuda_device::DynamicSharedArray::::get();" + new + load, + "unconstrained native dynamic shared view", + ), +} +cases.update( + { + "configure_before_use": ( + new + + "let mut scratch: library::StorageDescriptor<0,32,false,false>=library::storage_reconfigure_v1(scratch);" + + load, + None, + ), + "configure_after_use": ( + new + + load + + sync + + "let mut scratch: library::StorageDescriptor<0,32,false,false>=library::storage_reconfigure_v1(scratch);" + + load, + "storage policy cannot change after use", + ), + "configure_chain": ( + new + + "let scratch: library::StorageDescriptor<1024,0,false,false>=library::storage_reconfigure_v1(scratch);let mut scratch: library::StorageDescriptor<1024,32,false,false>=library::storage_reconfigure_v1(scratch);" + + load, + None, + ), + } +) +for capacity in [55, 56]: + cases[f"padded_capacity_{capacity}"] = ( + f"let mut scratch=library::storage_new_v1::<{capacity},0,false,true>();" + "let first=library::storage_reserve_v1::(&mut scratch);" + "let second=library::storage_reserve_v1::(&mut scratch);" + "library::storage_release_v1(first);library::storage_release_v1(second);", + "smaller than complete plan" if capacity == 55 else None, + ) +for dtype in ["i8", "u8", "i16", "u16", "i32", "u32", "i64", "u64", "f32", "f64"]: + cases[f"reserve_{dtype}"] = ( + new + reserve.replace("::", f"::<{dtype},_,3,0>") + release, + None, + ) +unknown = set(args.case) - cases.keys() +if unknown: + parser.error(f"unknown cases: {sorted(unknown)}") +results = {} +try: + for name, (body, diagnostic) in cases.items(): + if args.case and name not in args.case: + continue + source.write_text(prefix + body + suffix) + log_path = output / f"storage-{name}.log" + with log_path.open("w") as log: + result = subprocess.run( + command, + cwd=record["cwd"], + env=env, + stdout=log, + stderr=subprocess.STDOUT, + ) + text = log_path.read_text() + passed = ( + result.returncode == 0 + if diagnostic is None + else result.returncode != 0 and diagnostic in text + ) + results[name] = { + "returncode": result.returncode, + "expected": diagnostic, + "passed": passed, + } + print(f"{name}: {'PASS' if passed else 'FAIL'}", flush=True) + if not passed: + print(text[-5000:], flush=True) + break +finally: + source.write_text(original) + (output / "storage-compile-results.json").write_text( + json.dumps(results, indent=2) + "\n" + ) +raise SystemExit(0 if all(item["passed"] for item in results.values()) else 1) diff --git a/rust/cuda-coop-codegen/tests/oxide-fixture/README.md b/rust/cuda-coop-codegen/tests/oxide-fixture/README.md index cb3e919abe50..caeb786e742c 100644 --- a/rust/cuda-coop-codegen/tests/oxide-fixture/README.md +++ b/rust/cuda-coop-codegen/tests/oxide-fixture/README.md @@ -105,3 +105,17 @@ controls, descriptor/reference mutation, ambiguous origins and mixed markers, with positive cases for independent allocations and barriers in helpers. The native carrier binary also accepts `overlap` or `bounds`; each invocation must fail through a device trap, before an invalid native access. + +The `--storage` runner route exercises the internal full-block storage boundary +with the unchanged numeric provider: shared/manual/automatic and exclusive +policies, distinct helper-created descriptors, typed raw reservations, and +Direct/Transpose operations at one and four items per thread. It checks three +blocks, boundary valid counts and inferred launch bytes through owned async +completion. The wrapper remains private and unsafe for its stated memory and +initialization obligations. + +Add `--compile-only` to prepare an isolated fixture without GPU execution. +`check_storage_rejections.py OUTPUT` uses that copy for paired positive and +negative layout/reuse/output-epoch cases and restores its source afterward. +It checks intended diagnostics, including whole-plan capacity, live views, +release versus completion and scratch-free output conflicts. diff --git a/rust/cuda-coop-codegen/tests/oxide-fixture/api/Cargo.toml b/rust/cuda-coop-codegen/tests/oxide-fixture/api/Cargo.toml index 2dce4d7120fc..a20c32465f28 100644 --- a/rust/cuda-coop-codegen/tests/oxide-fixture/api/Cargo.toml +++ b/rust/cuda-coop-codegen/tests/oxide-fixture/api/Cargo.toml @@ -16,3 +16,4 @@ cuda-coop-codegen = { path = "../../.." } [features] checked = [] carrier = [] +storage = [] 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 index 63ae2fe618d9..4cd1eb1a8fd5 100644 --- a/rust/cuda-coop-codegen/tests/oxide-fixture/api/src/lib.rs +++ b/rust/cuda-coop-codegen/tests/oxide-fixture/api/src/lib.rs @@ -184,3 +184,6 @@ mod checked { } #[cfg(feature = "checked")] pub use checked::{checked_load, checked_store}; + +#[cfg(feature = "storage")] +pub mod storage; diff --git a/rust/cuda-coop-codegen/tests/oxide-fixture/api/src/storage.rs b/rust/cuda-coop-codegen/tests/oxide-fixture/api/src/storage.rs new file mode 100644 index 000000000000..6b5478c358d5 --- /dev/null +++ b/rust/cuda-coop-codegen/tests/oxide-fixture/api/src/storage.rs @@ -0,0 +1,108 @@ +// Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +//! Private storage integration using the unchanged R01 numeric provider. +use cuda_device::{CollectiveTiles, DisjointSlice, library}; + +/// # Safety +/// All lanes agree on source, offset, valid count and storage origin. The source +/// contains the full accessed block range; the filled operation initializes all +/// lane payload slots. The root is exactly one 128-thread block per owned tile. +#[inline(never)] +pub unsafe fn load( + source: &[i32], + offset: usize, + valid: usize, + values: &mut [i32; ITEMS], + storage: Storage, +) { + assert!(valid <= 128 * ITEMS && offset <= source.len() && valid <= source.len() - offset); + unsafe { + if TRANSPOSE { + library::checked_block_storage_call_v1::( + super::transpose_load_filled, + source, + storage, + (source.as_ptr().add(offset), valid as i32, -7_i32), + (values.as_mut_ptr(),), + ); + } else { + library::checked_block_storage_call_v1::( + super::direct_load_filled, + source, + storage, + (source.as_ptr().add(offset), valid as i32, -7_i32), + (values.as_mut_ptr(),), + ); + } + } +} + +/// # Safety +/// All lanes agree on target, offset, valid count and storage origin. Offset is +/// this block's canonical 512-element tile start; payload values are initialized. +#[inline(never)] +pub unsafe fn store( + target: &mut DisjointSlice>, + offset: usize, + valid: usize, + values: &mut [i32; ITEMS], + storage: Storage, +) { + assert!(valid <= 128 * ITEMS && offset <= target.len() && valid <= target.len() - offset); + assert!(offset == cuda_device::thread::blockIdx_x() as usize * 512); + let pointer = target.as_mut_ptr(); + unsafe { + if TRANSPOSE { + library::checked_block_storage_call_v1::( + super::transpose_store_partial, + target, + storage, + (pointer.add(offset), valid as i32), + (values.as_mut_ptr(),), + ); + } else { + library::checked_block_storage_call_v1::( + super::direct_store_partial, + target, + storage, + (pointer.add(offset), valid as i32), + (values.as_mut_ptr(),), + ); + } + } +} + +/// # Safety +/// The carrier and range obey `store`'s ownership contract. A prior write must +/// be completed by an explicit or genuinely required planned block barrier. +#[inline(never)] +pub unsafe fn reload( + source: &mut DisjointSlice>, + offset: usize, + valid: usize, + values: &mut [i32; ITEMS], + storage: Storage, +) { + assert!(valid <= 128 * ITEMS && offset <= source.len() && valid <= source.len() - offset); + let pointer = source.as_mut_ptr(); + unsafe { + if TRANSPOSE { + library::checked_block_storage_call_v1::( + super::transpose_load_filled, + &*source, + storage, + (pointer.add(offset).cast_const(), valid as i32, -7_i32), + (values.as_mut_ptr(),), + ); + } else { + library::checked_block_storage_call_v1::( + super::direct_load_filled, + &*source, + storage, + (pointer.add(offset).cast_const(), valid as i32, -7_i32), + (values.as_mut_ptr(),), + ); + } + } +} diff --git a/rust/cuda-coop-codegen/tests/oxide-fixture/consumer/Cargo.toml b/rust/cuda-coop-codegen/tests/oxide-fixture/consumer/Cargo.toml index ff096fc7ba97..1d496140d001 100644 --- a/rust/cuda-coop-codegen/tests/oxide-fixture/consumer/Cargo.toml +++ b/rust/cuda-coop-codegen/tests/oxide-fixture/consumer/Cargo.toml @@ -11,9 +11,11 @@ publish = false cuda-device.workspace = true cuda-host.workspace = true cuda-core.workspace = true +cuda-async.workspace = true oxide-artifacts.workspace = true coop-marker-fixture-helper = { path = "../helper" } [features] checked = ["coop-marker-fixture-helper/checked"] root-context = ["checked"] +storage = ["coop-marker-fixture-helper/storage", "cuda-host/async"] 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 index fb2f5eff8a13..3f2b53b5f283 100644 --- a/rust/cuda-coop-codegen/tests/oxide-fixture/consumer/src/main.rs +++ b/rust/cuda-coop-codegen/tests/oxide-fixture/consumer/src/main.rs @@ -1,9 +1,12 @@ // Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception +#[cfg(not(feature = "storage"))] use cuda_device::{kernel, launch_bounds, launch_contract}; +#[cfg(not(feature = "storage"))] use cuda_host::cuda_module; +#[cfg(not(feature = "storage"))] #[cuda_module] mod kernels { use super::*; @@ -109,6 +112,7 @@ mod kernels { } } +#[cfg(not(feature = "storage"))] fn main() -> Result<(), Box> { use cuda_core::{CudaContext, DeviceBuffer}; use cuda_core::{LaunchConfig3D, LaunchContractError}; @@ -340,3 +344,10 @@ fn main() -> Result<(), Box> { ); Ok(()) } + +#[cfg(feature = "storage")] +mod storage; +#[cfg(feature = "storage")] +fn main() -> Result<(), Box> { + storage::run() +} diff --git a/rust/cuda-coop-codegen/tests/oxide-fixture/consumer/src/storage.rs b/rust/cuda-coop-codegen/tests/oxide-fixture/consumer/src/storage.rs new file mode 100644 index 000000000000..9b7feab4bb01 --- /dev/null +++ b/rust/cuda-coop-codegen/tests/oxide-fixture/consumer/src/storage.rs @@ -0,0 +1,161 @@ +// Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +use cuda_device::{CollectiveTiles, DisjointSlice, kernel, launch_contract}; +use cuda_host::cuda_module; + +#[cuda_module] +mod kernels { + use super::*; + #[kernel] + #[launch_contract(domain = 1, block = (128, 1, 1), dynamic_shared = inferred)] + pub fn copy( + input: &[i32], + mut output: DisjointSlice>, + valid: usize, + ) { + // SAFETY: host-owned buffers cover every canonical tile; root geometry + // and uniform controls satisfy the private wrapper's stated contract. + unsafe { + coop_marker_fixture_helper::storage::roundtrip::( + input, + &mut output, + valid, + ); + } + } +} + +pub fn run() -> Result<(), Box> { + use cuda_async::simt::device_context::{init_device_contexts, with_cuda_context}; + use cuda_async::simt::device_operation::DeviceOperation; + use cuda_core::{DeviceBuffer, LaunchConfig1D}; + let image = std::fs::read(std::env::current_exe()?)?; + let bundles = oxide_artifacts::read_artifact_bundles_from_object_bytes(&image)?; + let bundle = bundles + .iter() + .find(|bundle| { + bundle + .payload(oxide_artifacts::ArtifactPayloadKind::Cubin) + .is_some() + }) + .ok_or("missing native storage artifact")?; + std::fs::write( + "storage.cubin", + bundle + .payload(oxide_artifacts::ArtifactPayloadKind::Cubin) + .unwrap(), + )?; + for entry in &bundle.entries { + if let Some(facts) = entry.launch_facts { + println!("ENTRY {} {:?}", entry.symbol, facts); + } + } + if std::env::var_os("OXIDE_STORAGE_INSPECT_ONLY").is_some() { + return Ok(()); + } + init_device_contexts(0, 1)?; + let context = with_cuda_context(0, |c| c.clone())?; + let stream = context.default_stream(); + let module = unsafe { kernels::load_async(0)? }; + let input: Vec<_> = (0..1536).map(|i| i * 3 + 17).collect(); + let mut cases = 0; + macro_rules! check { + ($i:literal,$transpose:literal,$mode:literal) => {{ + let plan = module + .prepare_copy::<$i, $transpose, $mode>(LaunchConfig1D::auto_shared(3, 128))?; + println!( + "storage I={} transpose={} mode={} bytes={}", + $i, + $transpose, + $mode, + plan.__raw_config().shared_mem_bytes + ); + assert!( + module + .prepare_copy::<$i, $transpose, $mode>(LaunchConfig1D::auto_shared(3, 64)) + .is_err() + ); + let single = if $i == 1 { 528 } else { 2064 }; + let expected = if !$transpose { + 0 + } else if $mode == 3 { + single + 2 * 512 * $i + } else if $mode == 4 { + single + 512 * $i + } else if $mode == 6 { + 1536 + single + 512 * $i + } else if $mode == 7 { + ((64 + single + 63) & !63) + 512 * $i + } else { + single + }; + assert_eq!(plan.__raw_config().shared_mem_bytes, expected); + if expected > 0 { + assert!( + module + .prepare_copy::<$i, $transpose, $mode>(LaunchConfig1D::new( + 3, + 128, + expected - 1 + )) + .is_err() + ); + } + for valid in [0, 1, 128 * $i - 1, 128 * $i] { + let source = DeviceBuffer::from_host(&stream, &input)?; + let target = DeviceBuffer::from_host(&stream, &vec![-99; input.len()])?; + let (_source, target) = module + .copy_async_owned::<$i, $transpose, $mode, _, _>(&plan, source, target, valid) + .sync()?; + let mut result = vec![0; input.len()]; + target.copy_to_host(&stream, &mut result)?; + stream.synchronize()?; + for (index, value) in result.into_iter().enumerate() { + assert_eq!( + value, + if index % 512 < valid { + input[index] + + if $mode == 5 { + ((index % 512 / $i + 1) % 128) as i32 + } else if $mode == 6 { + (((index % 512 / $i + 1) % 128) * 3 + 1) as i32 + } else if $mode == 7 { + (index % 512 / $i % 8 + 10) as i32 + } else { + 0 + } + } else { + -99 + }, + "I={} transpose={} mode={} valid={valid} index={index}", + $i, + $transpose, + $mode + ); + } + cases += 1; + } + }}; + } + check!(1, false, 0); + check!(4, false, 0); + check!(1, true, 0); + check!(4, true, 0); + check!(1, true, 1); + check!(4, true, 1); + check!(1, true, 2); + check!(4, true, 2); + check!(1, true, 3); + check!(4, true, 3); + check!(1, true, 4); + check!(4, true, 4); + check!(1, true, 5); + check!(4, true, 5); + check!(1, true, 6); + check!(4, true, 6); + check!(1, true, 7); + check!(4, true, 7); + println!("PASS: {cases} storage cases with three blocks and boundary counts"); + 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 index abb1aba763ac..bce66cd3101c 100644 --- a/rust/cuda-coop-codegen/tests/oxide-fixture/helper/Cargo.toml +++ b/rust/cuda-coop-codegen/tests/oxide-fixture/helper/Cargo.toml @@ -13,3 +13,4 @@ coop-marker-fixture-api = { path = "../api" } [features] checked = ["coop-marker-fixture-api/checked"] +storage = ["coop-marker-fixture-api/storage"] 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 index 8f0145b72ba3..77f9296ed30d 100644 --- a/rust/cuda-coop-codegen/tests/oxide-fixture/helper/src/lib.rs +++ b/rust/cuda-coop-codegen/tests/oxide-fixture/helper/src/lib.rs @@ -84,3 +84,6 @@ pub unsafe fn roundtrip( mod checked; #[cfg(feature = "checked")] pub use checked::roundtrip; + +#[cfg(feature = "storage")] +pub mod storage; diff --git a/rust/cuda-coop-codegen/tests/oxide-fixture/helper/src/storage.rs b/rust/cuda-coop-codegen/tests/oxide-fixture/helper/src/storage.rs new file mode 100644 index 000000000000..4c55b0de622e --- /dev/null +++ b/rust/cuda-coop-codegen/tests/oxide-fixture/helper/src/storage.rs @@ -0,0 +1,121 @@ +// Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +pub use coop_marker_fixture_api::storage::{load, reload, store}; +use cuda_device::{CollectiveTiles, DisjointSlice, library, thread}; + +/// # Safety +/// Called uniformly by a block. Each rooted constructor call is independent. +#[inline(never)] +pub unsafe fn descriptor() -> library::StorageDescriptor { + unsafe { library::storage_new_v1::<0, 0, false, false>() } +} + +/// # Safety +/// Exact 128-thread blocks own disjoint 512-element target tiles. Input covers +/// all source ranges; every control and descriptor agrees across block lanes. +#[inline(never)] +pub unsafe fn roundtrip( + input: &[i32], + output: &mut DisjointSlice>, + valid: usize, +) { + let offset = thread::blockIdx_x() as usize * 512; + let mut values = [0; ITEMS]; + unsafe { + match MODE { + 0 => { + load::(input, offset, valid, &mut values, ()); + store::(output, offset, valid, &mut values, ()); + } + 1 => { + let mut scratch = descriptor(); + load::(input, offset, valid, &mut values, &mut scratch); + thread::sync_threads(); + store::(output, offset, valid, &mut values, &mut scratch); + } + 2 => { + let mut scratch = library::storage_new_v1::<0, 0, true, false>(); + load::(input, offset, valid, &mut values, &mut scratch); + store::(output, offset, valid, &mut values, &mut scratch); + } + 3 => { + let mut scratch = library::storage_new_v1::<0, 64, false, true>(); + // Same generic helper, different static uses and assigned offsets. + load::(input, offset, valid, &mut values, &mut scratch); + load::(input, offset, valid, &mut values, &mut scratch); + store::(output, offset, valid, &mut values, &mut scratch); + } + 4 => { + let mut first = descriptor(); + let mut second = descriptor(); + load::(input, offset, valid, &mut values, &mut first); + load::(input, offset, valid, &mut values, &mut second); + thread::sync_threads(); + store::(output, offset, valid, &mut values, &mut first); + } + 5 => { + let mut scratch = descriptor(); + let view = library::storage_reserve_v1::(&mut scratch); + let lane = thread::threadIdx_x() as usize; + view.as_mut_ptr().add(lane).write(lane as i32); + thread::sync_threads(); + let saved = view.as_mut_ptr().add((lane + 1) % 128).read(); + library::storage_release_v1(view); + thread::sync_threads(); + load::(input, offset, valid, &mut values, &mut scratch); + // A cross-lane read remains observable after this region is reused. + let mut item = 0; + while item < ITEMS { + values[item] += saved; + item += 1; + } + thread::sync_threads(); + store::(output, offset, valid, &mut values, &mut scratch); + } + 6 => { + let mut scratch = library::storage_new_v1::<0, 32, false, true>(); + let first = library::storage_reserve_v1::(&mut scratch); + let second = library::storage_reserve_v1::(&mut scratch); + let lane = thread::threadIdx_x() as usize; + first.as_mut_ptr().add(lane).write(lane as i32); + second.as_mut_ptr().add(lane).write((lane * 2 + 1) as u64); + thread::sync_threads(); + let neighbor = (lane + 1) % 128; + let saved = first.as_mut_ptr().add(neighbor).read() + + second.as_mut_ptr().add(neighbor).read() as i32; + library::storage_release_v1(first); + library::storage_release_v1(second); + thread::sync_threads(); + load::(input, offset, valid, &mut values, &mut scratch); + let mut item = 0; + while item < ITEMS { + values[item] += saved; + item += 1; + } + store::(output, offset, valid, &mut values, &mut scratch); + } + _ => { + let prefix = + cuda_device::DynamicSharedArray::<[u64; 8], 32>::get_fixed().cast::(); + let lane = thread::threadIdx_x() as usize; + if lane < 8 { + prefix.add(lane).write(lane as u64 + 10); + } + thread::sync_threads(); + let mut first = descriptor(); + let mut second = library::storage_new_v1::<0, 64, false, false>(); + load::(input, offset, valid, &mut values, &mut first); + load::(input, offset, valid, &mut values, &mut second); + let saved = prefix.add(lane % 8).read() as i32; + let mut item = 0; + while item < ITEMS { + values[item] += saved; + item += 1; + } + thread::sync_threads(); + store::(output, offset, valid, &mut values, &mut first); + } + } + } +} diff --git a/rust/cuda-coop-codegen/tests/run_oxide_fixture.py b/rust/cuda-coop-codegen/tests/run_oxide_fixture.py index c2899825874b..2e4ede054860 100644 --- a/rust/cuda-coop-codegen/tests/run_oxide_fixture.py +++ b/rust/cuda-coop-codegen/tests/run_oxide_fixture.py @@ -28,6 +28,12 @@ parser.add_argument( "--carrier", action="store_true", help="exercise private collective-only carriers" ) +parser.add_argument( + "--storage", action="store_true", help="exercise planned full-block storage" +) +parser.add_argument( + "--compile-only", action="store_true", help="compile without executing the fixture" +) args = parser.parse_args() source = Path(__file__).resolve().parent checkout = source.parents[2] @@ -69,12 +75,14 @@ command = [ str(args.cargo_oxide.resolve()), "oxide", - "run", + "build" if args.compile_only else "run", "--arch", args.arch, "--materialize-cubin", ] -if args.root_context: +if args.storage: + command.extend(["--features", "storage"]) +elif args.root_context: command.extend(["--features", "root-context"]) elif args.checked: command.extend(["--features", "checked"]) From a9b49d60cd1a5f70a8a2758b7c779a34d23ca8f1 Mon Sep 17 00:00:00 2001 From: Trent Nelson Date: Tue, 20 Oct 2026 09:10:51 -0700 Subject: [PATCH 17/33] test(rust): check MaybeUninit collective compatibility Signed-off-by: Trent Nelson --- .../tests/check_collective_rejections.py | 56 +++++++++++++++++++ 1 file changed, 56 insertions(+) diff --git a/rust/cuda-coop-codegen/tests/check_collective_rejections.py b/rust/cuda-coop-codegen/tests/check_collective_rejections.py index d67393be7fe2..7689ec4308eb 100644 --- a/rust/cuda-coop-codegen/tests/check_collective_rejections.py +++ b/rust/cuda-coop-codegen/tests/check_collective_rejections.py @@ -100,6 +100,48 @@ def in_helper(code): in_helper(" let input = unsafe { input.add(lane) };"), "not proved block-uniform", ), + "maybe_uninit_value": ( + { + "helper": in_helper( + " let initialized = initialized_value((valid, 0_u8));" + )["helper"].replace( + "payload[i] += CONTEXT_INCREMENT;", + "payload[i] += CONTEXT_INCREMENT + initialized.0;", + ) + + """ +#[inline(never)] +fn initialized_value(value: T) -> T { + unsafe { core::mem::MaybeUninit::new(value).assume_init() } +} +""" + }, + None, + ), + "maybe_uninit_uninhabited": ( + in_helper( + " if mode & 8 != 0 {\n" + " #[allow(invalid_value)]\n" + " let _ = unsafe {\n" + " core::mem::MaybeUninit::::uninit().assume_init()\n" + " };\n" + " }" + ), + "assert_inhabited requires a proved inhabited concrete type", + ), + "maybe_uninit_count": ( + in_helper( + " let initialized = core::mem::MaybeUninit::new(lane as i32);\n" + " let valid = unsafe { initialized.assume_init() };" + ), + "not proved block-uniform", + ), + "maybe_uninit_pointer": ( + in_helper( + " let initialized = unsafe { core::mem::MaybeUninit::new(input.add(lane)) };\n" + " let input = unsafe { initialized.assume_init() };" + ), + "not proved block-uniform", + ), "early_return": (in_helper(" if lane & 1 == 0 { return; }"), "lane-varying"), "cycle": ( in_helper( @@ -276,6 +318,20 @@ def in_helper(code): }, None, ) +cases["maybe_uninit_dead"] = ( + in_helper( + " match ITEMS_PER_THREAD {\n" + " 0 => {\n" + " #[allow(invalid_value)]\n" + " let _ = unsafe {\n" + " core::mem::MaybeUninit::::uninit().assume_init()\n" + " };\n" + " }\n" + " _ => {}\n" + " }" + ), + None, +) unknown = set(args.case) - cases.keys() if unknown: parser.error(f"unknown cases: {', '.join(sorted(unknown))}") From 2119cae9b6d53fa31d2d5786421874d588765f6d Mon Sep 17 00:00:00 2001 From: Trent Nelson Date: Tue, 20 Oct 2026 09:25:24 -0700 Subject: [PATCH 18/33] Add validity-aware Direct block Load and Store Signed-off-by: Trent Nelson --- docs/rust/index.rst | 24 +- rust/cuda-coop-codegen/src/block_io.cuh | 62 ++ rust/cuda-coop-codegen/src/lib.rs | 286 ++++++++-- rust/cuda-coop-codegen/src/oxide.rs | 121 +++- .../tests/check_validity_rejections.py | 120 ++++ rust/cuda-coop-codegen/tests/native.rs | 72 ++- .../tests/oxide-fixture/README.md | 15 + .../tests/oxide-fixture/api/Cargo.toml | 1 + .../tests/oxide-fixture/api/build.rs | 10 + .../tests/oxide-fixture/api/src/carrier.rs | 8 +- .../tests/oxide-fixture/api/src/lib.rs | 57 ++ .../tests/oxide-fixture/api/src/validity.rs | 174 ++++++ .../tests/oxide-fixture/carrier/Cargo.toml | 9 + .../oxide-fixture/carrier/src/validity.rs | 528 ++++++++++++++++++ .../tests/run_oxide_fixture.py | 12 +- rust/cuda-coop/src/thread_data.rs | 3 +- rust/cuda-coop/src/thread_data/block_io.rs | 21 + 17 files changed, 1449 insertions(+), 74 deletions(-) create mode 100644 rust/cuda-coop-codegen/src/block_io.cuh create mode 100644 rust/cuda-coop-codegen/tests/check_validity_rejections.py create mode 100644 rust/cuda-coop-codegen/tests/oxide-fixture/api/src/validity.rs create mode 100644 rust/cuda-coop-codegen/tests/oxide-fixture/carrier/src/validity.rs create mode 100644 rust/cuda-coop/src/thread_data/block_io.rs diff --git a/docs/rust/index.rst b/docs/rust/index.rst index 25f80abf7c14..e2950e6e8e8f 100644 --- a/docs/rust/index.rst +++ b/docs/rust/index.rst @@ -23,8 +23,9 @@ 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: +Each request describes a complete, one-dimensional block and a positive +compile-time item count per thread. The original raw ``i32`` recipes support +these forms: .. list-table:: :header-rows: 1 @@ -50,9 +51,22 @@ 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. +ownership, valid bounds, and synchronization before scratch reuse. + +The separate validity-aware Direct recipes support all ten numeric payload +types: ``i8``, ``u8``, ``i16``, ``u16``, ``i32``, ``u32``, ``i64``, ``u64``, +``f32`` and ``f64``. ``Request::new_typed`` selects the exact type. These recipes +take opaque value storage and separate per-slot ``u8`` flags. Full and +default-filled Load initialize every slot. Masked Load without a default copies +only the loaded prefix and resets all other flags to zero. Full and partial +Store require flags equal to exactly one in their write domain and preserve +all payload bytes and flags. A failing Store traps; it is not transactional +across lanes. Partial Store forms destination pointers only within its valid +prefix, including when the allocation is shorter than a complete tile. + +Validity-aware Transpose and the public collective API remain under development. +The private integration fixture uses the actual ``ThreadData`` implementation +without exposing an initialized slice or a public native layout. Compile a provider ------------------ diff --git a/rust/cuda-coop-codegen/src/block_io.cuh b/rust/cuda-coop-codegen/src/block_io.cuh new file mode 100644 index 000000000000..d404f17e0396 --- /dev/null +++ b/rust/cuda-coop-codegen/src/block_io.cuh @@ -0,0 +1,62 @@ +// Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +#include +#include + +namespace cuda_coop_detail +{ +// Copying this cell never evaluates a possibly uninitialized numeric value. +// Conversion to T is used only after the caller's complete Store preflight. +template +struct alignas(T) opaque_cell +{ + unsigned char bytes[sizeof(T)]; + + __device__ opaque_cell() = default; + __device__ opaque_cell(const opaque_cell& other) + { + cuda::std::memcpy(bytes, other.bytes, sizeof(T)); + } + __device__ opaque_cell& operator=(const opaque_cell& other) + { + cuda::std::memcpy(bytes, other.bytes, sizeof(T)); + return *this; + } + __device__ opaque_cell& operator=(T value) + { + cuda::std::memcpy(bytes, &value, sizeof(T)); + return *this; + } + __device__ operator T() const + { + T value; + cuda::std::memcpy(&value, bytes, sizeof(T)); + return value; + } +}; + +// CUB forms a per-lane iterator before testing its valid count. Keep that +// displacement as an integer; form a pointer only for an actually stored slot. +template +struct guarded_output +{ + using value_type = T; + using difference_type = int; + using pointer = T*; + using reference = T&; + using iterator_category = cuda::std::random_access_iterator_tag; + + T* base; + int offset; + + __device__ guarded_output operator+(int delta) const + { + return {base, offset + delta}; + } + __device__ T& operator[](int index) const + { + return base[offset + index]; + } +}; +} // namespace cuda_coop_detail diff --git a/rust/cuda-coop-codegen/src/lib.rs b/rust/cuda-coop-codegen/src/lib.rs index 872adf16c76a..a16180ba5362 100644 --- a/rust/cuda-coop-codegen/src/lib.rs +++ b/rust/cuda-coop-codegen/src/lib.rs @@ -3,7 +3,7 @@ //! Native, host-only CUB specialization for Rust cooperative algorithms. //! -//! This private development crate compiles i32 block Load/Store providers with +//! This private development crate compiles 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. @@ -12,11 +12,13 @@ //! //! [`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. +//! The raw i32 recipes support Direct and Transpose. Their Store payloads must +//! be fully initialized. The validity-aware Direct recipes support all ten +//! [`Dtype`] variants with separate value storage and per-slot validity bytes. +//! Masked Load resets validity outside its prefix; Direct Store preserves both +//! arrays and checks initialization only in its write domain. Transpose Store +//! replaces each lane's blocked payload with striped items across the block; +//! validity-aware Transpose recipes are not yet supported. //! //! [`Provider::from_checkout`] snapshots the selected CCCL and CUDA headers; //! [`Provider::compile`] uses those exact bytes and loaded NVRTC to emit an @@ -101,7 +103,48 @@ impl Algorithm { } } -/// Device operation and its validity controls, all with i32 elements. +/// Exact scalar type used in a native specialization. +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub enum Dtype { + /// Rust `i8`, C++ `signed char`. + I8, + /// Rust `u8`, C++ `unsigned char`. + U8, + /// Rust `i16`, C++ `short`. + I16, + /// Rust `u16`, C++ `unsigned short`. + U16, + /// Rust `i32`, C++ `int`. + I32, + /// Rust `u32`, C++ `unsigned int`. + U32, + /// Rust `i64`, C++ `long long`. + I64, + /// Rust `u64`, C++ `unsigned long long`. + U64, + /// Rust `f32`, C++ `float`. + F32, + /// Rust `f64`, C++ `double`. + F64, +} +impl Dtype { + fn native(self) -> (&'static str, usize) { + match self { + Self::I8 => ("signed char", 1), + Self::U8 => ("unsigned char", 1), + Self::I16 => ("short", 2), + Self::U16 => ("unsigned short", 2), + Self::I32 => ("int", 4), + Self::U32 => ("unsigned int", 4), + Self::I64 => ("long long", 8), + Self::U64 => ("unsigned long long", 8), + Self::F32 => ("float", 4), + Self::F64 => ("double", 8), + } + } +} + +/// Device operation and its validity controls. #[derive(Clone, Copy, Debug, Eq, PartialEq)] pub enum Operation { /// Load every slot. ABI: `(const int* memory, int* payload, unsigned char* scratch)`. @@ -112,10 +155,50 @@ pub enum Operation { StoreFull, /// Store a prefix only. ABI adds `(int valid)`; other destination elements remain unchanged. StorePartial, + /// Full Load with separate value storage and u8 flags; marks every slot initialized. + LoadFullValidity, + /// Counted Load without a default; resets flags outside the loaded prefix to zero. + LoadMaskedValidity, + /// Counted Load with an exact typed default; marks every slot initialized. + LoadFilledValidity, + /// Full Store after checking all flags equal one; preserves values and flags. + StoreFullValidity, + /// Counted Store after checking its write domain; preserves all values and flags. + StorePartialValidity, } impl Operation { fn is_load(self) -> bool { - matches!(self, Self::LoadFull | Self::LoadFilled) + matches!( + self, + Self::LoadFull + | Self::LoadFilled + | Self::LoadFullValidity + | Self::LoadMaskedValidity + | Self::LoadFilledValidity + ) + } + fn validity(self) -> bool { + matches!( + self, + Self::LoadFullValidity + | Self::LoadMaskedValidity + | Self::LoadFilledValidity + | Self::StoreFullValidity + | Self::StorePartialValidity + ) + } + fn counted(self) -> bool { + matches!( + self, + Self::LoadFilled + | Self::StorePartial + | Self::LoadMaskedValidity + | Self::LoadFilledValidity + | Self::StorePartialValidity + ) + } + fn filled(self) -> bool { + matches!(self, Self::LoadFilled | Self::LoadFilledValidity) } fn token(self) -> &'static str { match self { @@ -123,11 +206,16 @@ impl Operation { Self::LoadFilled => "load_filled", Self::StoreFull => "store_full", Self::StorePartial => "store_partial", + Self::LoadFullValidity => "load_full_validity_v1", + Self::LoadMaskedValidity => "load_masked_validity_v1", + Self::LoadFilledValidity => "load_filled_validity_v1", + Self::StoreFullValidity => "store_full_validity_v1", + Self::StorePartialValidity => "store_partial_validity_v1", } } } -/// A validated i32, one-dimensional block specialization. +/// A validated, 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 @@ -136,6 +224,7 @@ impl Operation { pub struct Request { operation: Operation, algorithm: Algorithm, + dtype: Dtype, block_threads: u32, items_per_thread: u32, } @@ -156,6 +245,41 @@ impl Request { block_threads: u32, items_per_thread: u32, ) -> Result { + Self::new_typed( + Dtype::I32, + operation, + algorithm, + block_threads, + items_per_thread, + ) + } + + /// Validate an exact typed block specialization. + /// + /// The five validity-aware recipes support every [`Dtype`] with Direct. + /// Legacy raw recipes retain their i32 ABI and algorithm support. + /// + /// # Errors + /// + /// Rejects the invalid shapes documented by [`Self::new`], a non-i32 raw + /// recipe, or a validity-aware Transpose recipe. + pub fn new_typed( + dtype: Dtype, + operation: Operation, + algorithm: Algorithm, + block_threads: u32, + items_per_thread: u32, + ) -> Result { + if !operation.validity() && dtype != Dtype::I32 { + return Err(Error::message( + "legacy raw recipes require signed i32 elements", + )); + } + if operation.validity() && algorithm != Algorithm::Direct { + return Err(Error::message( + "validity-aware recipes currently require Direct", + )); + } if !(1..=1024).contains(&block_threads) || items_per_thread == 0 || u64::from(block_threads) * u64::from(items_per_thread) > i32::MAX as u64 @@ -167,6 +291,7 @@ impl Request { Ok(Self { operation, algorithm, + dtype, block_threads, items_per_thread, }) @@ -179,11 +304,15 @@ impl Request { pub fn algorithm(self) -> Algorithm { self.algorithm } + /// Exact numeric element type. + pub fn dtype(self) -> Dtype { + self.dtype + } /// 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. + /// Number of payload elements per thread. pub fn items_per_thread(self) -> u32 { self.items_per_thread } @@ -205,7 +334,8 @@ pub struct Layout { /// 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. +/// number of slots of the requested type. Raw 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. @@ -218,6 +348,18 @@ pub struct Layout { /// 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. +/// +/// Validity-aware Direct ABIs add `unsigned char* flags` immediately after +/// scratch. `payload` addresses aligned storage for the requested number of +/// elements; it need not contain initialized values. Flags occupy a separate, +/// initialized byte array of the same slot count. Only exactly one denotes an +/// initialized value. Both arrays are private to the lane, disjoint from +/// memory/scratch and one another, and remain exclusively borrowed for the +/// call. Native code accesses payload bytes without evaluating invalid values. +/// Masked Load writes zero flags outside its prefix; filled/full Load write +/// every value before setting its flag. Store checks its entire lane-local +/// write domain before numeric conversion and preserves both arrays. Failure +/// traps; it does not promise rollback of stores by other lanes. #[derive(Debug)] pub struct Specialization { /// SHA-256 of request/source, target/options, captured headers and native toolchain inputs. @@ -232,7 +374,8 @@ pub struct Specialization { /// 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. + /// True only for raw Transpose stores; every slot remains initialized. + /// This describes Store rearrangement, not whether Load writes its payload. pub mutates_payload: bool, } @@ -347,6 +490,7 @@ impl Provider { include_bytes!("lib.rs"), include_bytes!("headers.rs"), include_bytes!("nvrtc.rs"), + include_bytes!("block_io.cuh"), ] { hash_field(&mut inputs, bytes); } @@ -354,9 +498,8 @@ impl Provider { 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 source = String::from(include_str!("block_io.cuh")); + source.push_str("\ntemplate \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 { @@ -408,56 +551,111 @@ impl Provider { } fn source(request: Request, symbol: &str) -> String { - let (header, primitive, algorithm, pointer) = if request.operation.is_load() { - ("load", "BlockLoad", "BLOCK_LOAD", "const int*") + let op = request.operation; + let (header, primitive, algorithm, qualifier, method) = if op.is_load() { + ("load", "BlockLoad", "BLOCK_LOAD", "const ", "Load") } else { - ("store", "BlockStore", "BLOCK_STORE", "int*") + ("store", "BlockStore", "BLOCK_STORE", "", "Store") }; - 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 (dtype, size) = request.dtype.native(); + let items = request.items_per_thread; + let tile = request.block_threads * items; + let mut controls = String::new(); + let mut operands = String::new(); + let check = if op.counted() { + controls.push_str(", int valid"); + operands.push_str(", valid"); + format!("if (valid < 0 || valid > {tile}) __trap();") + } else { + String::new() }; + if op.filled() { + controls.push_str(&format!(", {dtype} default_value")); + operands.push_str(", default_value"); + } let constructor = if request.algorithm == Algorithm::Direct { "primitive operation;" } else { "primitive operation{*reinterpret_cast(scratch)};" }; - let method = if request.operation.is_load() { - "Load" + let memory = if matches!( + op, + Operation::StorePartial | Operation::StorePartialValidity + ) { + "cuda_coop_detail::guarded_output{memory, 0}" } else { - "Store" + "memory" + }; + let (cell, flags, body) = if op.validity() { + let domain = if matches!( + op, + Operation::LoadMaskedValidity | Operation::StorePartialValidity + ) { + format!("int(threadIdx.x) * {items} + i < valid") + } else { + "true".to_owned() + }; + let body = if op.is_load() { + format!( + r#"cell values[{items}]; + operation.Load(memory, values{operands}); +#pragma unroll + for (int i = 0; i < {items}; ++i) {{ + if ({domain}) {{ + cuda::std::memcpy(bytes + sizeof(value_type) * i, values[i].bytes, sizeof(value_type)); + flags[i] = 1; + }} else {{ + flags[i] = 0; + }} + }}"# + ) + } else { + format!( + r#"// Complete this lane's preflight before any numeric conversion or store. +#pragma unroll + for (int i = 0; i < {items}; ++i) {{ + if (({domain}) && flags[i] != 1) __trap(); + }} + cell values[{items}]; +#pragma unroll + for (int i = 0; i < {items}; ++i) {{ + if ({domain}) cuda::std::memcpy(values[i].bytes, bytes + sizeof(value_type) * i, sizeof(value_type)); + }} + operation.Store({memory}, values{operands});"# + ) + }; + ( + "cuda_coop_detail::opaque_cell", + ", unsigned char* flags", + format!("auto* const bytes = reinterpret_cast(payload);\n {body}"), + ) + } else { + ( + "value_type", + "", + format!( + "operation.{method}({memory}, *reinterpret_cast(payload){operands});" + ), + ) }; 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}) {{ +using value_type = {dtype}; +using cell = {cell}; +using primitive = cub::{primitive}; +static_assert(sizeof(value_type) == {size} && alignof(value_type) == {size}); +static_assert(sizeof(cell) == sizeof(value_type) && alignof(cell) == alignof(value_type)); +extern "C" __device__ void {symbol}({qualifier}value_type* memory, value_type* payload, unsigned char* scratch{flags}{controls}) {{ {check} {constructor} - operation.{method}(memory, *reinterpret_cast(payload){operands}); + {body} }} }} "#, threads = request.block_threads, - items = request.items_per_thread, choice = request.algorithm.token() ) } diff --git a/rust/cuda-coop-codegen/src/oxide.rs b/rust/cuda-coop-codegen/src/oxide.rs index 291e3f90004f..511f7a33bfe5 100644 --- a/rust/cuda-coop-codegen/src/oxide.rs +++ b/rust/cuda-coop-codegen/src/oxide.rs @@ -10,11 +10,12 @@ //! 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`]. +//! The protocol supports the raw i32 and typed Direct validity recipes +//! 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 crate::{Algorithm, Artifact, Dtype, Error, Layout, Operation, Provider, Request, Result}; use serde::{Deserialize, Serialize}; use std::io::{Read, Write}; use std::path::{Path, PathBuf}; @@ -33,7 +34,7 @@ struct Registration<'a> { algorithm: &'a str, } -#[derive(Debug, Deserialize, Serialize, PartialEq)] +#[derive(Clone, Copy, Debug, Deserialize, Serialize, PartialEq)] #[serde( tag = "kind", content = "bits", @@ -46,6 +47,40 @@ enum Scalar { Float(u16), } +impl Scalar { + fn dtype(self) -> Result { + match self { + Self::Signed(8) => Ok(Dtype::I8), + Self::Unsigned(8) => Ok(Dtype::U8), + Self::Signed(16) => Ok(Dtype::I16), + Self::Unsigned(16) => Ok(Dtype::U16), + Self::Signed(32) => Ok(Dtype::I32), + Self::Unsigned(32) => Ok(Dtype::U32), + Self::Signed(64) => Ok(Dtype::I64), + Self::Unsigned(64) => Ok(Dtype::U64), + Self::Float(32) => Ok(Dtype::F32), + Self::Float(64) => Ok(Dtype::F64), + _ => Err(Error::message("unsupported scalar kind or width")), + } + } +} +impl From for Scalar { + fn from(dtype: Dtype) -> Self { + match dtype { + Dtype::I8 => Self::Signed(8), + Dtype::U8 => Self::Unsigned(8), + Dtype::I16 => Self::Signed(16), + Dtype::U16 => Self::Unsigned(16), + Dtype::I32 => Self::Signed(32), + Dtype::U32 => Self::Unsigned(32), + Dtype::I64 => Self::Signed(64), + Dtype::U64 => Self::Unsigned(64), + Dtype::F32 => Self::Float(32), + Dtype::F64 => Self::Float(64), + } + } +} + #[derive(Debug, Deserialize, Serialize, PartialEq)] #[serde( tag = "kind", @@ -80,11 +115,6 @@ struct SpecializationRequest { 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]", @@ -95,6 +125,11 @@ impl SpecializationRequest { "load_filled" => Operation::LoadFilled, "store_full" => Operation::StoreFull, "store_partial" => Operation::StorePartial, + "load_full_validity_v1" => Operation::LoadFullValidity, + "load_masked_validity_v1" => Operation::LoadMaskedValidity, + "load_filled_validity_v1" => Operation::LoadFilledValidity, + "store_full_validity_v1" => Operation::StoreFullValidity, + "store_partial_validity_v1" => Operation::StorePartialValidity, _ => { return Err(Error::message(format!( "unsupported operation {:?}", @@ -112,7 +147,13 @@ impl SpecializationRequest { ))); } }; - Request::new(operation, algorithm, self.block[0], self.items_per_thread) + Request::new_typed( + self.dtype.dtype()?, + operation, + algorithm, + self.block[0], + self.items_per_thread, + ) } } @@ -153,21 +194,26 @@ struct BatchResponse { specializations: Vec, } -fn arguments(operation: Operation) -> Vec { +fn arguments(request: Request) -> Vec { + let operation = request.operation(); + let dtype = Scalar::from(request.dtype()); let mut arguments = vec![ if operation.is_load() { - Argument::ConstPointer(Scalar::Signed(32)) + Argument::ConstPointer(dtype) } else { - Argument::MutPointer(Scalar::Signed(32)) + Argument::MutPointer(dtype) }, - Argument::MutPointer(Scalar::Signed(32)), + Argument::MutPointer(dtype), ]; - if matches!(operation, Operation::LoadFilled | Operation::StorePartial) { - arguments.push(Argument::Value(Scalar::Signed(32))); + if operation.validity() { + arguments.push(Argument::MutPointer(Scalar::Unsigned(8))); } - if operation == Operation::LoadFilled { + if operation.counted() { arguments.push(Argument::Value(Scalar::Signed(32))); } + if operation.filled() { + arguments.push(Argument::Value(dtype)); + } arguments } @@ -334,7 +380,7 @@ fn write_response(requests: &[Request], artifact: Artifact, output: &mut impl Wr .map_err(|_| Error::message("request index exceeds u32"))?, identity: entry.identity, symbol: entry.symbol, - arguments: arguments(request.operation()), + arguments: arguments(*request), cub_layout: entry.cub_layout.try_into()?, scratch: entry.scratch.try_into()?, scratch_free: entry.scratch_free, @@ -363,6 +409,47 @@ fn write_response(requests: &[Request], artifact: Artifact, output: &mut impl Wr mod tests { use super::*; + #[test] + fn typed_validity_abi_preserves_width_and_separate_flags() { + for dtype in [ + Dtype::I8, + Dtype::U8, + Dtype::I16, + Dtype::U16, + Dtype::I32, + Dtype::U32, + Dtype::I64, + Dtype::U64, + Dtype::F32, + Dtype::F64, + ] { + let scalar = Scalar::from(dtype); + assert_eq!(scalar.dtype().unwrap(), dtype); + let request = SpecializationRequest { + operation: "load_filled_validity_v1".into(), + algorithm: "direct".into(), + dtype: scalar, + block: [128, 1, 1], + items_per_thread: 4, + } + .native() + .unwrap(); + assert_eq!( + arguments(request), + vec![ + Argument::ConstPointer(scalar), + Argument::MutPointer(scalar), + Argument::MutPointer(Scalar::Unsigned(8)), + Argument::Value(Scalar::Signed(32)), + Argument::Value(scalar), + ] + ); + } + for scalar in [Scalar::Signed(1), Scalar::Unsigned(128), Scalar::Float(16)] { + assert!(scalar.dtype().is_err()); + } + } + #[test] fn input_query_needs_no_installed_toolkit() { let checkout = Path::new(env!("CARGO_MANIFEST_DIR")); diff --git a/rust/cuda-coop-codegen/tests/check_validity_rejections.py b/rust/cuda-coop-codegen/tests/check_validity_rejections.py new file mode 100644 index 000000000000..97b9399c70ed --- /dev/null +++ b/rust/cuda-coop-codegen/tests/check_validity_rejections.py @@ -0,0 +1,120 @@ +# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. +# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +"""Compile private Direct validity controls and resource-epoch contrasts.""" + +import argparse +import json +import os +import subprocess +from pathlib import Path + +parser = argparse.ArgumentParser(description=__doc__) +parser.add_argument("output", type=Path) +args = parser.parse_args() +output = args.output.resolve() +record = json.loads((output / "command.json").read_text()) +source = output / "fixture/carrier/src/validity.rs" +original = source.read_text() +command = record["argv"].copy() +command[command.index("run")] = "build" +# build forwards binary selection to Cargo; run has its own --bin option. +command.insert(command.index("--bin"), "--") +env = os.environ.copy() +env.update( + CCCL_SOURCE=record.get("cccl_source", str(Path(__file__).resolve().parents[3])), + CUDA_OXIDE_BACKEND=record["backend"], + CUDA_HOME=record["cuda_root"], + CUDA_PATH=record["cuda_root"], +) +prefix = """// Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception +use coop_marker_fixture_api::{ThreadData, this_block, validity}; +use cuda_device::{CollectiveTiles, DisjointSlice, kernel, launch_contract, thread}; +#[kernel] +#[launch_contract(domain = 1, block = (128, 1, 1), dynamic_shared = inferred)] +pub fn probe(input: &[i32], mut output: DisjointSlice>, valid: usize) { + let group = this_block(); + let offset = thread::blockIdx_x() as usize * 512; + let mut values = ThreadData::<_, 4>::new(); +""" +suffix = "\n}\nfn main() {}\n" +cases = { + "varying_count": ( + "validity::load_masked(group, input, offset, thread::threadIdx_x() as usize, &mut values);", + "not proved block-uniform", + ), + "varying_default": ( + "validity::load_filled(group, input, offset, valid, thread::threadIdx_x() as i32, &mut values);", + "not proved block-uniform", + ), + "varying_offset": ( + "validity::load_masked(group, input, offset + thread::threadIdx_x() as usize, valid, &mut values);", + "not proved block-uniform", + ), + "missing_barrier": ( + "validity::load_filled(group, input, offset, valid, 0, &mut values);\n" + "validity::store_partial(group, &mut output, offset, valid, &mut values);\n" + "validity::reload_filled(group, &mut output, offset, valid, 0, &mut values);", + "requires a full-block barrier", + ), + "varying_participation": ( + "if thread::threadIdx_x() & 1 == 0 {\n" + "validity::load_masked(group, input, offset, valid, &mut values);\n}", + "lane-varying", + ), + "loop_without_completion": ( + "validity::load_filled(group, input, offset, valid, 0, &mut values);\n" + "let mut count = 2; while count > 0 {\n" + "validity::store_partial(group, &mut output, offset, valid, &mut values);\n" + "count -= 1; }", + "requires a full-block barrier", + ), + "explicit_completion": ( + "validity::load_filled(group, input, offset, valid, 0, &mut values);\n" + "validity::store_partial(group, &mut output, offset, valid, &mut values);\n" + "thread::sync_threads();\n" + "validity::reload_filled(group, &mut output, offset, valid, 0, &mut values);", + None, + ), + "reconverged_payload": ( + "let mut other = ThreadData::<_, 4>::new();\n" + "values.write(0, 7); other.write(0, 9);\n" + "let mut selected = if thread::threadIdx_x() & 1 == 0 { values } else { other };\n" + "validity::store_partial(group, &mut output, offset, 1, &mut selected);", + None, + ), +} +results = [] +try: + for name, (body, diagnostic) in cases.items(): + source.write_text(prefix + body + suffix) + result = subprocess.run( + command, + cwd=record["cwd"], + env=env, + text=True, + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + ) + (output / f"validity-{name}.log").write_text(result.stdout) + passed = ( + result.returncode == 0 + if diagnostic is None + else result.returncode != 0 and diagnostic in result.stdout + ) + results.append( + { + "case": name, + "expected": diagnostic, + "returncode": result.returncode, + "passed": passed, + } + ) + print(f"{'PASS' if passed else 'FAIL'} {name}", flush=True) + (output / "validity-rejections.json").write_text( + json.dumps(results, indent=2) + "\n" + ) +finally: + source.write_text(original) +raise SystemExit(0 if all(row["passed"] for row in results) else 1) diff --git a/rust/cuda-coop-codegen/tests/native.rs b/rust/cuda-coop-codegen/tests/native.rs index 4f8e16d38eb8..e61d2c2609cf 100644 --- a/rust/cuda-coop-codegen/tests/native.rs +++ b/rust/cuda-coop-codegen/tests/native.rs @@ -3,7 +3,7 @@ //! Real native compiler coverage; ordinary host tests need no CUDA installation. -use cuda_coop_codegen::{Algorithm, Layout, Operation, Provider, Request}; +use cuda_coop_codegen::{Algorithm, Dtype, Layout, Operation, Provider, Request}; use std::{env, path::Path, time::Instant}; #[test] @@ -11,6 +11,76 @@ 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()); } + assert!( + Request::new_typed(Dtype::U64, Operation::LoadFull, Algorithm::Direct, 128, 4).is_err() + ); + assert!(Request::new(Operation::LoadMaskedValidity, Algorithm::Transpose, 128, 4).is_err()); +} + +#[test] +#[ignore = "requires CUDA_HOME and a Linux CUDA toolkit with NVRTC LTO support"] +fn native_typed_validity_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 provider = Provider::from_checkout(&checkout, &toolkit)?; + let mut requests = Vec::new(); + for dtype in [ + Dtype::I8, + Dtype::U8, + Dtype::I16, + Dtype::U16, + Dtype::I32, + Dtype::U32, + Dtype::I64, + Dtype::U64, + Dtype::F32, + Dtype::F64, + ] { + for items in [1, 4] { + for operation in [ + Operation::LoadFullValidity, + Operation::LoadMaskedValidity, + Operation::LoadFilledValidity, + Operation::StoreFullValidity, + Operation::StorePartialValidity, + ] { + requests.push(Request::new_typed( + dtype, + operation, + Algorithm::Direct, + 128, + items, + )?); + } + } + } + let first = provider.compile(&requests, &target)?; + let repeated = provider.compile(&requests, &target)?; + assert_eq!(first.ltoir, repeated.ltoir); + assert_eq!(first.identity, repeated.identity); + assert_eq!(first.specializations.len(), requests.len()); + let mut identities = std::collections::HashSet::new(); + for entry in first.specializations { + assert!(identities.insert(entry.identity)); + assert_eq!( + entry.cub_layout, + Layout { + size: 1, + alignment: 1 + } + ); + assert_eq!( + entry.scratch, + Layout { + size: 0, + alignment: 1 + } + ); + assert!(entry.scratch_free); + assert!(!entry.mutates_payload); + } + Ok(()) } #[test] diff --git a/rust/cuda-coop-codegen/tests/oxide-fixture/README.md b/rust/cuda-coop-codegen/tests/oxide-fixture/README.md index cb3e919abe50..22afed6ab498 100644 --- a/rust/cuda-coop-codegen/tests/oxide-fixture/README.md +++ b/rust/cuda-coop-codegen/tests/oxide-fixture/README.md @@ -105,3 +105,18 @@ controls, descriptor/reference mutation, ambiguous origins and mixed markers, with positive cases for independent allocations and barriers in helpers. The native carrier binary also accepts `overlap` or `bounds`; each invocation must fail through a device trap, before an invalid native access. + +Add `--validity` to exercise the five typed Direct recipes with the actual +`ThreadData` source, through a test-only facade. The private bridge passes +separate addresses for opaque cells and exact-one flags; it exposes no +initialized slice. This route uses the checked resource marker, canonical +full-block carrier and ordinary owned async completion. It requires compiler +support for the `MaybeUninit` intrinsics reached by safe payload access. + +The cases cover all ten numeric types, items 1/4, exact typed defaults, +full and no-default masked Load, arbitrary flags outside a partial Store's +write domain, preservation after Store, genuinely short allocations and +Store/barrier/Load. A full Load followed by masked Load checks that old +initialization is reset. The binary accepts `hole0`, `hole2`, `hole255`, +`count` and `bounds`; each must fail in a separate process. The facade is +private test infrastructure, not a published safe collective API. diff --git a/rust/cuda-coop-codegen/tests/oxide-fixture/api/Cargo.toml b/rust/cuda-coop-codegen/tests/oxide-fixture/api/Cargo.toml index 2dce4d7120fc..d05b9378354c 100644 --- a/rust/cuda-coop-codegen/tests/oxide-fixture/api/Cargo.toml +++ b/rust/cuda-coop-codegen/tests/oxide-fixture/api/Cargo.toml @@ -16,3 +16,4 @@ cuda-coop-codegen = { path = "../../.." } [features] checked = [] carrier = [] +validity = ["carrier"] diff --git a/rust/cuda-coop-codegen/tests/oxide-fixture/api/build.rs b/rust/cuda-coop-codegen/tests/oxide-fixture/api/build.rs index 6c1abfb85736..6f0a40201f83 100644 --- a/rust/cuda-coop-codegen/tests/oxide-fixture/api/build.rs +++ b/rust/cuda-coop-codegen/tests/oxide-fixture/api/build.rs @@ -38,4 +38,14 @@ fn main() { println!("cargo::rustc-env=COOP_{tag}_{name}={registration}"); } } + for (operation, name) in [ + (Operation::LoadFullValidity, "LOAD_FULL"), + (Operation::LoadMaskedValidity, "LOAD_MASKED"), + (Operation::LoadFilledValidity, "LOAD_FILLED"), + (Operation::StoreFullValidity, "STORE_FULL"), + (Operation::StorePartialValidity, "STORE_PARTIAL"), + ] { + let registration = oxide::registration(&executable, operation, Algorithm::Direct).unwrap(); + println!("cargo::rustc-env=COOP_VALIDITY_{name}={registration}"); + } } diff --git a/rust/cuda-coop-codegen/tests/oxide-fixture/api/src/carrier.rs b/rust/cuda-coop-codegen/tests/oxide-fixture/api/src/carrier.rs index e46c515f0d61..d23d219467e5 100644 --- a/rust/cuda-coop-codegen/tests/oxide-fixture/api/src/carrier.rs +++ b/rust/cuda-coop-codegen/tests/oxide-fixture/api/src/carrier.rs @@ -6,8 +6,8 @@ //! epochs, while this trusted boundary checks the unchanged offset and extent. use cuda_device::{CollectiveTiles, DisjointSlice, library, thread}; -#[inline(never)] -fn extent(len: usize, offset: usize, valid: usize) { +#[inline(always)] +pub(crate) fn extent(len: usize, offset: usize, valid: usize) { let threads = thread::blockDim_x() as usize * thread::blockDim_y() as usize * thread::blockDim_z() as usize; @@ -16,8 +16,8 @@ fn extent(len: usize, offset: usize, valid: usize) { assert!(offset <= len && valid <= len - offset); } -#[inline(never)] -fn containment(offset: usize, valid: usize) { +#[inline(always)] +pub(crate) fn containment(offset: usize, valid: usize) { let threads = thread::blockDim_x() as usize * thread::blockDim_y() as usize * thread::blockDim_z() as usize; 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 index 63ae2fe618d9..df61f93fc982 100644 --- a/rust/cuda-coop-codegen/tests/oxide-fixture/api/src/lib.rs +++ b/rust/cuda-coop-codegen/tests/oxide-fixture/api/src/lib.rs @@ -7,6 +7,53 @@ #[cfg(feature = "carrier")] pub mod carrier; +// Test-source facade: compile the exact payload/group implementation without +// introducing a public operation or a second representation of ThreadData. +#[cfg(feature = "validity")] +mod alignment { + include!(concat!( + env!("CCCL_SOURCE"), + "/rust/cuda-coop/src/alignment.rs" + )); +} +#[cfg(feature = "validity")] +pub use alignment::{Alignment, SupportedAlignment}; +#[cfg(feature = "validity")] +pub mod thread_data { + include!(concat!( + env!("CCCL_SOURCE"), + "/rust/cuda-coop/src/thread_data.rs" + )); + mod block_io { + include!(concat!( + env!("CCCL_SOURCE"), + "/rust/cuda-coop/src/thread_data/block_io.rs" + )); + } + + // Fault injection and flag observation belong only to this private fixture. + pub fn flag(values: &ThreadData, index: usize) -> u8 { + values.initialized[index] + } + pub fn corrupt_flag( + values: &mut ThreadData, + index: usize, + flag: u8, + ) { + values.initialized[index] = flag; + } +} +#[cfg(feature = "validity")] +pub use thread_data::{Numeric, ThreadData}; +#[cfg(feature = "validity")] +mod group { + include!(concat!(env!("CCCL_SOURCE"), "/rust/cuda-coop/src/group.rs")); +} +#[cfg(feature = "validity")] +pub use group::*; +#[cfg(feature = "validity")] +pub mod validity; + macro_rules! registration { ($name:ident, $variable:literal) => { fn $name() -> &'static str { @@ -22,6 +69,16 @@ 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"); +#[cfg(feature = "validity")] +registration!(validity_load_full, "COOP_VALIDITY_LOAD_FULL"); +#[cfg(feature = "validity")] +registration!(validity_load_masked, "COOP_VALIDITY_LOAD_MASKED"); +#[cfg(feature = "validity")] +registration!(validity_load_filled, "COOP_VALIDITY_LOAD_FILLED"); +#[cfg(feature = "validity")] +registration!(validity_store_full, "COOP_VALIDITY_STORE_FULL"); +#[cfg(feature = "validity")] +registration!(validity_store_partial, "COOP_VALIDITY_STORE_PARTIAL"); /// # Safety /// All threads participate with one block-wide input origin, valid count and diff --git a/rust/cuda-coop-codegen/tests/oxide-fixture/api/src/validity.rs b/rust/cuda-coop-codegen/tests/oxide-fixture/api/src/validity.rs new file mode 100644 index 000000000000..7891d885e423 --- /dev/null +++ b/rust/cuda-coop-codegen/tests/oxide-fixture/api/src/validity.rs @@ -0,0 +1,174 @@ +// Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +//! Private Direct integration boundary. This is not a public safe operation. +//! The checked marker proves participation, controls, allocation origin and +//! epochs. CUB uses opaque cells; the provider owns mask reset and preflight. +use crate::{Alignment, Block, Numeric, SupportedAlignment, ThreadData, ThreadGroup}; +use cuda_device::{CollectiveTiles, DisjointSlice, library, thread}; + +#[inline(always)] +fn extent(len: usize, offset: usize, valid: usize) { + crate::carrier::extent::(len, offset, valid); + assert!(offset <= isize::MAX as usize / core::mem::size_of::()); +} + +#[inline(always)] +fn tile() -> usize { + let threads = thread::blockDim_x() as usize + * thread::blockDim_y() as usize + * thread::blockDim_z() as usize; + assert!(ITEMS > 0 && threads <= usize::MAX / ITEMS); + threads * ITEMS +} + +#[inline(always)] +pub fn load_full( + _group: ThreadGroup, + source: &[T], + offset: usize, + values: &mut ThreadData, +) where + Alignment: SupportedAlignment, +{ + extent::(source.len(), offset, tile::()); + let payload = values.block_io_parts(); + // SAFETY: checked bounds precede pointer formation, source is the pointer's + // allocation witness, and the provider initializes every cell and flag. + unsafe { + library::checked_block_resource_call_v1::( + crate::validity_load_full, + source, + (source.as_ptr().add(offset),), + payload, + ); + } +} + +#[inline(always)] +pub fn load_masked( + _group: ThreadGroup, + source: &[T], + offset: usize, + valid: usize, + values: &mut ThreadData, +) where + Alignment: SupportedAlignment, +{ + extent::(source.len(), offset, valid); + let payload = values.block_io_parts(); + // SAFETY: only the checked prefix is read. The provider replaces all flags + // and copies numeric bytes only for initialized cells. + unsafe { + library::checked_block_resource_call_v1::( + crate::validity_load_masked, + source, + (source.as_ptr().add(offset), valid as i32), + payload, + ); + } +} + +#[inline(always)] +pub fn load_filled( + _group: ThreadGroup, + source: &[T], + offset: usize, + valid: usize, + default: T, + values: &mut ThreadData, +) where + Alignment: SupportedAlignment, +{ + extent::(source.len(), offset, valid); + let payload = values.block_io_parts(); + // SAFETY: bounds and typed default match the exact native ABI. The provider + // initializes every cell, including the out-of-prefix defaults. + unsafe { + library::checked_block_resource_call_v1::( + crate::validity_load_filled, + source, + (source.as_ptr().add(offset), valid as i32, default), + payload, + ); + } +} + +#[inline(always)] +pub fn store_full( + _group: ThreadGroup, + target: &mut DisjointSlice>, + offset: usize, + values: &mut ThreadData, +) where + Alignment: SupportedAlignment, +{ + let valid = tile::(); + extent::(target.len(), offset, valid); + crate::carrier::containment::(offset, valid); + let pointer = target.as_mut_ptr(); + let payload = values.block_io_parts(); + // SAFETY: this carrier owns the checked range. The provider checks every + // written slot's exact-one flag before interpreting its bytes as T. + unsafe { + library::checked_block_resource_call_v1::( + crate::validity_store_full, + target, + (pointer.add(offset),), + payload, + ); + } +} + +#[inline(always)] +pub fn store_partial( + _group: ThreadGroup, + target: &mut DisjointSlice>, + offset: usize, + valid: usize, + values: &mut ThreadData, +) where + Alignment: SupportedAlignment, +{ + extent::(target.len(), offset, valid); + crate::carrier::containment::(offset, valid); + let pointer = target.as_mut_ptr(); + let payload = values.block_io_parts(); + // SAFETY: the guarded native iterator forms an effective pointer only for + // a written slot. Preflight covers that same domain, with no tail reads. + unsafe { + library::checked_block_resource_call_v1::( + crate::validity_store_partial, + target, + (pointer.add(offset), valid as i32), + payload, + ); + } +} + +#[inline(always)] +pub fn reload_filled( + _group: ThreadGroup, + source: &mut DisjointSlice>, + offset: usize, + valid: usize, + default: T, + values: &mut ThreadData, +) where + Alignment: SupportedAlignment, +{ + extent::(source.len(), offset, valid); + crate::carrier::containment::(offset, valid); + let pointer = source.as_mut_ptr(); + let payload = values.block_io_parts(); + // SAFETY: the immutable witness binds the same allocation. A prior Store + // requires explicit block completion, enforced by the compiler. + unsafe { + library::checked_block_resource_call_v1::( + crate::validity_load_filled, + &*source, + (pointer.add(offset).cast_const(), valid as i32, default), + payload, + ); + } +} diff --git a/rust/cuda-coop-codegen/tests/oxide-fixture/carrier/Cargo.toml b/rust/cuda-coop-codegen/tests/oxide-fixture/carrier/Cargo.toml index fe068a69bb79..0f96fdca727d 100644 --- a/rust/cuda-coop-codegen/tests/oxide-fixture/carrier/Cargo.toml +++ b/rust/cuda-coop-codegen/tests/oxide-fixture/carrier/Cargo.toml @@ -6,6 +6,7 @@ name = "coop-carrier-fixture" version = "0.0.0" edition = "2024" publish = false +default-run = "coop-carrier-fixture" [dependencies] cuda-device.workspace = true @@ -13,3 +14,11 @@ cuda-host = { workspace = true, features = ["async"] } cuda-core.workspace = true cuda-async.workspace = true coop-marker-fixture-api = { path = "../api", features = ["carrier"] } + +[features] +validity = ["coop-marker-fixture-api/validity"] + +[[bin]] +name = "coop-validity-fixture" +path = "src/validity.rs" +required-features = ["validity"] diff --git a/rust/cuda-coop-codegen/tests/oxide-fixture/carrier/src/validity.rs b/rust/cuda-coop-codegen/tests/oxide-fixture/carrier/src/validity.rs new file mode 100644 index 000000000000..d28290bbe97f --- /dev/null +++ b/rust/cuda-coop-codegen/tests/oxide-fixture/carrier/src/validity.rs @@ -0,0 +1,528 @@ +// Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +use coop_marker_fixture_api::{Numeric, ThreadData, this_block, thread_data, validity}; +use cuda_device::{CollectiveTiles, DisjointSlice, kernel, launch_contract, thread}; +use cuda_host::cuda_module; + +#[cuda_module] +mod kernels { + use super::*; + + #[inline(never)] + fn exercise( + input: &[T], + output: &mut DisjointSlice>, + valid: usize, + shift: usize, + default: T, + mode: u32, + bad_flag: u8, + ) { + let block = this_block(); + let threads = thread::blockDim_x() as usize; + let lane = thread::threadIdx_x() as usize; + let offset = thread::blockIdx_x() as usize * threads * 4 + shift; + let mut values = ThreadData::<_, ITEMS_PER_THREAD>::new(); + match mode { + 0 => validity::load_full(block, input, offset, &mut values), + 1 | 5 | 6 => { + validity::load_full(block, input, offset, &mut values); + validity::load_masked(block, input, offset, valid, &mut values); + } + 2 | 4 => validity::load_filled(block, input, offset, valid, default, &mut values), + _ => { + // Arbitrary masks and a lane-divergent move/selection. Only the + // written prefix is initialized; flags outside it are noncanonical. + let mut other = ThreadData::<_, ITEMS_PER_THREAD>::new(); + let mut item = 0; + while item < ITEMS_PER_THREAD { + if lane * ITEMS_PER_THREAD + item < valid { + values.write(item, default); + other.write(item, default); + } else { + thread_data::corrupt_flag(&mut values, item, 2); + thread_data::corrupt_flag(&mut other, item, 255); + } + item += 1; + } + values = if lane % 2 == 0 { values } else { other }; + } + } + if bad_flag != 1 { + thread_data::corrupt_flag(&mut values, 0, bad_flag); + } + let before = values; + if mode == 0 || mode == 2 { + validity::store_full(block, output, offset, &mut values); + } else { + let written = if mode == 5 { + valid + 1 + } else if mode == 6 { + valid / 2 + } else { + valid + }; + validity::store_partial(block, output, offset, written, &mut values); + } + let mut observations_ok = true; + let mut item = 0; + while item < ITEMS_PER_THREAD { + observations_ok &= thread_data::flag(&values, item) == thread_data::flag(&before, item); + if let Some(value) = before.get(item) { + observations_ok &= values.get(item) == Some(value); + } else { + observations_ok &= values.get(item).is_none(); + } + if mode == 1 || mode == 6 { + observations_ok &= + values.get(item).is_some() == (lane * ITEMS_PER_THREAD + item < valid); + } + item += 1; + } + if mode == 4 { + thread::sync_threads(); + validity::reload_filled(block, output, offset, valid, default, &mut values); + thread::sync_threads(); + validity::store_partial(block, output, offset, valid, &mut values); + } + assert!(observations_ok); + // Final writes complete through the owned async host operation. + } + + #[kernel] + #[launch_contract(domain = 1, block = (128, 1, 1), dynamic_shared = inferred)] + pub fn copy_i8_128( + input: &[i8], + mut output: DisjointSlice>, + valid: usize, + shift: usize, + default: i8, + mode: u32, + bad_flag: u8, + ) { + exercise::<_, ITEMS_PER_THREAD>(input, &mut output, valid, shift, default, mode, bad_flag); + } + + #[kernel] + #[launch_contract(domain = 1, block = (128, 1, 1), dynamic_shared = inferred)] + pub fn copy_u8_128( + input: &[u8], + mut output: DisjointSlice>, + valid: usize, + shift: usize, + default: u8, + mode: u32, + bad_flag: u8, + ) { + exercise::<_, ITEMS_PER_THREAD>(input, &mut output, valid, shift, default, mode, bad_flag); + } + + #[kernel] + #[launch_contract(domain = 1, block = (128, 1, 1), dynamic_shared = inferred)] + pub fn copy_i16_128( + input: &[i16], + mut output: DisjointSlice>, + valid: usize, + shift: usize, + default: i16, + mode: u32, + bad_flag: u8, + ) { + exercise::<_, ITEMS_PER_THREAD>(input, &mut output, valid, shift, default, mode, bad_flag); + } + + #[kernel] + #[launch_contract(domain = 1, block = (128, 1, 1), dynamic_shared = inferred)] + pub fn copy_u16_128( + input: &[u16], + mut output: DisjointSlice>, + valid: usize, + shift: usize, + default: u16, + mode: u32, + bad_flag: u8, + ) { + exercise::<_, ITEMS_PER_THREAD>(input, &mut output, valid, shift, default, mode, bad_flag); + } + + #[kernel] + #[launch_contract(domain = 1, block = (128, 1, 1), dynamic_shared = inferred)] + pub fn copy_i32_128( + input: &[i32], + mut output: DisjointSlice>, + valid: usize, + shift: usize, + default: i32, + mode: u32, + bad_flag: u8, + ) { + exercise::<_, ITEMS_PER_THREAD>(input, &mut output, valid, shift, default, mode, bad_flag); + } + + #[kernel] + #[launch_contract(domain = 1, block = (128, 1, 1), dynamic_shared = inferred)] + pub fn copy_u32_128( + input: &[u32], + mut output: DisjointSlice>, + valid: usize, + shift: usize, + default: u32, + mode: u32, + bad_flag: u8, + ) { + exercise::<_, ITEMS_PER_THREAD>(input, &mut output, valid, shift, default, mode, bad_flag); + } + + #[kernel] + #[launch_contract(domain = 1, block = (128, 1, 1), dynamic_shared = inferred)] + pub fn copy_i64_128( + input: &[i64], + mut output: DisjointSlice>, + valid: usize, + shift: usize, + default: i64, + mode: u32, + bad_flag: u8, + ) { + exercise::<_, ITEMS_PER_THREAD>(input, &mut output, valid, shift, default, mode, bad_flag); + } + + #[kernel] + #[launch_contract(domain = 1, block = (128, 1, 1), dynamic_shared = inferred)] + pub fn copy_u64_128( + input: &[u64], + mut output: DisjointSlice>, + valid: usize, + shift: usize, + default: u64, + mode: u32, + bad_flag: u8, + ) { + exercise::<_, ITEMS_PER_THREAD>(input, &mut output, valid, shift, default, mode, bad_flag); + } + + #[kernel] + #[launch_contract(domain = 1, block = (128, 1, 1), dynamic_shared = inferred)] + pub fn copy_f32_128( + input: &[f32], + mut output: DisjointSlice>, + valid: usize, + shift: usize, + default: f32, + mode: u32, + bad_flag: u8, + ) { + exercise::<_, ITEMS_PER_THREAD>(input, &mut output, valid, shift, default, mode, bad_flag); + } + + #[kernel] + #[launch_contract(domain = 1, block = (128, 1, 1), dynamic_shared = inferred)] + pub fn copy_f64_128( + input: &[f64], + mut output: DisjointSlice>, + valid: usize, + shift: usize, + default: f64, + mode: u32, + bad_flag: u8, + ) { + exercise::<_, ITEMS_PER_THREAD>(input, &mut output, valid, shift, default, mode, bad_flag); + } + + #[kernel] + #[launch_contract(domain = 1, block = (64, 1, 1), dynamic_shared = inferred)] + pub fn copy_i32_64( + input: &[i32], + mut output: DisjointSlice>, + valid: usize, + shift: usize, + default: i32, + mode: u32, + bad_flag: u8, + ) { + exercise::<_, ITEMS_PER_THREAD>(input, &mut output, valid, shift, default, mode, bad_flag); + } +} + +fn main() -> Result<(), Box> { + use cuda_async::simt::device_context::{init_device_contexts, with_cuda_context}; + use cuda_async::simt::device_operation::DeviceOperation; + use cuda_core::{DeviceBuffer, LaunchConfig1D}; + init_device_contexts(0, 1)?; + let context = with_cuda_context(0, |context| context.clone())?; + let stream = context.default_stream(); + // SAFETY: embedded kernels and these typed declarations are compiled together. + let module = unsafe { kernels::load_async(0)? }; + let trap = std::env::args().nth(1); + let mut cases = 0; + macro_rules! run { + ($t:ty, $prepare:ident, $launch:ident, $block:literal, $items:literal, $default:expr) => {{ + let default: $t = $default; + for groups in [1_u32, 3] { + let prepared = + module.$prepare::<$items>(LaunchConfig1D::auto_shared(groups, $block))?; + assert_eq!(prepared.__raw_config().shared_mem_bytes, 0); + for mode in [0, 1, 2, 3, 4, 6] { + for valid in [0, 1, $block * $items - 1, $block * $items] { + if (mode == 0) && valid != $block * $items { + continue; + } + let shift = if $items == 1 { 7 } else { 0 }; + let end = (groups as usize - 1) * $block * 4 + shift; + let source_len = end + + if mode <= 1 || mode == 6 { + $block * $items + } else { + valid + }; + let target_len = end + if mode <= 2 { $block * $items } else { valid }; + let input: Vec<$t> = (0..source_len).map(|i| (i % 61 + 1) as $t).collect(); + let source = DeviceBuffer::from_host(&stream, &input)?; + let target = DeviceBuffer::from_host(&stream, &vec![0 as $t; target_len])?; + let (_, target) = module + .$launch::<$items, _, _>( + &prepared, source, target, valid, shift, default, mode, 1, + ) + .sync()?; + let mut result = vec![0 as $t; target_len]; + target.copy_to_host(&stream, &mut result)?; + stream.synchronize()?; + for (index, &value) in result.iter().enumerate() { + let within = index % ($block * 4); + let written = if mode == 0 || mode == 2 { + $block * $items + } else if mode == 6 { + valid / 2 + } else { + valid + }; + let expected = if within >= shift && within - shift < written { + if mode == 3 || (mode == 2 && within - shift >= valid) { + default + } else { + input[index] + } + } else { + 0 as $t + }; + assert_eq!( + value.to_ne_bytes(), + expected.to_ne_bytes(), + "type={} items={} mode={} valid={} index={}", + stringify!($t), + $items, + mode, + valid, + index + ); + } + cases += 1; + } + } + } + }}; + } + if let Some(trap) = trap { + let plan = module.prepare_copy_i32_128::<4>(LaunchConfig1D::auto_shared(1, 128))?; + let len = if trap == "ownership" { 1024 } else { 512 }; + let source = DeviceBuffer::from_host(&stream, &vec![11_i32; len])?; + let target = DeviceBuffer::from_host(&stream, &vec![0_i32; len])?; + let (valid, shift, mode, flag) = match trap.as_str() { + "hole0" => (512, 0, 0, 0), + "hole2" => (512, 0, 0, 2), + "hole255" => (512, 0, 0, 255), + "count" => (513, 0, 1, 1), + "bounds" => (512, 1, 1, 1), + "ownership" => (512, 1, 0, 1), + "different-count" => (127, 0, 5, 1), + _ => panic!("unknown trap mode"), + }; + let _ = module + .copy_i32_128_async_owned::<4, _, _>( + &plan, source, target, valid, shift, -7, mode, flag, + ) + .sync()?; + panic!("expected trap returned successfully"); + } + run!( + i8, + prepare_copy_i8_128, + copy_i8_128_async_owned, + 128, + 1, + i8::MIN + ); + run!( + i8, + prepare_copy_i8_128, + copy_i8_128_async_owned, + 128, + 4, + i8::MIN + ); + run!( + u8, + prepare_copy_u8_128, + copy_u8_128_async_owned, + 128, + 1, + u8::MAX + ); + run!( + u8, + prepare_copy_u8_128, + copy_u8_128_async_owned, + 128, + 4, + u8::MAX + ); + run!( + i16, + prepare_copy_i16_128, + copy_i16_128_async_owned, + 128, + 1, + i16::MIN + ); + run!( + i16, + prepare_copy_i16_128, + copy_i16_128_async_owned, + 128, + 4, + i16::MIN + ); + run!( + u16, + prepare_copy_u16_128, + copy_u16_128_async_owned, + 128, + 1, + u16::MAX + ); + run!( + u16, + prepare_copy_u16_128, + copy_u16_128_async_owned, + 128, + 4, + u16::MAX + ); + run!( + i32, + prepare_copy_i32_128, + copy_i32_128_async_owned, + 128, + 1, + i32::MIN + ); + run!( + i32, + prepare_copy_i32_128, + copy_i32_128_async_owned, + 128, + 4, + i32::MIN + ); + run!( + u32, + prepare_copy_u32_128, + copy_u32_128_async_owned, + 128, + 1, + u32::MAX + ); + run!( + u32, + prepare_copy_u32_128, + copy_u32_128_async_owned, + 128, + 4, + u32::MAX + ); + run!( + i64, + prepare_copy_i64_128, + copy_i64_128_async_owned, + 128, + 1, + i64::MIN + 17 + ); + run!( + i64, + prepare_copy_i64_128, + copy_i64_128_async_owned, + 128, + 4, + i64::MIN + 17 + ); + run!( + u64, + prepare_copy_u64_128, + copy_u64_128_async_owned, + 128, + 1, + u64::MAX - 17 + ); + run!( + u64, + prepare_copy_u64_128, + copy_u64_128_async_owned, + 128, + 4, + u64::MAX - 17 + ); + run!( + f32, + prepare_copy_f32_128, + copy_f32_128_async_owned, + 128, + 1, + -0.0_f32 + ); + run!( + f32, + prepare_copy_f32_128, + copy_f32_128_async_owned, + 128, + 4, + -0.0_f32 + ); + run!( + f64, + prepare_copy_f64_128, + copy_f64_128_async_owned, + 128, + 1, + -0.0_f64 + ); + run!( + f64, + prepare_copy_f64_128, + copy_f64_128_async_owned, + 128, + 4, + -0.0_f64 + ); + run!( + i32, + prepare_copy_i32_64, + copy_i32_64_async_owned, + 64, + 1, + -1234567 + ); + run!( + i32, + prepare_copy_i32_64, + copy_i32_64_async_owned, + 64, + 4, + -1234567 + ); + println!( + "PASS: {cases} Direct validity cases, ten numeric types, exact flags, short tails, owned async" + ); + Ok(()) +} diff --git a/rust/cuda-coop-codegen/tests/run_oxide_fixture.py b/rust/cuda-coop-codegen/tests/run_oxide_fixture.py index c2899825874b..b5176c2bad92 100644 --- a/rust/cuda-coop-codegen/tests/run_oxide_fixture.py +++ b/rust/cuda-coop-codegen/tests/run_oxide_fixture.py @@ -28,6 +28,11 @@ parser.add_argument( "--carrier", action="store_true", help="exercise private collective-only carriers" ) +parser.add_argument( + "--validity", + action="store_true", + help="exercise private typed Direct payload validity", +) args = parser.parse_args() source = Path(__file__).resolve().parent checkout = source.parents[2] @@ -74,11 +79,13 @@ args.arch, "--materialize-cubin", ] -if args.root_context: +if args.validity: + command.extend(["--features", "validity", "--bin", "coop-validity-fixture"]) +elif args.root_context: command.extend(["--features", "root-context"]) elif args.checked: command.extend(["--features", "checked"]) -consumer = fixture / ("carrier" if args.carrier else "consumer") +consumer = fixture / ("carrier" if args.carrier or args.validity else "consumer") with (output / "build-run.log").open("w") as log: result = subprocess.run( command, cwd=consumer, env=env, stdout=log, stderr=subprocess.STDOUT @@ -89,6 +96,7 @@ "argv": command, "cwd": str(consumer), "oxide_root": str(oxide), + "cccl_source": str(checkout), "backend": env["CUDA_OXIDE_BACKEND"], "cuda_root": env["CUDA_HOME"], "returncode": result.returncode, diff --git a/rust/cuda-coop/src/thread_data.rs b/rust/cuda-coop/src/thread_data.rs index 4d955cd2cbb2..95f5170d63ec 100644 --- a/rust/cuda-coop/src/thread_data.rs +++ b/rust/cuda-coop/src/thread_data.rs @@ -126,7 +126,8 @@ where #[inline] pub fn get(&self, index: usize) -> Option { if index < ITEMS && self.initialized[index] == 1 { - // SAFETY: Only write sets this flag, after initializing the value. + // SAFETY: write and the private native boundary set this flag only + // after initializing the value. // Copies preserve both arrays. Numeric types are Copy. Some(unsafe { self.values[index].assume_init() }) } else { diff --git a/rust/cuda-coop/src/thread_data/block_io.rs b/rust/cuda-coop/src/thread_data/block_io.rs new file mode 100644 index 000000000000..1bbc9b447fc7 --- /dev/null +++ b/rust/cuda-coop/src/thread_data/block_io.rs @@ -0,0 +1,21 @@ +// Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +// Private native-boundary bridge, compiled only by the coordinated fixture +// until the complete collective API is qualified. These addresses describe +// opaque MaybeUninit cells and separate exact-one flags, never a T slice. +use super::{Alignment, Numeric, SupportedAlignment, ThreadData}; + +impl ThreadData +where + Alignment: SupportedAlignment, +{ + #[inline(always)] + pub(crate) fn block_io_parts(&mut self) -> (*mut T, *mut u8) { + let () = Self::POSITIVE_EXTENT; + ( + self.values.as_mut_ptr().cast::(), + self.initialized.as_mut_ptr(), + ) + } +} From 951b7b7a92431a9e9c62596a9aa5f161fe3937ce Mon Sep 17 00:00:00 2001 From: Trent Nelson Date: Tue, 20 Oct 2026 09:49:23 -0700 Subject: [PATCH 19/33] test(rust): exercise storage reuse across runtime loops Signed-off-by: Trent Nelson --- .../tests/check_storage_rejections.py | 5 ++-- .../tests/oxide-fixture/README.md | 5 +++- .../oxide-fixture/consumer/src/storage.rs | 21 +++++++++++++--- .../tests/oxide-fixture/helper/src/storage.rs | 24 ++++++++++++++++++- 4 files changed, 48 insertions(+), 7 deletions(-) diff --git a/rust/cuda-coop-codegen/tests/check_storage_rejections.py b/rust/cuda-coop-codegen/tests/check_storage_rejections.py index f67f47730ab5..dde9b2417f83 100644 --- a/rust/cuda-coop-codegen/tests/check_storage_rejections.py +++ b/rust/cuda-coop-codegen/tests/check_storage_rejections.py @@ -204,11 +204,13 @@ continue source.write_text(prefix + body + suffix) log_path = output / f"storage-{name}.log" + case_env = env.copy() + case_env["CARGO_TARGET_DIR"] = str(output / "storage-targets" / name) with log_path.open("w") as log: result = subprocess.run( command, cwd=record["cwd"], - env=env, + env=case_env, stdout=log, stderr=subprocess.STDOUT, ) @@ -226,7 +228,6 @@ print(f"{name}: {'PASS' if passed else 'FAIL'}", flush=True) if not passed: print(text[-5000:], flush=True) - break finally: source.write_text(original) (output / "storage-compile-results.json").write_text( diff --git a/rust/cuda-coop-codegen/tests/oxide-fixture/README.md b/rust/cuda-coop-codegen/tests/oxide-fixture/README.md index a5a99baa5e82..e595bcc77fd0 100644 --- a/rust/cuda-coop-codegen/tests/oxide-fixture/README.md +++ b/rust/cuda-coop-codegen/tests/oxide-fixture/README.md @@ -115,7 +115,9 @@ would incorrectly change ownership if the fixed carrier capacity were ignored. The `--storage` runner route exercises the internal full-block storage boundary with the unchanged numeric provider: shared/manual/automatic and exclusive policies, distinct helper-created descriptors, typed raw reservations, and -Direct/Transpose operations at one and four items per thread. It checks three +Direct/Transpose operations at one and four items per thread. Runtime loops +exercise reuse across iterations; reconfiguration before use checks alignment, +and automatic Direct operations remain scratch-free. It checks three blocks, boundary valid counts and inferred launch bytes through owned async completion. The wrapper remains private and unsafe for its stated memory and initialization obligations. @@ -123,5 +125,6 @@ initialization obligations. Add `--compile-only` to prepare an isolated fixture without GPU execution. `check_storage_rejections.py OUTPUT` uses that copy for paired positive and negative layout/reuse/output-epoch cases and restores its source afterward. +Each case uses its own Cargo target directory. It checks intended diagnostics, including whole-plan capacity, live views, release versus completion and scratch-free output conflicts. diff --git a/rust/cuda-coop-codegen/tests/oxide-fixture/consumer/src/storage.rs b/rust/cuda-coop-codegen/tests/oxide-fixture/consumer/src/storage.rs index 9b7feab4bb01..c1ab62f93231 100644 --- a/rust/cuda-coop-codegen/tests/oxide-fixture/consumer/src/storage.rs +++ b/rust/cuda-coop-codegen/tests/oxide-fixture/consumer/src/storage.rs @@ -90,14 +90,21 @@ pub fn run() -> Result<(), Box> { } else { single }; - assert_eq!(plan.__raw_config().shared_mem_bytes, expected); - if expected > 0 { + let actual = plan.__raw_config().shared_mem_bytes; + if $mode == 7 { + // Either descriptor may be placed first after the native + // prefix. Both orders obey the same alignment requirements. + assert!(actual == expected || actual == 64 + 512 * $i + single); + } else { + assert_eq!(actual, expected); + } + if actual > 0 { assert!( module .prepare_copy::<$i, $transpose, $mode>(LaunchConfig1D::new( 3, 128, - expected - 1 + actual - 1 )) .is_err() ); @@ -122,6 +129,8 @@ pub fn run() -> Result<(), Box> { (((index % 512 / $i + 1) % 128) * 3 + 1) as i32 } else if $mode == 7 { (index % 512 / $i % 8 + 10) as i32 + } else if $mode == 8 { + (1 + (valid & 3)) as i32 } else { 0 } @@ -156,6 +165,12 @@ pub fn run() -> Result<(), Box> { check!(4, true, 6); check!(1, true, 7); check!(4, true, 7); + check!(1, true, 8); + check!(4, true, 8); + check!(1, true, 9); + check!(4, true, 9); + check!(1, false, 2); + check!(4, false, 2); println!("PASS: {cases} storage cases with three blocks and boundary counts"); Ok(()) } diff --git a/rust/cuda-coop-codegen/tests/oxide-fixture/helper/src/storage.rs b/rust/cuda-coop-codegen/tests/oxide-fixture/helper/src/storage.rs index 4c55b0de622e..b51cad4c48e5 100644 --- a/rust/cuda-coop-codegen/tests/oxide-fixture/helper/src/storage.rs +++ b/rust/cuda-coop-codegen/tests/oxide-fixture/helper/src/storage.rs @@ -95,7 +95,7 @@ pub unsafe fn roundtrip(output, offset, valid, &mut values, &mut scratch); } - _ => { + 7 => { let prefix = cuda_device::DynamicSharedArray::<[u64; 8], 32>::get_fixed().cast::(); let lane = thread::threadIdx_x() as usize; @@ -116,6 +116,28 @@ pub unsafe fn roundtrip(output, offset, valid, &mut values, &mut first); } + 8 => { + // Runtime trip count keeps reuse of these same static sites + // observable, including the planned barrier on the backedge. + let mut iteration = 0; + while iteration < 2 + (valid & 3) { + load::(input, offset, valid, &mut values, ()); + let mut item = 0; + while item < ITEMS { + values[item] += iteration as i32; + item += 1; + } + store::(output, offset, valid, &mut values, ()); + iteration += 1; + } + } + _ => { + let original = descriptor(); + let mut scratch: library::StorageDescriptor<0, 64, true, false> = + library::storage_reconfigure_v1(original); + load::(input, offset, valid, &mut values, &mut scratch); + store::(output, offset, valid, &mut values, &mut scratch); + } } } } From 5ce6a050ca5a2ceca95aea891b77dfabde2d7fb0 Mon Sep 17 00:00:00 2001 From: Trent Nelson Date: Tue, 20 Oct 2026 10:35:16 -0700 Subject: [PATCH 20/33] Keep payload indexing guards device-compatible Signed-off-by: Trent Nelson --- rust/cuda-coop/src/thread_data.rs | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/rust/cuda-coop/src/thread_data.rs b/rust/cuda-coop/src/thread_data.rs index 4d955cd2cbb2..1fef6629f299 100644 --- a/rust/cuda-coop/src/thread_data.rs +++ b/rust/cuda-coop/src/thread_data.rs @@ -165,8 +165,8 @@ where #[inline] fn index(&self, index: usize) -> &T { - assert!(index < ITEMS, "ThreadData index out of bounds"); - assert!(self.initialized[index] == 1, "ThreadData slot is unwritten"); + assert!(index < ITEMS); + assert!(self.initialized[index] == 1); // SAFETY: The bounds and exact initialization flag were checked above. unsafe { self.values[index].assume_init_ref() } } @@ -179,8 +179,8 @@ where { #[inline] fn index_mut(&mut self, index: usize) -> &mut T { - assert!(index < ITEMS, "ThreadData index out of bounds"); - assert!(self.initialized[index] == 1, "ThreadData slot is unwritten"); + assert!(index < ITEMS); + assert!(self.initialized[index] == 1); // SAFETY: The checks prove initialization; &mut self is exclusive. unsafe { self.values[index].assume_init_mut() } } From 63d58b83d0bc577070e975bad1fe4abad30622d4 Mon Sep 17 00:00:00 2001 From: Trent Nelson Date: Tue, 20 Oct 2026 10:38:16 -0700 Subject: [PATCH 21/33] Pack payload validity bits for device access Signed-off-by: Trent Nelson --- rust/cuda-coop/src/thread_data.rs | 46 ++++++++++--------------------- rust/cuda-coop/tests/payload.rs | 30 ++++++++++++++++++++ 2 files changed, 44 insertions(+), 32 deletions(-) diff --git a/rust/cuda-coop/src/thread_data.rs b/rust/cuda-coop/src/thread_data.rs index 1fef6629f299..1992ec8939e9 100644 --- a/rust/cuda-coop/src/thread_data.rs +++ b/rust/cuda-coop/src/thread_data.rs @@ -71,6 +71,8 @@ where // without moving it away from offset zero. alignment: [ as crate::alignment::sealed::Alignment>::Marker; 0], values: [MaybeUninit; ITEMS], + // Packed low bits in the first ITEMS.div_ceil(8) bytes. The backing extent + // avoids generic const expressions and preserves this private layout size. initialized: [u8; ITEMS], } @@ -119,14 +121,17 @@ where #[inline] pub fn write(&mut self, index: usize, value: T) { self.values[index].write(value); - self.initialized[index] = 1; + self.initialized[index / 8] |= 1 << (index % 8); } /// Copy an initialized value, or return `None` if unwritten or out of bounds. - #[inline] + #[inline(always)] pub fn get(&self, index: usize) -> Option { - if index < ITEMS && self.initialized[index] == 1 { - // SAFETY: Only write sets this flag, after initializing the value. + if index >= ITEMS { + return None; + } + if (self.initialized[index / 8] & (1 << (index % 8))) != 0 { + // SAFETY: Only write sets this bit, after initializing the value. // Copies preserve both arrays. Numeric types are Copy. Some(unsafe { self.values[index].assume_init() }) } else { @@ -138,8 +143,8 @@ where /// out of bounds. The exclusive borrow prevents simultaneous payload access. #[inline] pub fn get_mut(&mut self, index: usize) -> Option<&mut T> { - if index < ITEMS && self.initialized[index] == 1 { - // SAFETY: The flag proves initialization, and &mut self is exclusive. + if index < ITEMS && (self.initialized[index / 8] & (1 << (index % 8))) != 0 { + // SAFETY: The bit proves initialization, and &mut self is exclusive. Some(unsafe { self.values[index].assume_init_mut() }) } else { None @@ -166,8 +171,8 @@ where #[inline] fn index(&self, index: usize) -> &T { assert!(index < ITEMS); - assert!(self.initialized[index] == 1); - // SAFETY: The bounds and exact initialization flag were checked above. + assert!((self.initialized[index / 8] & (1 << (index % 8))) != 0); + // SAFETY: The bounds and initialization bit were checked above. unsafe { self.values[index].assume_init_ref() } } } @@ -180,7 +185,7 @@ where #[inline] fn index_mut(&mut self, index: usize) -> &mut T { assert!(index < ITEMS); - assert!(self.initialized[index] == 1); + assert!((self.initialized[index / 8] & (1 << (index % 8))) != 0); // SAFETY: The checks prove initialization; &mut self is exclusive. unsafe { self.values[index].assume_init_mut() } } @@ -276,26 +281,3 @@ where Alignment: SupportedAlignment, { } - -#[cfg(test)] -mod tests { - extern crate std; - - use super::ThreadData; - use std::panic::{AssertUnwindSafe, catch_unwind}; - - #[test] - fn only_exactly_one_marks_a_slot_initialized() { - let mut items = ThreadData::<_, 1>::new(); - items.write(0, 42_i32); - for flag in [0, 2, 255] { - items.initialized[0] = flag; - assert_eq!(items.get(0), None); - assert!(items.get_mut(0).is_none()); - assert!(catch_unwind(|| items[0]).is_err()); - assert!(catch_unwind(AssertUnwindSafe(|| items[0] = 9)).is_err()); - } - items.write(0, 7); - assert_eq!(items.get(0), Some(7)); - } -} diff --git a/rust/cuda-coop/tests/payload.rs b/rust/cuda-coop/tests/payload.rs index 8b24d6a7c2c8..88e8c24b4a9e 100644 --- a/rust/cuda-coop/tests/payload.rs +++ b/rust/cuda-coop/tests/payload.rs @@ -61,6 +61,36 @@ fn arbitrary_writes_and_checked_mutation_preserve_unwritten_slots() { assert_eq!(items.get(2), Some(-1)); } +#[test] +fn sparse_writes_across_byte_boundaries_preserve_neighbors_and_copies() { + let mut items = coop::ThreadData::<_, 33>::new(); + for index in [0, 7, 8, 16, 31, 32] { + items.write(index, index as u16 + 1); + } + let copy = items; + *items.get_mut(8).unwrap() += 10; + items[31] += 20; + for index in 0..33 { + let expected = [0, 7, 8, 16, 31, 32] + .contains(&index) + .then_some(index as u16 + 1); + assert_eq!(copy.get(index), expected); + assert_eq!( + items.get(index), + expected.map(|value| value + + match index { + 8 => 10, + 31 => 20, + _ => 0, + }) + ); + } + for index in [33, usize::MAX] { + assert_eq!(items.get(index), None); + assert!(items.get_mut(index).is_none()); + } +} + #[test] fn invalid_indexing_and_writes_panic_without_initializing_other_slots() { let mut items = coop::ThreadData::<_, 4>::new(); From 8f2eeb4430482b8604235da954440873b482f2cc Mon Sep 17 00:00:00 2001 From: Trent Nelson Date: Tue, 20 Oct 2026 10:44:35 -0700 Subject: [PATCH 22/33] Cover enum controls in checked collective calls Signed-off-by: Trent Nelson --- .../tests/check_collective_rejections.py | 83 ++++++++++++++++++- 1 file changed, 82 insertions(+), 1 deletion(-) diff --git a/rust/cuda-coop-codegen/tests/check_collective_rejections.py b/rust/cuda-coop-codegen/tests/check_collective_rejections.py index 7689ec4308eb..bb64773b3081 100644 --- a/rust/cuda-coop-codegen/tests/check_collective_rejections.py +++ b/rust/cuda-coop-codegen/tests/check_collective_rejections.py @@ -332,6 +332,87 @@ def in_helper(code): ), None, ) +cases.update( + { + "enum_uniform": ( + { + "helper": in_helper( + " let choice = enum_choice(mode & 1 == 0, valid);\n" + " let moved = choice;\n" + " let valid = match moved { Some(count) => count, None => valid };" + )["helper"] + + """ +#[inline(never)] +fn enum_choice(present: bool, count: i32) -> Option { + if present { Some(count) } else { None } +} +""" + }, + None, + ), + "enum_uniform_tag": ( + in_helper( + " let choice = Some(lane);\n" + " match choice { Some(_) => cuda_device::thread::sync_threads(), None => {} }" + ), + None, + ), + "enum_presence": ( + in_helper( + " let choice = if lane & 1 == 0 { Some(valid) } else { None };\n" + " match choice { Some(_) => cuda_device::thread::sync_threads(), None => {} }" + ), + "lane-varying", + ), + "enum_fieldless": ( + in_helper( + " enum Choice { Present, Absent }\n" + " let choice = if lane & 1 == 0 { Choice::Present } else { Choice::Absent };\n" + " match choice { Choice::Present => cuda_device::thread::sync_threads(), Choice::Absent => {} }" + ), + "lane-varying", + ), + "enum_count": ( + in_helper( + " let choice = Some(lane as i32);\n" + " let valid = match choice { Some(count) => count, None => valid };" + ), + "not proved block-uniform", + ), + "enum_pointer": ( + in_helper( + " let choice = Some(unsafe { input.add(lane) });\n" + " let input = match choice { Some(pointer) => pointer, None => input };" + ), + "not proved block-uniform", + ), + "enum_default": ( + { + "api": original["api"].replace( + "(input, valid, -7_i32),", + "(input, valid, match Some(cuda_device::thread::threadIdx_x() as i32) { Some(value) => value, None => -7_i32 }),", + ) + }, + "not proved block-uniform", + ), + "enum_mutation": ( + { + "helper": in_helper( + " let mut choice = Some(valid);\n" + " replace_choice(&mut choice, lane as i32);\n" + " let valid = match choice { Some(count) => count, None => valid };" + )["helper"] + + """ +#[inline(never)] +fn replace_choice(choice: &mut Option, count: i32) { + *choice = Some(count); +} +""" + }, + "not proved block-uniform", + ), + } +) unknown = set(args.case) - cases.keys() if unknown: parser.error(f"unknown cases: {', '.join(sorted(unknown))}") @@ -363,7 +444,7 @@ def in_helper(code): case_env = env.copy() if name == "unchecked_bounds": case_env["CUDA_OXIDE_UNCHECKED_INDEXING"] = "1" - if name in [ + if name.startswith("enum_") or name in [ "aggregate_reference", "projected_reborrow", "reborrow_uniform", From 95ebc3df661e029bd0bf29f9b33b679006ba15e0 Mon Sep 17 00:00:00 2001 From: Trent Nelson Date: Tue, 20 Oct 2026 10:50:57 -0700 Subject: [PATCH 23/33] Expose Direct Load and Store with packed payload validity Signed-off-by: Trent Nelson --- docs/rust/index.rst | 15 +- docs/rust/load_store.rst | 125 ++ rust/Cargo.lock | 11 +- rust/cuda-coop-codegen/src/lib.rs | 33 +- rust/cuda-coop-codegen/src/oxide.rs | 70 +- .../tests/oxide-fixture/api/src/lib.rs | 6 - .../tests/oxide-fixture/api/src/validity.rs | 8 +- rust/cuda-coop/Cargo.toml | 5 +- rust/cuda-coop/build.rs | 8 + rust/cuda-coop/examples/load_store/.gitignore | 7 + rust/cuda-coop/examples/load_store/Cargo.lock | 1018 +++++++++++++++++ rust/cuda-coop/examples/load_store/Cargo.toml | 19 + .../examples/load_store/helper/Cargo.toml | 11 + .../examples/load_store/helper/src/lib.rs | 35 + .../cuda-coop/examples/load_store/src/main.rs | 69 ++ rust/cuda-coop/src/lib.rs | 6 +- rust/cuda-coop/src/load_store.rs | 310 +++++ rust/cuda-coop/src/load_store/boundary.rs | 148 +++ rust/cuda-coop/src/thread_data.rs | 3 + rust/cuda-coop/src/thread_data/block_io.rs | 11 +- rust/cuda-coop/tests/check_load_store.py | 147 +++ rust/cuda-coop/tests/load_store/.gitignore | 7 + rust/cuda-coop/tests/load_store/Cargo.lock | 1010 ++++++++++++++++ rust/cuda-coop/tests/load_store/Cargo.toml | 17 + rust/cuda-coop/tests/load_store/src/main.rs | 626 ++++++++++ 25 files changed, 3686 insertions(+), 39 deletions(-) create mode 100644 docs/rust/load_store.rst create mode 100644 rust/cuda-coop/build.rs create mode 100644 rust/cuda-coop/examples/load_store/.gitignore create mode 100644 rust/cuda-coop/examples/load_store/Cargo.lock create mode 100644 rust/cuda-coop/examples/load_store/Cargo.toml create mode 100644 rust/cuda-coop/examples/load_store/helper/Cargo.toml create mode 100644 rust/cuda-coop/examples/load_store/helper/src/lib.rs create mode 100644 rust/cuda-coop/examples/load_store/src/main.rs create mode 100644 rust/cuda-coop/src/load_store.rs create mode 100644 rust/cuda-coop/src/load_store/boundary.rs create mode 100644 rust/cuda-coop/tests/check_load_store.py create mode 100644 rust/cuda-coop/tests/load_store/.gitignore create mode 100644 rust/cuda-coop/tests/load_store/Cargo.lock create mode 100644 rust/cuda-coop/tests/load_store/Cargo.toml create mode 100644 rust/cuda-coop/tests/load_store/src/main.rs diff --git a/docs/rust/index.rst b/docs/rust/index.rst index e2950e6e8e8f..82d6bd1a140d 100644 --- a/docs/rust/index.rst +++ b/docs/rust/index.rst @@ -6,12 +6,12 @@ Rust development The Rust workspace contains per-thread payloads, CUDA group queries, and the native host compiler foundation for ``cuda.coop``. The ``cuda-coop`` crate supplies checked -``ThreadData`` values and physical/mapped hierarchy queries; see :doc:`payload` -and :doc:`groups`. Its private ``cuda-coop-codegen`` crate +``ThreadData`` values, physical/mapped hierarchy queries, and Direct block +Load/Store; see :doc:`payload`, :doc:`groups`, and :doc:`load_store`. 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; collective operations and scratch storage are -still under development. The crate is unpublished and its interface can change. +compiler integration development. Additional collective algorithms and public +scratch storage remain under development. The crate is unpublished and its interface can change. Device code uses the pinned Oxide dependency and ``nightly-2026-10-11``. The host provider supports Rust 1.85 independently; validate it with @@ -64,9 +64,9 @@ all payload bytes and flags. A failing Store traps; it is not transactional across lanes. Partial Store forms destination pointers only within its valid prefix, including when the allocation is shorter than a complete tile. -Validity-aware Transpose and the public collective API remain under development. -The private integration fixture uses the actual ``ThreadData`` implementation -without exposing an initialized slice or a public native layout. +Validity-aware Transpose remains under development. The public Direct API uses +the same ``ThreadData`` representation without exposing an initialized slice or +a public native layout. Compile a provider ------------------ @@ -144,3 +144,4 @@ timestamps may require an explicit rebuild. payload groups + load_store diff --git a/docs/rust/load_store.rst b/docs/rust/load_store.rst new file mode 100644 index 000000000000..8515808b6241 --- /dev/null +++ b/docs/rust/load_store.rst @@ -0,0 +1,125 @@ +.. Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. +.. SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +Direct Load and Store +===================== + +``cuda_coop::load`` fills an existing ``ThreadData`` and returns unit. +``cuda_coop::store`` writes its values and returns unit. Both take the group +first. This initial implementation supports complete, one-dimensional blocks +and the Direct algorithm: each thread owns consecutive elements in blocked +order. Other algorithms, warp groups, and public temporary storage remain +under development. + +Use ``use cuda_coop as coop;`` and keep the per-thread item count a const generic +named ``ITEMS_PER_THREAD``. The payload element type is inferred from its source. +The supported types are ``i8``, ``u8``, ``i16``, ``u16``, ``i32``, ``u32``, +``i64``, ``u64``, ``f32`` and ``f64``. + +A complete program +------------------ + +The example is an independent Cargo workspace with a separate generic helper +crate. It copies three block tiles at item counts 1 and 4, including a short last +tile. The consumer has no build script, native shim, provider, or scratch argument. +Its ``cuda-coop`` dependency supplies the native specialization during the normal +Oxide build. + +Use the Oxide revision selected by the manifests and ``nightly-2026-10-11``. +The initial native configuration is Linux with CUDA 13.3 and a supported NVIDIA +GPU. From a complete CCCL checkout, with that Oxide toolchain installed: + +.. code-block:: shell + + cd rust/cuda-coop/examples/load_store + CUDA_HOME=/usr/local/cuda-13.3 CUDA_PATH=/usr/local/cuda-13.3 \ + cargo oxide run --arch sm_120 --materialize-cubin + +Choose the architecture for your GPU. Materialization produces the native device +artifact during the Cargo build. No separate compilation or linking step is +required. The exact block launch contract is checked when preparing the launch; +``auto_shared`` uses the compiler's inferred requirement, which is zero for +Direct Load/Store. + +.. literalinclude:: ../../rust/cuda-coop/examples/load_store/src/main.rs + :language: rust + :start-after: // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +The reusable helper preserves the source element type and borrows its payload +read-only for Direct Store: + +.. literalinclude:: ../../rust/cuda-coop/examples/load_store/helper/src/lib.rs + :language: rust + :start-after: use cuda_coop as coop; + +Counts, offsets, and payload validity +------------------------------------- + +``LoadOptions::default()`` and ``StoreOptions::default()`` select the full tile +at element offset zero. Named setters select a partial count, a typed Load +default, and an offset. The count describes the whole block's prefix, not a +per-thread count. It must fit the tile and ``i32::MAX``. Offsets are element +indices from the start of the allocation; supply each block's origin explicitly. +Counts, defaults, offsets, and their presence must be uniform across the block. + +A masked Load without ``oob_default`` resets every out-of-prefix slot to +unwritten, even if initialized before the Load. ``get`` returns ``None`` for an +unwritten slot; indexing it traps. A default requires an explicit count and has +the source's numeric type. Filled Load initializes every slot, including the +missing tail. Partial Store writes only its selected prefix and requires only +those slots to be initialized. + +Direct Store accepts ``&ThreadData``, a shared reborrow of ``&mut ThreadData``, +or a scalar numeric value. A scalar contributes one item per thread. Store +preserves every payload value and validity flag, including initialized slots +outside its write domain. Algorithms that mutate payloads will require mutable +borrows when those algorithms are supported. + +Collective ownership and completion +----------------------------------- + +Writable kernel parameters use the canonical spelling +``DisjointSlice>``. The capacity defines each block's +ownership grain as ``block_size * CAPACITY`` elements; it is independent of the +operation's per-thread item count. The carrier keeps the buffer lifetime and +collective ownership without giving each lane an overlapping mutable slice. +The compiler's generated host launch consumes or borrows the actual device +buffers until the operation completes. + +The descriptor returned by ``this_block()`` is a query value. Safe collective +calls additionally require a supported Oxide build, an exact launch contract, +uniform control flow, and matching group participation. Compilation diagnoses +lane-varying controls or conflicting resource use. Runtime checks trap invalid +bounds, ownership ranges, counts, or unwritten values selected for Store. +Failures are not transactional across lanes. + +To load from a carrier after storing to it, call +``cuda_device::thread::sync_threads()`` and pass ``&mut carrier`` as the Load +source. The compiler requires this completion before conflicting accesses, +including reuse in loops. Ending a Rust borrow or dropping a value is not +completion. The host example calls ``sync()`` before inspecting its output. + +Host development +---------------- + +Host payload tests and rustdoc do not require a GPU or CUDA installation: + +.. code-block:: shell + + cd rust + cargo test --workspace + cargo clippy --workspace --all-targets -- -D warnings + RUSTDOCFLAGS="-D warnings" cargo doc --workspace --no-deps + +The codegen crate is only a host build dependency. Native headers come from the +same complete CCCL checkout as that dependency; a Python installation or the +consumer's working directory does not select headers. Native builds use +Oxide's selected toolkit. Cargo tracks the provider and declared native inputs; +replacing them with unchanged or older timestamps may require an explicit +rebuild. Crate publication is not supported yet. + +The rustdoc entries for ``load``, ``store``, ``LoadOptions``, ``StoreOptions``, +``LoadSource`` and ``StoreValue`` specify their full generic signatures and +supported argument forms. Native correctness and artifact checks are separate +from host compilation. Performance qualification for masked payload access is +ongoing; the current implementation makes no negligible-overhead claim. diff --git a/rust/Cargo.lock b/rust/Cargo.lock index e6a9256c8ed5..e8b7b6d4af62 100644 --- a/rust/Cargo.lock +++ b/rust/Cargo.lock @@ -40,6 +40,7 @@ dependencies = [ name = "cuda-coop" version = "0.0.0" dependencies = [ + "cuda-coop-codegen", "cuda-device", ] @@ -56,7 +57,7 @@ dependencies = [ [[package]] name = "cuda-device" version = "0.2.1" -source = "git+https://github.com/tpn/cuda-rust.git?rev=8c4a42e9c95b2a6bdc14a7531904d8469f0e7ca6#8c4a42e9c95b2a6bdc14a7531904d8469f0e7ca6" +source = "git+https://github.com/tpn/cuda-rust.git?rev=be42c33eeb15ea5639fdd7d793a76706f10cb19e#be42c33eeb15ea5639fdd7d793a76706f10cb19e" dependencies = [ "cuda-macros", ] @@ -64,12 +65,14 @@ dependencies = [ [[package]] name = "cuda-macros" version = "0.2.1" -source = "git+https://github.com/tpn/cuda-rust.git?rev=8c4a42e9c95b2a6bdc14a7531904d8469f0e7ca6#8c4a42e9c95b2a6bdc14a7531904d8469f0e7ca6" +source = "git+https://github.com/tpn/cuda-rust.git?rev=be42c33eeb15ea5639fdd7d793a76706f10cb19e#be42c33eeb15ea5639fdd7d793a76706f10cb19e" dependencies = [ "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", + "serde", + "serde_json", "syn 2.0.119", ] @@ -124,7 +127,7 @@ checksum = "cf8baf1c55e62ffcace7a9f06f4bd9cd3f0c4beb022d3b367256b91b87513d98" [[package]] name = "oxide-tile-export" version = "0.2.1" -source = "git+https://github.com/tpn/cuda-rust.git?rev=8c4a42e9c95b2a6bdc14a7531904d8469f0e7ca6#8c4a42e9c95b2a6bdc14a7531904d8469f0e7ca6" +source = "git+https://github.com/tpn/cuda-rust.git?rev=be42c33eeb15ea5639fdd7d793a76706f10cb19e#be42c33eeb15ea5639fdd7d793a76706f10cb19e" dependencies = [ "serde", "serde_json", @@ -152,7 +155,7 @@ dependencies = [ [[package]] name = "reserved-oxide-symbols" version = "0.2.1" -source = "git+https://github.com/tpn/cuda-rust.git?rev=8c4a42e9c95b2a6bdc14a7531904d8469f0e7ca6#8c4a42e9c95b2a6bdc14a7531904d8469f0e7ca6" +source = "git+https://github.com/tpn/cuda-rust.git?rev=be42c33eeb15ea5639fdd7d793a76706f10cb19e#be42c33eeb15ea5639fdd7d793a76706f10cb19e" [[package]] name = "serde" diff --git a/rust/cuda-coop-codegen/src/lib.rs b/rust/cuda-coop-codegen/src/lib.rs index 750b9fa4a00d..41b532965f6d 100644 --- a/rust/cuda-coop-codegen/src/lib.rs +++ b/rust/cuda-coop-codegen/src/lib.rs @@ -349,18 +349,18 @@ pub struct Layout { /// so a null pointer is permitted for that argument only. All providers return /// unit; loads fill payloads in place and stores write memory. /// -/// Validity-aware Direct ABIs add `unsigned char* flags` immediately after +/// Validity-aware Direct ABIs add a byte `flags` pointer immediately after /// scratch. `payload` addresses aligned storage for the requested number of /// elements; it need not contain initialized values. Flags occupy a separate, /// initialized array of `items_per_thread.div_ceil(8)` bytes. Bit `i % 8` of /// byte `i / 8` denotes whether slot `i` is initialized. Both arrays are private -/// to the lane, disjoint from memory/scratch and one another, and remain -/// exclusively borrowed for the call. Native code accesses payload bytes -/// without evaluating invalid values. Masked Load clears bits outside its -/// prefix; filled/full Load write every value before setting its bit. -/// Store checks its entire lane-local write domain before numeric conversion -/// and preserves both arrays. Failure traps; it does not promise rollback of -/// stores by other lanes. +/// to the lane, disjoint from memory/scratch and one another. Load exclusively +/// borrows both arrays; Direct Store receives const pointers and borrows them +/// read-only. Native code accesses payload bytes without evaluating invalid +/// values. Masked Load clears bits outside its prefix; filled/full Load write +/// every value before setting its bit. Store checks its entire lane-local +/// write domain before numeric conversion and preserves both arrays. Failure +/// traps; it does not promise rollback of stores by other lanes. #[derive(Debug)] pub struct Specialization { /// SHA-256 of request/source, target/options, captured headers and native toolchain inputs. @@ -588,6 +588,11 @@ fn source(request: Request, symbol: &str) -> String { } else { "memory" }; + let payload_const = if op.validity() && !op.is_load() { + "const " + } else { + "" + }; let (cell, flags, body) = if op.validity() { let domain = if matches!( op, @@ -633,8 +638,14 @@ fn source(request: Request, symbol: &str) -> String { }; ( "cuda_coop_detail::opaque_cell", - ", unsigned char* flags", - format!("auto* const bytes = reinterpret_cast(payload);\n {body}"), + if op.is_load() { + ", unsigned char* flags" + } else { + ", const unsigned char* flags" + }, + format!( + "auto* const bytes = reinterpret_cast<{payload_const}unsigned char*>(payload);\n {body}" + ), ) } else { ( @@ -655,7 +666,7 @@ using cell = {cell}; using primitive = cub::{primitive}; static_assert(sizeof(value_type) == {size} && alignof(value_type) == {size}); static_assert(sizeof(cell) == sizeof(value_type) && alignof(cell) == alignof(value_type)); -extern "C" __device__ void {symbol}({qualifier}value_type* memory, value_type* payload, unsigned char* scratch{flags}{controls}) {{ +extern "C" __device__ void {symbol}({qualifier}value_type* memory, {payload_const}value_type* payload, unsigned char* scratch{flags}{controls}) {{ {check} {constructor} {body} diff --git a/rust/cuda-coop-codegen/src/oxide.rs b/rust/cuda-coop-codegen/src/oxide.rs index 58d8590a91f3..1d2f10a04c2d 100644 --- a/rust/cuda-coop-codegen/src/oxide.rs +++ b/rust/cuda-coop-codegen/src/oxide.rs @@ -203,10 +203,18 @@ fn arguments(request: Request) -> Vec { } else { Argument::MutPointer(dtype) }, - Argument::MutPointer(dtype), + if operation.validity() && !operation.is_load() { + Argument::ConstPointer(dtype) + } else { + Argument::MutPointer(dtype) + }, ]; if operation.validity() { - arguments.push(Argument::MutPointer(Scalar::Unsigned(8))); + arguments.push(if operation.is_load() { + Argument::MutPointer(Scalar::Unsigned(8)) + } else { + Argument::ConstPointer(Scalar::Unsigned(8)) + }); } if operation.counted() { arguments.push(Argument::Value(Scalar::Signed(32))); @@ -217,6 +225,48 @@ fn arguments(request: Request) -> Vec { arguments } +/// Register the public Direct recipes from cuda-coop's owning build script. +/// +/// Normal Cargo builds only emit registrations for the current build-script +/// executable. Oxide later invokes that executable with the existing native +/// request/input protocol. Headers come from this codegen crate's complete CCCL +/// checkout, independent of the consumer's working directory. Host checks and +/// documentation do not load CUDA. Native compilation uses Oxide's selected +/// toolkit and target; no consumer build script is required. +/// +/// # Errors +/// +/// Returns errors locating this checkout or executable, encoding registrations, +/// or serving a native request. Native request errors are those of [`serve`] +/// and [`serve_inputs`]. +pub fn build_script() -> Result<()> { + let checkout = Path::new(env!("CARGO_MANIFEST_DIR")) + .join("../..") + .canonicalize()?; + match std::env::args().nth(1).as_deref() { + Some("--oxide-library-request-v1") => { + return serve(&checkout, std::io::stdin().lock(), std::io::stdout().lock()); + } + Some("--oxide-library-inputs-v1") => { + return serve_inputs(&checkout, std::io::stdin().lock(), std::io::stdout().lock()); + } + _ => {} + } + println!("cargo::rerun-if-changed=build.rs"); + let executable = std::env::current_exe()?; + for (operation, name) in [ + (Operation::LoadFullValidity, "LOAD_FULL"), + (Operation::LoadMaskedValidity, "LOAD_MASKED"), + (Operation::LoadFilledValidity, "LOAD_FILLED"), + (Operation::StoreFullValidity, "STORE_FULL"), + (Operation::StorePartialValidity, "STORE_PARTIAL"), + ] { + let value = registration(&executable, operation, Algorithm::Direct)?; + println!("cargo::rustc-env=COOP_VALIDITY_{name}={value}"); + } + Ok(()) +} + /// Encode a constant registration for one native recipe. /// /// `executable` is the live, already-built dependency build script, normally @@ -457,6 +507,22 @@ mod tests { ] ); } + for operation in [ + Operation::StoreFullValidity, + Operation::StorePartialValidity, + ] { + let request = + Request::new_typed(Dtype::I16, operation, Algorithm::Direct, 128, 4).unwrap(); + let mut expected = vec![ + Argument::MutPointer(Scalar::Signed(16)), + Argument::ConstPointer(Scalar::Signed(16)), + Argument::ConstPointer(Scalar::Unsigned(8)), + ]; + if operation.counted() { + expected.push(Argument::Value(Scalar::Signed(32))); + } + assert_eq!(arguments(request), expected); + } for scalar in [Scalar::Signed(1), Scalar::Unsigned(128), Scalar::Float(16)] { assert!(scalar.dtype().is_err()); } 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 index b4a21fbf5395..cb14c141806a 100644 --- a/rust/cuda-coop-codegen/tests/oxide-fixture/api/src/lib.rs +++ b/rust/cuda-coop-codegen/tests/oxide-fixture/api/src/lib.rs @@ -24,12 +24,6 @@ pub mod thread_data { env!("CCCL_SOURCE"), "/rust/cuda-coop/src/thread_data.rs" )); - mod block_io { - include!(concat!( - env!("CCCL_SOURCE"), - "/rust/cuda-coop/src/thread_data/block_io.rs" - )); - } // Fault injection and flag observation belong only to this private fixture. pub fn flag(values: &ThreadData, index: usize) -> u8 { diff --git a/rust/cuda-coop-codegen/tests/oxide-fixture/api/src/validity.rs b/rust/cuda-coop-codegen/tests/oxide-fixture/api/src/validity.rs index a0543a71b15d..3118fb4bbaf6 100644 --- a/rust/cuda-coop-codegen/tests/oxide-fixture/api/src/validity.rs +++ b/rust/cuda-coop-codegen/tests/oxide-fixture/api/src/validity.rs @@ -99,7 +99,7 @@ pub fn store_full, target: &mut DisjointSlice>, offset: usize, - values: &mut ThreadData, + values: &ThreadData, ) where Alignment: SupportedAlignment, { @@ -107,7 +107,7 @@ pub fn store_full(target.len(), offset, valid); crate::carrier::containment::(offset, valid); let pointer = target.as_mut_ptr(); - let payload = values.block_io_parts(); + let payload = values.block_io_read_parts(); // SAFETY: this carrier owns the checked range. The provider checks every // written slot's validity bit before interpreting its bytes as T. unsafe { @@ -126,14 +126,14 @@ pub fn store_partial>, offset: usize, valid: usize, - values: &mut ThreadData, + values: &ThreadData, ) where Alignment: SupportedAlignment, { extent::(target.len(), offset, valid); crate::carrier::containment::(offset, valid); let pointer = target.as_mut_ptr(); - let payload = values.block_io_parts(); + let payload = values.block_io_read_parts(); // SAFETY: the guarded native iterator forms an effective pointer only for // a written slot. Preflight covers that same domain, with no tail reads. unsafe { diff --git a/rust/cuda-coop/Cargo.toml b/rust/cuda-coop/Cargo.toml index 03e46ac0c204..0a903f2e1ed2 100644 --- a/rust/cuda-coop/Cargo.toml +++ b/rust/cuda-coop/Cargo.toml @@ -9,7 +9,10 @@ license.workspace = true publish = false [dependencies] -cuda-device = { git = "https://github.com/tpn/cuda-rust.git", rev = "8c4a42e9c95b2a6bdc14a7531904d8469f0e7ca6" } +cuda-device = { git = "https://github.com/tpn/cuda-rust.git", rev = "be42c33eeb15ea5639fdd7d793a76706f10cb19e" } + +[build-dependencies] +cuda-coop-codegen = { path = "../cuda-coop-codegen" } [lints] workspace = true diff --git a/rust/cuda-coop/build.rs b/rust/cuda-coop/build.rs new file mode 100644 index 000000000000..858fae4c0249 --- /dev/null +++ b/rust/cuda-coop/build.rs @@ -0,0 +1,8 @@ +// Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +//! Register the owning host provider for native Direct operations. + +fn main() { + cuda_coop_codegen::oxide::build_script().expect("cuda.coop native provider registration"); +} diff --git a/rust/cuda-coop/examples/load_store/.gitignore b/rust/cuda-coop/examples/load_store/.gitignore new file mode 100644 index 000000000000..74bffa8af4a4 --- /dev/null +++ b/rust/cuda-coop/examples/load_store/.gitignore @@ -0,0 +1,7 @@ +# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. +# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception +/target/ +/.oxide-artifacts/ +/*.ll +/*.options +/*.target diff --git a/rust/cuda-coop/examples/load_store/Cargo.lock b/rust/cuda-coop/examples/load_store/Cargo.lock new file mode 100644 index 000000000000..7e39af8b87c4 --- /dev/null +++ b/rust/cuda-coop/examples/load_store/Cargo.lock @@ -0,0 +1,1018 @@ +# This file is automatically @generated by Cargo. +# It is not intended for manual editing. +version = 4 + +[[package]] +name = "aho-corasick" +version = "1.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c982642fa9e8606056828ee9a8505737230110bb1099153c79efe865c59d12ba" +dependencies = [ + "memchr", +] + +[[package]] +name = "anyhow" +version = "1.0.104" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "330a5ed07fa54e4702c9d6c4174f74427fc0ef6e214bbd677ae50a5099946470" + +[[package]] +name = "bindgen" +version = "0.69.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "271383c67ccabffb7381723dea0672a673f292304fcb45c01cc648c7a8d58088" +dependencies = [ + "bitflags", + "cexpr", + "clang-sys", + "itertools", + "lazy_static", + "lazycell", + "log", + "prettyplease", + "proc-macro2", + "quote", + "regex", + "rustc-hash 1.1.0", + "shlex", + "syn 2.0.119", + "which", +] + +[[package]] +name = "bitflags" +version = "2.13.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3ded4057c258ba199e2d26386d3af3780957ecaee6c4ef4041c6b4b8b97c0b06" + +[[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 = "cexpr" +version = "0.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6fac387a98bb7c37292057cffc56d62ecb629900026402633ae9160df93a8766" +dependencies = [ + "nom", +] + +[[package]] +name = "cfg-if" +version = "1.0.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4e7648175b45a9a48536d676f68d918270699102aa8dab5496df06904c914600" + +[[package]] +name = "clang-sys" +version = "1.9.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "157a8ba7b480713b56f4c09fd13fc3e0a22a5dfab8097ba61cbc5feef950788a" +dependencies = [ + "glob", + "libc", + "libloading", +] + +[[package]] +name = "coop-load-store" +version = "0.0.0" +dependencies = [ + "coop-load-store-helper", + "cuda-async", + "cuda-coop", + "cuda-core", + "cuda-device", + "cuda-host", +] + +[[package]] +name = "coop-load-store-helper" +version = "0.0.0" +dependencies = [ + "cuda-coop", +] + +[[package]] +name = "cpufeatures" +version = "0.2.17" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "59ed5838eebb26a2bb2e58f6d5b5316989ae9d08bab10e0e6d103e656d1b0280" +dependencies = [ + "libc", +] + +[[package]] +name = "crc32fast" +version = "1.5.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0ba106e844045ddc02c20768ac4b0b7cb45f2cc2ecbd833ac9f5dacb183c2055" +dependencies = [ + "cfg-if", +] + +[[package]] +name = "crossbeam-utils" +version = "0.8.23" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a31eee39dddec8330830986fcd7625edb5a24ec90ea038215273bbc3adb08ac6" + +[[package]] +name = "crunchy" +version = "0.2.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "460fbee9c2c2f33933d720630a6a0bac33ba7053db5344fac858d4b8952d77d5" + +[[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-artifact-finalizer" +version = "0.2.1" +source = "git+https://github.com/tpn/cuda-rust.git?rev=be42c33eeb15ea5639fdd7d793a76706f10cb19e#be42c33eeb15ea5639fdd7d793a76706f10cb19e" +dependencies = [ + "libnvvm-sys", + "nvjitlink-sys", + "serde", + "sha2", + "thiserror 2.0.21", +] + +[[package]] +name = "cuda-async" +version = "0.4.0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=be42c33eeb15ea5639fdd7d793a76706f10cb19e#be42c33eeb15ea5639fdd7d793a76706f10cb19e" +dependencies = [ + "anyhow", + "cuda-bindings", + "cuda-core", + "dashmap", + "futures", + "half", + "once_cell", + "rustc-hash 2.1.3", + "thiserror 1.0.69", +] + +[[package]] +name = "cuda-bindings" +version = "0.4.0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=be42c33eeb15ea5639fdd7d793a76706f10cb19e#be42c33eeb15ea5639fdd7d793a76706f10cb19e" +dependencies = [ + "bindgen", + "libloading", + "prettyplease", + "proc-macro2", + "quote", + "syn 2.0.119", +] + +[[package]] +name = "cuda-coop" +version = "0.0.0" +dependencies = [ + "cuda-coop-codegen", + "cuda-device", +] + +[[package]] +name = "cuda-coop-codegen" +version = "0.0.0" +dependencies = [ + "libloading", + "serde", + "serde_json", + "sha2", +] + +[[package]] +name = "cuda-core" +version = "0.4.0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=be42c33eeb15ea5639fdd7d793a76706f10cb19e#be42c33eeb15ea5639fdd7d793a76706f10cb19e" +dependencies = [ + "anyhow", + "cuda-bindings", + "cuda-core-derive", + "half", + "oxide-artifacts", +] + +[[package]] +name = "cuda-core-derive" +version = "0.4.0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=be42c33eeb15ea5639fdd7d793a76706f10cb19e#be42c33eeb15ea5639fdd7d793a76706f10cb19e" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", +] + +[[package]] +name = "cuda-device" +version = "0.2.1" +source = "git+https://github.com/tpn/cuda-rust.git?rev=be42c33eeb15ea5639fdd7d793a76706f10cb19e#be42c33eeb15ea5639fdd7d793a76706f10cb19e" +dependencies = [ + "cuda-macros", +] + +[[package]] +name = "cuda-host" +version = "0.2.1" +source = "git+https://github.com/tpn/cuda-rust.git?rev=be42c33eeb15ea5639fdd7d793a76706f10cb19e#be42c33eeb15ea5639fdd7d793a76706f10cb19e" +dependencies = [ + "cuda-artifact-finalizer", + "cuda-async", + "cuda-core", + "cuda-macros", + "half", + "libc", + "oxide-artifacts", + "ptx-parse", + "sha2", + "thiserror 2.0.21", +] + +[[package]] +name = "cuda-macros" +version = "0.2.1" +source = "git+https://github.com/tpn/cuda-rust.git?rev=be42c33eeb15ea5639fdd7d793a76706f10cb19e#be42c33eeb15ea5639fdd7d793a76706f10cb19e" +dependencies = [ + "oxide-tile-export", + "proc-macro2", + "quote", + "reserved-oxide-symbols", + "serde", + "serde_json", + "syn 2.0.119", +] + +[[package]] +name = "cuda-target-spec" +version = "0.1.0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=be42c33eeb15ea5639fdd7d793a76706f10cb19e#be42c33eeb15ea5639fdd7d793a76706f10cb19e" + +[[package]] +name = "dashmap" +version = "6.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e6361d5c062261c78a176addb82d4c821ae42bed6089de0e12603cd25de2059c" +dependencies = [ + "cfg-if", + "crossbeam-utils", + "hashbrown 0.14.5", + "lock_api", + "once_cell", + "parking_lot_core", +] + +[[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 = "either" +version = "1.19.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0e9c71c2167ca323c882b99918929403426e2373ea17242ff5653e0d5e1058be" + +[[package]] +name = "equivalent" +version = "1.0.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "00d174d5400e5e8fd687ad1049e2f578285fa914201b1af7e8b112a4546bd826" + +[[package]] +name = "errno" +version = "0.3.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "39cab71617ae0d63f51a36d69f866391735b51691dbda63cf6f96d042b63efeb" +dependencies = [ + "libc", + "windows-sys 0.61.2", +] + +[[package]] +name = "foldhash" +version = "0.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d9c4f5dac5e15c24eb999c26181a6ca40b39fe946cbe4c263c7209467bc83af2" + +[[package]] +name = "futures" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9a31d2a3fbaaeb2af2368bbdd904aa8e812d3c04a1ee10d3171f52d556e5d0a3" +dependencies = [ + "futures-channel", + "futures-core", + "futures-executor", + "futures-io", + "futures-sink", + "futures-task", + "futures-util", +] + +[[package]] +name = "futures-channel" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b1f9e3d69d39e4862ffed03ed071a76f9a13ba1d9109d355b0f0aa6b15e393c4" +dependencies = [ + "futures-core", + "futures-sink", +] + +[[package]] +name = "futures-core" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "92d699e522242e69e3003b94ecc1f960f3a5e015aa7c5d7486e65ad01dd94f5e" + +[[package]] +name = "futures-executor" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "031b47cf1a3c6cc8bc2fc76cd437f521619387907d469316e7c0bc278f1f5432" +dependencies = [ + "futures-core", + "futures-task", + "futures-util", +] + +[[package]] +name = "futures-io" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "53c0fa8157de1303bfffdaa1cc2a673bfffb60102f76b0ef4441659124373fed" + +[[package]] +name = "futures-macro" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9fb9654ba8355388abeb8dcb4fc62f511300867002afc858860463bdd9fe0c44" +dependencies = [ + "proc-macro2", + "quote", + "syn 3.0.9", +] + +[[package]] +name = "futures-sink" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1944426bf7d03f1d14f708785e4b33efd750b36d48a157b836b3efc15ede8e1d" + +[[package]] +name = "futures-task" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cd417de3d1d015fc3bfd2b1ea46dfc7bab72ef86f1cc7cc9c78e728b34a6d1fd" + +[[package]] +name = "futures-util" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0d50a92467f8ba5dd6e3ee5d4bd04d73ab2e4e1c44474a0674821dfce14b79bc" +dependencies = [ + "futures-channel", + "futures-core", + "futures-io", + "futures-macro", + "futures-sink", + "futures-task", + "memchr", + "pin-project-lite", + "slab", +] + +[[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 = "glob" +version = "0.3.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e4eba85ea1d0a966a983acd07deee566e67395d2d96b6fb39e62b5a833f1eb0b" + +[[package]] +name = "half" +version = "2.7.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" +dependencies = [ + "cfg-if", + "crunchy", + "zerocopy", +] + +[[package]] +name = "hashbrown" +version = "0.14.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e5274423e17b7c9fc20b6e7e208532f9b19825d82dfd615708b70edd83df41f1" + +[[package]] +name = "hashbrown" +version = "0.15.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9229cfe53dfd69f0609a49f65461bd93001ea1ef889cd5529dd176593f5338a1" +dependencies = [ + "foldhash", +] + +[[package]] +name = "hashbrown" +version = "0.17.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ed5909b6e89a2db4456e54cd5f673791d7eca6732202bbf2a9cc504fe2f9b84a" + +[[package]] +name = "home" +version = "0.5.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cc627f471c528ff0c4a49e1d5e60450c8f6461dd6d10ba9dcd3a61d3dff7728d" +dependencies = [ + "windows-sys 0.61.2", +] + +[[package]] +name = "indexmap" +version = "2.14.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cc4e190f5d26ca7051642629da2c52fc03bde85a03197c99408dcd291734c855" +dependencies = [ + "equivalent", + "hashbrown 0.17.1", +] + +[[package]] +name = "itertools" +version = "0.12.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ba291022dbbd398a455acf126c1e341954079855bc60dfdda641363bd6922569" +dependencies = [ + "either", +] + +[[package]] +name = "itoa" +version = "1.0.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8f42a60cbdf9a97f5d2305f08a87dc4e09308d1276d28c869c684d7777685682" + +[[package]] +name = "lazy_static" +version = "1.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "20870f649af7073d53e38067b2a84312175d56ea15217e1b15bc83506ec50afb" + +[[package]] +name = "lazycell" +version = "1.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "830d08ce1d1d941e6b30645f1a0eb5643013d835ce3779a5fc208261dbe10f55" + +[[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 = "libnvvm-sys" +version = "0.2.1" +source = "git+https://github.com/tpn/cuda-rust.git?rev=be42c33eeb15ea5639fdd7d793a76706f10cb19e#be42c33eeb15ea5639fdd7d793a76706f10cb19e" +dependencies = [ + "cuda-target-spec", + "libloading", + "thiserror 2.0.21", +] + +[[package]] +name = "linux-raw-sys" +version = "0.4.15" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d26c52dbd32dccf2d10cac7725f8eae5296885fb5703b261f7d0a0739ec807ab" + +[[package]] +name = "lock_api" +version = "0.4.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "224399e74b87b5f3557511d98dff8b14089b3dadafcab6bb93eab67d3aace965" +dependencies = [ + "scopeguard", +] + +[[package]] +name = "log" +version = "0.4.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9f8bd3e56ce4dfc153cf470fffbfa98c7620958b312ca5c3a4b8d5181fd13c6" + +[[package]] +name = "memchr" +version = "2.8.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cf8baf1c55e62ffcace7a9f06f4bd9cd3f0c4beb022d3b367256b91b87513d98" + +[[package]] +name = "minimal-lexical" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "68354c5c6bd36d73ff3feceb05efa59b6acb7626617f4962be322a825e61f79a" + +[[package]] +name = "nom" +version = "7.1.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d273983c5a657a70a3e8f2a01329822f3b8c8172b73826411a55751e404a0a4a" +dependencies = [ + "memchr", + "minimal-lexical", +] + +[[package]] +name = "nvjitlink-sys" +version = "0.2.1" +source = "git+https://github.com/tpn/cuda-rust.git?rev=be42c33eeb15ea5639fdd7d793a76706f10cb19e#be42c33eeb15ea5639fdd7d793a76706f10cb19e" +dependencies = [ + "libloading", + "thiserror 2.0.21", +] + +[[package]] +name = "object" +version = "0.36.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "62948e14d923ea95ea2c7c86c71013138b66525b86bdc08d2dcc262bdb497b87" +dependencies = [ + "crc32fast", + "hashbrown 0.15.5", + "indexmap", + "memchr", +] + +[[package]] +name = "once_cell" +version = "1.21.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9f7c3e4beb33f85d45ae3e3a1792185706c8e16d043238c593331cc7cd313b50" + +[[package]] +name = "oxide-artifacts" +version = "0.2.1" +source = "git+https://github.com/tpn/cuda-rust.git?rev=be42c33eeb15ea5639fdd7d793a76706f10cb19e#be42c33eeb15ea5639fdd7d793a76706f10cb19e" +dependencies = [ + "object", +] + +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +source = "git+https://github.com/tpn/cuda-rust.git?rev=be42c33eeb15ea5639fdd7d793a76706f10cb19e#be42c33eeb15ea5639fdd7d793a76706f10cb19e" +dependencies = [ + "serde", + "serde_json", + "thiserror 2.0.21", +] + +[[package]] +name = "parking_lot_core" +version = "0.9.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2621685985a2ebf1c516881c026032ac7deafcda1a2c9b7850dc81e3dfcb64c1" +dependencies = [ + "cfg-if", + "libc", + "redox_syscall", + "smallvec", + "windows-link", +] + +[[package]] +name = "pin-project-lite" +version = "0.2.17" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a89322df9ebe1c1578d689c92318e070967d1042b512afbe49518723f4e6d5cd" + +[[package]] +name = "prettyplease" +version = "0.2.37" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "479ca8adacdd7ce8f1fb39ce9ecccbfe93a3f1344b3d0d97f20bc0196208f62b" +dependencies = [ + "proc-macro2", + "syn 2.0.119", +] + +[[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 = "ptx-parse" +version = "0.1.0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=be42c33eeb15ea5639fdd7d793a76706f10cb19e#be42c33eeb15ea5639fdd7d793a76706f10cb19e" + +[[package]] +name = "quote" +version = "1.0.47" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1fbf4db142a473a8d80c26bbf18454ed458bf8d26c8219c331daecfdbd079001" +dependencies = [ + "proc-macro2", +] + +[[package]] +name = "redox_syscall" +version = "0.5.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ed2bf2547551a7053d6fdfafda3f938979645c44812fbfcda098faae3f1a362d" +dependencies = [ + "bitflags", +] + +[[package]] +name = "regex" +version = "1.13.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f020237b6c8eed93db2e2cb53c00c60a8e1bc73da7d073199a1180401450218d" +dependencies = [ + "aho-corasick", + "memchr", + "regex-automata", + "regex-syntax", +] + +[[package]] +name = "regex-automata" +version = "0.4.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ad8553b9b26413251cbf30e620595c7a41b3887f03da04579c0e6b0d6a06b4b2" +dependencies = [ + "aho-corasick", + "memchr", + "regex-syntax", +] + +[[package]] +name = "regex-syntax" +version = "0.8.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d6f6ff9a378485b298a5286656da665ba74413d36db0979633275d2e708145d4" + +[[package]] +name = "reserved-oxide-symbols" +version = "0.2.1" +source = "git+https://github.com/tpn/cuda-rust.git?rev=be42c33eeb15ea5639fdd7d793a76706f10cb19e#be42c33eeb15ea5639fdd7d793a76706f10cb19e" + +[[package]] +name = "rustc-hash" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "08d43f7aa6b08d49f382cde6a7982047c3426db949b1424bc4b7ec9ae12c6ce2" + +[[package]] +name = "rustc-hash" +version = "2.1.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6b1e7f9a428571be2dc5bc0505c13fb6bf936822b894ec87abf8a08a4e51742d" + +[[package]] +name = "rustix" +version = "0.38.44" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fdb5bc1ae2baa591800df16c9ca78619bf65c0488b41b96ccec5d11220d8c154" +dependencies = [ + "bitflags", + "errno", + "libc", + "linux-raw-sys", + "windows-sys 0.59.0", +] + +[[package]] +name = "scopeguard" +version = "1.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "94143f37725109f92c262ed2cf5e59bce7498c01bcc1502d7b9afe439a4e9f49" + +[[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 3.0.9", +] + +[[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 = "shlex" +version = "1.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0fda2ff0d084019ba4d7c6f371c95d8fd75ce3524c3cb8fb653a3023f6323e64" + +[[package]] +name = "slab" +version = "0.4.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0c790de23124f9ab44544d7ac05d60440adc586479ce501c1d6d7da3cd8c9cf5" + +[[package]] +name = "smallvec" +version = "1.16.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5b3dc8af474f516a851ff4bd12db780f948b9250ad37211e4eec0bccea54e01b" + +[[package]] +name = "syn" +version = "2.0.119" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "872831b642d1a07999a962a351ed35b955ea2cfc8f3862091e2a240a84f17297" +dependencies = [ + "proc-macro2", + "quote", + "unicode-ident", +] + +[[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 = "thiserror" +version = "1.0.69" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b6aaf5339b578ea85b50e080feb250a3e8ae8cfcdff9a461c9ec2904bc923f52" +dependencies = [ + "thiserror-impl 1.0.69", +] + +[[package]] +name = "thiserror" +version = "2.0.21" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "09e52cb86a36cede5cb101bf8908837b3e4c6e5e59fe7fd85c23fb56200d189e" +dependencies = [ + "thiserror-impl 2.0.21", +] + +[[package]] +name = "thiserror-impl" +version = "1.0.69" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4fee6c4efc90059e10f81e6d42c60a18f76588c3d74cb83a0b242a2b6c7504c1" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", +] + +[[package]] +name = "thiserror-impl" +version = "2.0.21" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fe5197923287db20a58125f0bc85c062f7f2c892de97b18c356f9efb14b28524" +dependencies = [ + "proc-macro2", + "quote", + "syn 3.0.9", +] + +[[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 = "which" +version = "4.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "87ba24419a2078cd2b0f2ede2691b6c66d8e47836da3b6db8265ebad47afbfc7" +dependencies = [ + "either", + "home", + "once_cell", + "rustix", +] + +[[package]] +name = "windows-link" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f0805222e57f7521d6a62e36fa9163bc891acd422f971defe97d64e70d0a4fe5" + +[[package]] +name = "windows-sys" +version = "0.59.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1e38bc4d79ed67fd075bcc251a1c39b32a1776bbe92e5bef1f0bf1f8c531853b" +dependencies = [ + "windows-targets", +] + +[[package]] +name = "windows-sys" +version = "0.61.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ae137229bcbd6cdf0f7b80a31df61766145077ddf49416a728b02cb3921ff3fc" +dependencies = [ + "windows-link", +] + +[[package]] +name = "windows-targets" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9b724f72796e036ab90c1021d4780d4d3d648aca59e491e6b98e725b84e99973" +dependencies = [ + "windows_aarch64_gnullvm", + "windows_aarch64_msvc", + "windows_i686_gnu", + "windows_i686_gnullvm", + "windows_i686_msvc", + "windows_x86_64_gnu", + "windows_x86_64_gnullvm", + "windows_x86_64_msvc", +] + +[[package]] +name = "windows_aarch64_gnullvm" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "32a4622180e7a0ec044bb555404c800bc9fd9ec262ec147edd5989ccd0c02cd3" + +[[package]] +name = "windows_aarch64_msvc" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "09ec2a7bb152e2252b53fa7803150007879548bc709c039df7627cabbd05d469" + +[[package]] +name = "windows_i686_gnu" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8e9b5ad5ab802e97eb8e295ac6720e509ee4c243f69d781394014ebfe8bbfa0b" + +[[package]] +name = "windows_i686_gnullvm" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0eee52d38c090b3caa76c563b86c3a4bd71ef1a819287c19d586d7334ae8ed66" + +[[package]] +name = "windows_i686_msvc" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "240948bc05c5e7c6dabba28bf89d89ffce3e303022809e73deaefe4f6ec56c66" + +[[package]] +name = "windows_x86_64_gnu" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "147a5c80aabfbf0c7d901cb5895d1de30ef2907eb21fbbab29ca94c5b08b1a78" + +[[package]] +name = "windows_x86_64_gnullvm" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "24d5b23dc417412679681396f2b49f3de8c1473deb516bd34410872eff51ed0d" + +[[package]] +name = "windows_x86_64_msvc" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "589f6da84c646204747d1270a2a5661ea66ed1cced2631d546fdfb155959f9ec" + +[[package]] +name = "zerocopy" +version = "0.8.63" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e5fe1f8f1b06191a00962174c61aa5005e0bb391a6d80d07e24d115c01a92ed8" +dependencies = [ + "zerocopy-derive", +] + +[[package]] +name = "zerocopy-derive" +version = "0.8.63" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "863ad3ac83293fb4d740aedbfdc9240dd8d1a50c1099acd76ce80ce7c7230c7f" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", +] + +[[package]] +name = "zmij" +version = "1.0.23" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "29666d0abbfad1e3dc4dcf6144730dd3a3ab225bbbdac83319345b1b44ccfc1b" diff --git a/rust/cuda-coop/examples/load_store/Cargo.toml b/rust/cuda-coop/examples/load_store/Cargo.toml new file mode 100644 index 000000000000..0cf131efc562 --- /dev/null +++ b/rust/cuda-coop/examples/load_store/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 = "coop-load-store" +version = "0.0.0" +edition = "2024" +publish = false + +[workspace] +members = ["helper"] + +[dependencies] +cuda-coop = { path = "../.." } +coop-load-store-helper = { path = "helper" } +cuda-device = { git = "https://github.com/tpn/cuda-rust.git", rev = "be42c33eeb15ea5639fdd7d793a76706f10cb19e" } +cuda-host = { git = "https://github.com/tpn/cuda-rust.git", rev = "be42c33eeb15ea5639fdd7d793a76706f10cb19e", features = ["async"] } +cuda-core = { git = "https://github.com/tpn/cuda-rust.git", rev = "be42c33eeb15ea5639fdd7d793a76706f10cb19e" } +cuda-async = { git = "https://github.com/tpn/cuda-rust.git", rev = "be42c33eeb15ea5639fdd7d793a76706f10cb19e" } diff --git a/rust/cuda-coop/examples/load_store/helper/Cargo.toml b/rust/cuda-coop/examples/load_store/helper/Cargo.toml new file mode 100644 index 000000000000..1c6b59094851 --- /dev/null +++ b/rust/cuda-coop/examples/load_store/helper/Cargo.toml @@ -0,0 +1,11 @@ +# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. +# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +[package] +name = "coop-load-store-helper" +version = "0.0.0" +edition = "2024" +publish = false + +[dependencies] +cuda-coop = { path = "../../.." } diff --git a/rust/cuda-coop/examples/load_store/helper/src/lib.rs b/rust/cuda-coop/examples/load_store/helper/src/lib.rs new file mode 100644 index 000000000000..986ac5ec32bc --- /dev/null +++ b/rust/cuda-coop/examples/load_store/helper/src/lib.rs @@ -0,0 +1,35 @@ +// Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +//! A reusable generic helper shared by downstream CUDA kernels. +#![no_std] +use cuda_coop as coop; + +/// Copy a partial block tile, preserving the source's element width. +pub fn copy( + block: coop::ThreadGroup, + input: &[T], + output: &mut coop::DisjointSlice<'_, T, coop::CollectiveTiles>, + offset: usize, + valid: usize, + default: T, +) { + let mut values = coop::ThreadData::<_, ITEMS_PER_THREAD>::new(); + coop::load( + block, + input, + &mut values, + coop::LoadOptions::default() + .valid_items(valid) + .oob_default(default) + .offset(offset), + ); + coop::store( + block, + output, + &values, + coop::StoreOptions::default() + .valid_items(valid) + .offset(offset), + ); +} diff --git a/rust/cuda-coop/examples/load_store/src/main.rs b/rust/cuda-coop/examples/load_store/src/main.rs new file mode 100644 index 000000000000..3f534fcc3ccb --- /dev/null +++ b/rust/cuda-coop/examples/load_store/src/main.rs @@ -0,0 +1,69 @@ +// Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +use cuda_coop as coop; +use cuda_device::{kernel, launch_contract}; +use cuda_host::cuda_module; + +#[cuda_module] +mod kernels { + use super::*; + + #[kernel] + #[launch_contract(domain = 1, block = (128, 1, 1), dynamic_shared = inferred)] + pub fn copy( + input: &[i16], + mut output: coop::DisjointSlice>, + ) { + let block = coop::this_block(); + let tile = block.count() as usize * ITEMS_PER_THREAD; + let offset = block.rank_at::() as usize * tile; + let remaining = if offset <= input.len() { + input.len() - offset + } else { + 0 + }; + let valid = remaining.min(tile); + coop_load_store_helper::copy::<_, ITEMS_PER_THREAD>( + block, + input, + &mut output, + offset, + valid, + -7_i16, + ); + } +} + +fn main() -> Result<(), Box> { + use cuda_async::simt::device_context::{init_device_contexts, with_cuda_context}; + use cuda_async::simt::device_operation::DeviceOperation; + use cuda_core::{DeviceBuffer, LaunchConfig1D}; + + init_device_contexts(0, 1)?; + let context = with_cuda_context(0, |context| context.clone())?; + let stream = context.default_stream(); + // SAFETY: these typed declarations and embedded kernels compile together. + let module = unsafe { kernels::load_async(0)? }; + macro_rules! run { + ($items:literal) => {{ + let input: Vec = (0..3 * 128 * $items - 7) + .map(|i| (i % 251) as i16) + .collect(); + let source = DeviceBuffer::from_host(&stream, &input)?; + let target = DeviceBuffer::from_host(&stream, &vec![-1_i16; input.len()])?; + let plan = module.prepare_copy::<$items>(LaunchConfig1D::auto_shared(3, 128))?; + let (_, target) = module + .copy_async_owned::<$items, _, _>(&plan, source, target) + .sync()?; + let mut output = vec![0_i16; input.len()]; + target.copy_to_host(&stream, &mut output)?; + stream.synchronize()?; + assert_eq!(output, input); + println!("items/thread={}: {} values copied", $items, output.len()); + }}; + } + run!(1); + run!(4); + Ok(()) +} diff --git a/rust/cuda-coop/src/lib.rs b/rust/cuda-coop/src/lib.rs index 3d575be24c8a..a73e14c85920 100644 --- a/rust/cuda-coop/src/lib.rs +++ b/rust/cuda-coop/src/lib.rs @@ -30,16 +30,20 @@ //! [`this_block`] and the other dimension-free group constructors describe //! the executing thread's current launch. Queries grant no collective //! participation or synchronization guarantee. This crate supplies no -//! collective operations or scratch storage. Host payload tests need no GPU; +//! scratch storage. Direct [`load`] and [`store`] use complete blocks with +//! compiler-checked participation and collective memory ownership. Host payload tests need no GPU; //! executing hardware queries requires an Oxide-compiled CUDA kernel. #![no_std] mod alignment; mod group; +mod load_store; mod thread_data; pub use alignment::{Alignment, SupportedAlignment}; +pub use cuda_device::{CollectiveTiles, DisjointSlice}; +pub use load_store::{LoadOptions, LoadSource, StoreOptions, StoreValue, load, store}; pub use thread_data::{Numeric, ThreadData, ThreadDataLike}; pub use group::{ diff --git a/rust/cuda-coop/src/load_store.rs b/rust/cuda-coop/src/load_store.rs new file mode 100644 index 000000000000..b73135d3c4ff --- /dev/null +++ b/rust/cuda-coop/src/load_store.rs @@ -0,0 +1,310 @@ +// Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +use crate::{Alignment, Block, Numeric, SupportedAlignment, ThreadData, ThreadGroup}; +use cuda_device::{CollectiveTiles, DisjointSlice}; + +mod boundary; + +/// Controls for a Direct [`load`]. +/// +/// `T` is inferred from the source and output. Defaults load the entire group +/// tile, leave no unspecified slots, and use element offset zero. A partial +/// load without an out-of-bounds default makes every slot outside its prefix +/// unwritten, including slots initialized before the call. +/// +/// Controls, including whether a count/default is present, must be uniform +/// across the block. They are checked when the operation executes. +#[derive(Clone, Copy)] +pub struct LoadOptions { + valid_items: Option, + oob_default: Option, + offset: usize, +} + +impl Default for LoadOptions { + fn default() -> Self { + Self { + valid_items: None, + oob_default: None, + offset: 0, + } + } +} + +impl LoadOptions { + /// Load this many elements of the whole group tile, starting at its origin. + /// + /// The count must not exceed the tile size or `i32::MAX`. Zero is permitted. + pub fn valid_items(mut self, count: usize) -> Self { + self.valid_items = Some(count); + self + } + + /// Initialize slots outside the valid prefix with this typed value. + /// + /// Requires an explicit [`valid_items`](Self::valid_items) count. The value + /// must have the source's element type and be uniform across the block. + pub fn oob_default(mut self, value: T) -> Self { + self.oob_default = Some(value); + self + } + + /// Set an element offset from the start of the source, defaulting to zero. + /// + /// Include the block's origin explicitly. The accessed prefix must fit the + /// source and, for a collective carrier, this block's owned region. + pub fn offset(mut self, offset: usize) -> Self { + self.offset = offset; + self + } +} + +/// Controls for a Direct [`store`]. +/// +/// Defaults write the entire group tile at element offset zero. A partial +/// store leaves destination elements outside its prefix unchanged. Controls, +/// including whether a count is present, must be uniform across the block. +#[derive(Clone, Copy, Default)] +pub struct StoreOptions { + valid_items: Option, + offset: usize, +} + +impl StoreOptions { + /// Write this many elements of the whole group tile. + /// + /// The count must not exceed the tile size or `i32::MAX`. A scalar payload + /// has one item per thread. Zero writes no elements. + pub fn valid_items(mut self, count: usize) -> Self { + self.valid_items = Some(count); + self + } + + /// Set an element offset from the start of the destination, default zero. + /// + /// Include the block's origin explicitly. The written prefix must fit the + /// allocation and this block's collective ownership region. + pub fn offset(mut self, offset: usize) -> Self { + self.offset = offset; + self + } +} + +/// Load a block tile in blocked order into an existing per-thread payload. +/// +/// `T` is inferred from `source`; `ITEMS` is the positive per-thread extent and +/// `ALIGN` is the payload's minimum alignment. The tile has +/// `group.count() * ITEMS` elements. Each thread receives consecutive items. +/// [`LoadSource`] accepts an immutable slice or a mutable borrow of a collective +/// carrier. [`LoadOptions`] supplies the count, typed default and element offset. +/// The operation returns unit and uses no shared scratch. +/// +/// Only complete one-dimensional blocks and the Direct algorithm are supported. +/// Compile with the supported Oxide backend and an exact block launch contract. +/// Every block thread must reach matching calls with uniform source, controls +/// and participation. The compiler checks these requirements and resource +/// conflicts; a group value or Rust borrow alone does not prove them. +/// +/// Loading a carrier after Store requires explicit block completion, such as +/// `cuda_device::thread::sync_threads()`. Ending a borrow is not completion. +/// Full/filled Load initializes every output slot. Masked Load without a default +/// resets every out-of-prefix slot to unwritten; access through [`ThreadData`] +/// remains checked. No initialized slice is exposed. +/// +/// # Panics +/// +/// Invalid counts, offsets, bounds, ownership regions, and a default without a +/// count trap on device. Group/control violations are compilation errors. +/// +/// ```no_run +/// use cuda_coop as coop; +/// fn copy_tile(input: &[i16], count: usize) { +/// let mut items = coop::ThreadData::<_, ITEMS_PER_THREAD>::new(); +/// coop::load(coop::this_block(), input, &mut items, +/// coop::LoadOptions::default().valid_items(count).oob_default(-7_i16)); +/// } +/// ``` +#[inline(always)] +pub fn load, const ITEMS: usize, const ALIGN: usize>( + group: ThreadGroup, + source: S, + output: &mut ThreadData, + options: LoadOptions, +) where + Alignment: SupportedAlignment, +{ + source.load(group, output, options); +} + +/// Store a block tile in blocked order, preserving the per-thread payload. +/// +/// `T` is inferred from the destination and value. [`StoreValue`] accepts a +/// shared or mutable borrow of [`ThreadData`], or a numeric scalar (one item per +/// thread). `CAPACITY` describes the carrier's ownership grain; it need not equal +/// the payload extent. This Direct operation borrows payload values and flags +/// read-only, returns unit, and uses no shared scratch. +/// +/// The supported groups, compiler/launch and uniformity requirements are those +/// of [`load`]. The destination is the canonical collective carrier: it grants +/// disjoint block regions without exposing duplicate mutable lane references. +/// Every written payload slot must be initialized. Slots outside the written +/// prefix may be unwritten and are preserved. A following conflicting access +/// to this carrier needs explicit block completion; Drop is not completion. +/// +/// # Panics +/// +/// Invalid counts, offsets, bounds, ownership regions or an unwritten value in +/// the write domain trap on device. Failure does not roll back other lanes' +/// stores. Group/control/resource violations are compilation errors. +/// +/// ```no_run +/// use cuda_coop as coop; +/// fn write_scalar(output: &mut coop::DisjointSlice<'_, i16, coop::CollectiveTiles<1>>) { +/// coop::store(coop::this_block(), output, 7_i16, Default::default()); +/// } +/// ``` +#[inline(always)] +pub fn store, const CAPACITY: usize>( + group: ThreadGroup, + destination: &mut DisjointSlice<'_, T, CollectiveTiles>, + value: V, + options: StoreOptions, +) { + value.store(group, destination, options); +} + +/// A sealed source accepted by [`load`] for numeric element type `T`. +/// +/// Implemented for `&[T]` and `&mut DisjointSlice>`. +/// The latter supports collective readback after explicit completion of a prior +/// Store. Implementations retain the original allocation witness and bounds; +/// this trait does not expose pointers or grant lane references. +pub trait LoadSource: sealed::LoadSource { + #[doc(hidden)] + fn load( + self, + group: ThreadGroup, + output: &mut ThreadData, + options: LoadOptions, + ) where + Alignment: SupportedAlignment; +} + +impl sealed::LoadSource for &[T] {} +impl LoadSource for &[T] { + #[inline(always)] + fn load( + self, + _group: ThreadGroup, + output: &mut ThreadData, + options: LoadOptions, + ) where + Alignment: SupportedAlignment, + { + // SAFETY: the pointer, length and immutable witness describe this slice. + unsafe { + boundary::load(self, self.as_ptr(), self.len(), output, options); + } + } +} + +impl sealed::LoadSource + for &mut DisjointSlice<'_, T, CollectiveTiles> +{ +} +impl LoadSource for &mut DisjointSlice<'_, T, CollectiveTiles> { + #[inline(always)] + fn load( + self, + _group: ThreadGroup, + output: &mut ThreadData, + options: LoadOptions, + ) where + Alignment: SupportedAlignment, + { + let valid = match options.valid_items { + Some(count) => count, + None => boundary::tile::(), + }; + boundary::containment::(options.offset, valid); + let pointer = self.as_mut_ptr().cast_const(); + let len = self.len(); + // SAFETY: the pointer and witness share an allocation; containment was + // checked. The compiler requires completion after conflicting accesses. + unsafe { + boundary::load(&*self, pointer, len, output, options); + } + } +} + +/// A sealed per-thread value accepted by Direct [`store`] for numeric type `T`. +/// +/// Implemented for numeric scalars and shared/mutable borrows of [`ThreadData`]. +/// A scalar contributes one item per thread. Payload borrows preserve every +/// value and flag, including slots outside the selected write domain. +pub trait StoreValue: sealed::StoreValue { + #[doc(hidden)] + fn store( + self, + group: ThreadGroup, + destination: &mut DisjointSlice<'_, T, CollectiveTiles>, + options: StoreOptions, + ); +} + +impl sealed::StoreValue for &ThreadData where + Alignment: SupportedAlignment +{ +} +impl StoreValue for &ThreadData +where + Alignment: SupportedAlignment, +{ + #[inline(always)] + fn store( + self, + _group: ThreadGroup, + destination: &mut DisjointSlice<'_, T, CollectiveTiles>, + options: StoreOptions, + ) { + boundary::store(destination, self, options); + } +} +impl sealed::StoreValue for &mut ThreadData where + Alignment: SupportedAlignment +{ +} +impl StoreValue for &mut ThreadData +where + Alignment: SupportedAlignment, +{ + #[inline(always)] + fn store( + self, + group: ThreadGroup, + destination: &mut DisjointSlice<'_, T, CollectiveTiles>, + options: StoreOptions, + ) { + (&*self).store(group, destination, options); + } +} +impl sealed::StoreValue for T {} +impl StoreValue for T { + #[inline(always)] + fn store( + self, + group: ThreadGroup, + destination: &mut DisjointSlice<'_, T, CollectiveTiles>, + options: StoreOptions, + ) { + let mut value = ThreadData::<_, 1>::new(); + value.write(0, self); + (&value).store(group, destination, options); + } +} + +mod sealed { + pub trait LoadSource {} + pub trait StoreValue {} +} diff --git a/rust/cuda-coop/src/load_store/boundary.rs b/rust/cuda-coop/src/load_store/boundary.rs new file mode 100644 index 000000000000..9ae7f1e54ab1 --- /dev/null +++ b/rust/cuda-coop/src/load_store/boundary.rs @@ -0,0 +1,148 @@ +// Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +use super::{LoadOptions, StoreOptions}; +use crate::{Alignment, Numeric, SupportedAlignment, ThreadData}; +use cuda_device::{CollectiveTiles, DisjointSlice, library, thread}; + +macro_rules! registration { + ($name:ident, $variable:literal) => { + fn $name() -> &'static str { + library::registration!($variable) + } + }; +} +registration!(load_full, "COOP_VALIDITY_LOAD_FULL"); +registration!(load_masked, "COOP_VALIDITY_LOAD_MASKED"); +registration!(load_filled, "COOP_VALIDITY_LOAD_FILLED"); +registration!(store_full, "COOP_VALIDITY_STORE_FULL"); +registration!(store_partial, "COOP_VALIDITY_STORE_PARTIAL"); + +#[inline(always)] +pub(super) fn tile() -> usize { + let threads = thread::blockDim_x() as usize + * thread::blockDim_y() as usize + * thread::blockDim_z() as usize; + assert!(ITEMS > 0 && threads <= usize::MAX / ITEMS); + threads * ITEMS +} + +#[inline(always)] +fn extent(len: usize, offset: usize, valid: usize) { + assert!(valid <= tile::() && valid <= i32::MAX as usize); + assert!(offset <= len && valid <= len - offset); + assert!(offset <= i64::MAX as usize); + assert!(offset <= isize::MAX as usize / core::mem::size_of::()); +} + +#[inline(always)] +pub(super) fn containment(offset: usize, valid: usize) { + let capacity = tile::(); + let groups = + thread::gridDim_x() as usize * thread::gridDim_y() as usize * thread::gridDim_z() as usize; + let rank = thread::blockIdx_x() as usize + + thread::gridDim_x() as usize + * (thread::blockIdx_y() as usize + + thread::gridDim_y() as usize * thread::blockIdx_z() as usize); + assert!(groups > 0 && groups <= usize::MAX / capacity); + let period = groups * capacity; + if valid != 0 { + let within = offset % period; + let begin = rank * capacity; + assert!(within >= begin && within - begin <= capacity); + assert!(valid <= capacity - (within - begin)); + } +} + +// The caller binds pointer/length to this immutable allocation witness and +// checks collective containment when the source is a writable carrier. +#[inline(always)] +pub(super) unsafe fn load( + memory: Memory, + pointer: *const T, + len: usize, + output: &mut ThreadData, + options: LoadOptions, +) where + Alignment: SupportedAlignment, +{ + let valid = match options.valid_items { + Some(count) => count, + None => tile::(), + }; + extent::(len, options.offset, valid); + // SAFETY: the caller binds pointer/length to the source allocation; bounds + // are checked before forming the pointer, including an empty prefix. + let pointer = unsafe { pointer.add(options.offset) }; + let payload = output.block_io_parts(); + // SAFETY: the checked marker enforces participation, uniform controls and + // resource epochs. The exact typed provider initializes values before flags + // and never interprets a masked invalid slot as T. No initialized slice is + // exposed. Each branch consumes the same source witness. + unsafe { + match (options.valid_items, options.oob_default) { + (None, None) => library::checked_block_resource_call_v1::( + load_full, + memory, + (pointer,), + payload, + ), + (Some(count), None) => { + library::checked_block_resource_call_v1::( + load_masked, + memory, + (pointer, count as i32), + payload, + ) + } + (Some(count), Some(default)) => { + library::checked_block_resource_call_v1::( + load_filled, + memory, + (pointer, count as i32, default), + payload, + ) + } + (None, Some(_)) => panic!(), + } + } +} + +#[inline(always)] +pub(super) fn store( + destination: &mut DisjointSlice<'_, T, CollectiveTiles>, + value: &ThreadData, + options: StoreOptions, +) where + Alignment: SupportedAlignment, +{ + let valid = match options.valid_items { + Some(count) => count, + None => tile::(), + }; + extent::(destination.len(), options.offset, valid); + containment::(options.offset, valid); + let pointer = destination.as_mut_ptr(); + let payload = value.block_io_read_parts(); + // SAFETY: this carrier owns the checked write range. The provider checks + // the entire lane-local write domain before reading numeric bytes and takes + // const payload/flag pointers. It preserves all slots, including invalid + // ones. The compiler checks uniform controls, participation and epochs. + unsafe { + let pointer = pointer.add(options.offset); + match options.valid_items { + None => library::checked_block_resource_call_v1::( + store_full, + destination, + (pointer,), + payload, + ), + Some(count) => library::checked_block_resource_call_v1::( + store_partial, + destination, + (pointer, count as i32), + payload, + ), + } + } +} diff --git a/rust/cuda-coop/src/thread_data.rs b/rust/cuda-coop/src/thread_data.rs index 11880d852c5f..76203c847736 100644 --- a/rust/cuda-coop/src/thread_data.rs +++ b/rust/cuda-coop/src/thread_data.rs @@ -282,3 +282,6 @@ where Alignment: SupportedAlignment, { } + +#[path = "thread_data/block_io.rs"] +mod block_io; diff --git a/rust/cuda-coop/src/thread_data/block_io.rs b/rust/cuda-coop/src/thread_data/block_io.rs index 48da4cf37b34..a155164ad459 100644 --- a/rust/cuda-coop/src/thread_data/block_io.rs +++ b/rust/cuda-coop/src/thread_data/block_io.rs @@ -1,15 +1,20 @@ // Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception -// Private native-boundary bridge, compiled only by the coordinated fixture -// until the complete collective API is qualified. These addresses describe -// opaque MaybeUninit cells and a separate packed validity bitmap, never a T slice. +// Private native-boundary bridge. These addresses describe opaque +// MaybeUninit cells and a separate packed validity bitmap, never a T slice. use super::{Alignment, Numeric, SupportedAlignment, ThreadData}; impl ThreadData where Alignment: SupportedAlignment, { + #[inline(always)] + pub(crate) fn block_io_read_parts(&self) -> (*const T, *const u8) { + let () = Self::POSITIVE_EXTENT; + (self.values.as_ptr().cast::(), self.initialized.as_ptr()) + } + #[inline(always)] pub(crate) fn block_io_parts(&mut self) -> (*mut T, *mut u8) { let () = Self::POSITIVE_EXTENT; diff --git a/rust/cuda-coop/tests/check_load_store.py b/rust/cuda-coop/tests/check_load_store.py new file mode 100644 index 000000000000..d29d8c0b397b --- /dev/null +++ b/rust/cuda-coop/tests/check_load_store.py @@ -0,0 +1,147 @@ +# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. +# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +"""Compile public Load/Store participation and resource-epoch contrasts. + +Use the same Cargo config/backend as the qualified public example. The config +is optional for an installed compatible Oxide revision; local source integration +can supply explicit Cargo patches without changing checked-in manifests. +""" + +import argparse +import json +import os +import shutil +import subprocess +from pathlib import Path + +parser = argparse.ArgumentParser(description=__doc__) +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("--cargo-config", 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() +output = args.output.resolve() +output.mkdir(parents=True, exist_ok=False) +(output / "src").mkdir() +package = Path(__file__).resolve().parents[1] +manifest = (package / "tests/load_store/Cargo.toml").read_text() +manifest = manifest.replace('path = "../.."', f"path = {json.dumps(str(package))}") +(output / "Cargo.toml").write_text(manifest) +shutil.copyfile(package / "tests/load_store/Cargo.lock", output / "Cargo.lock") +source = output / "src/main.rs" +prefix = """// Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception +use cuda_coop as coop; +use cuda_device::{DisjointSlice,CollectiveTiles,kernel,launch_contract,thread}; +#[kernel] +#[launch_contract(domain=1,block=(128,1,1),dynamic_shared=inferred)] +pub fn probe(input:&[i32],mut output:DisjointSlice>,valid:usize) { + let block=coop::this_block(); + let offset=thread::blockIdx_x() as usize*512; + let mut values=coop::ThreadData::<_,4>::new(); +""" +load = "coop::load(block,input,&mut values,coop::LoadOptions::default().offset(offset).valid_items(valid).oob_default(0));\n" +store = "coop::store(block,&mut output,&values,coop::StoreOptions::default().offset(offset).valid_items(valid));\n" +reload = "coop::load(block,&mut output,&mut values,coop::LoadOptions::default().offset(offset).valid_items(valid));\n" +cases = { + "varying_presence": ( + "let options=if thread::threadIdx_x()&1==0 { coop::LoadOptions::default().valid_items(valid) } else { coop::LoadOptions::default() };coop::load(block,input,&mut values,options);", + "lane-varying", + ), + "varying_count": ( + load.replace( + ".valid_items(valid)", ".valid_items(thread::threadIdx_x() as usize)" + ), + "not proved block-uniform", + ), + "varying_default": ( + load.replace(".oob_default(0)", ".oob_default(thread::threadIdx_x() as i32)"), + "not proved block-uniform", + ), + "varying_offset": ( + load.replace( + ".offset(offset)", ".offset(offset+thread::threadIdx_x() as usize)" + ), + "not proved block-uniform", + ), + "varying_participation": ( + "if thread::threadIdx_x()&1==0 {" + load + "}", + "lane-varying", + ), + "early_return": ("if thread::threadIdx_x()!=0 { return; }" + load, "lane-varying"), + "missing_completion": (load + store + reload, "requires a full-block barrier"), + "loop_without_completion": ( + load + "let mut iteration=0;while iteration<2 {" + store + "iteration+=1;}", + "requires a full-block barrier", + ), + "loop_with_completion": ( + load + + "let mut iteration=0;while iteration<2 {" + + store + + "thread::sync_threads();iteration+=1;}", + None, + ), + "explicit_completion": (load + store + "thread::sync_threads();" + reload, None), + "shared_payload": ( + load + "let shared=&values;" + store.replace("&values", "shared"), + None, + ), + "reconverged_selection": ( + "let mut other=coop::ThreadData::<_,4>::new();values.write(0,7);other.write(0,9);" + "let selected=if thread::threadIdx_x()&1==0 {values} else {other};" + "coop::store(block,&mut output,&selected,coop::StoreOptions::default().offset(offset).valid_items(1));", + None, + ), +} +command = [ + str(args.cargo_oxide.resolve()), + "oxide", + "build", + "--arch", + args.arch, + "--materialize-cubin", + "--", + "--locked", +] +if args.cargo_config: + command += ["--config", str(args.cargo_config.resolve())] +env = os.environ.copy() +env.update( + CUDA_OXIDE_BACKEND=str(args.backend.resolve()), + CUDA_HOME=str(args.cuda_root.resolve()), + CUDA_PATH=str(args.cuda_root.resolve()), +) +results = [] +for name, (body, diagnostic) in cases.items(): + text = prefix + body + "\n}\nfn main() {}\n" + source.write_text(text) + (output / f"{name}.rs").write_text(text) + run = subprocess.run( + command, + cwd=output, + env=env, + text=True, + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + ) + (output / f"{name}.log").write_text(run.stdout) + passed = ( + run.returncode == 0 + if diagnostic is None + else run.returncode != 0 and diagnostic in run.stdout + ) + results.append( + {"case": name, "expected": diagnostic, "exit": run.returncode, "pass": passed} + ) + print(name, "PASS" if passed else "FAIL", flush=True) +(output / "results.json").write_text( + json.dumps( + {"command": command, "backend": env["CUDA_OXIDE_BACKEND"], "cases": results}, + indent=2, + ) + + "\n" +) +raise SystemExit(0 if all(row["pass"] for row in results) else 1) diff --git a/rust/cuda-coop/tests/load_store/.gitignore b/rust/cuda-coop/tests/load_store/.gitignore new file mode 100644 index 000000000000..74bffa8af4a4 --- /dev/null +++ b/rust/cuda-coop/tests/load_store/.gitignore @@ -0,0 +1,7 @@ +# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. +# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception +/target/ +/.oxide-artifacts/ +/*.ll +/*.options +/*.target diff --git a/rust/cuda-coop/tests/load_store/Cargo.lock b/rust/cuda-coop/tests/load_store/Cargo.lock new file mode 100644 index 000000000000..133f595cdf96 --- /dev/null +++ b/rust/cuda-coop/tests/load_store/Cargo.lock @@ -0,0 +1,1010 @@ +# This file is automatically @generated by Cargo. +# It is not intended for manual editing. +version = 4 + +[[package]] +name = "aho-corasick" +version = "1.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c982642fa9e8606056828ee9a8505737230110bb1099153c79efe865c59d12ba" +dependencies = [ + "memchr", +] + +[[package]] +name = "anyhow" +version = "1.0.104" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "330a5ed07fa54e4702c9d6c4174f74427fc0ef6e214bbd677ae50a5099946470" + +[[package]] +name = "bindgen" +version = "0.69.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "271383c67ccabffb7381723dea0672a673f292304fcb45c01cc648c7a8d58088" +dependencies = [ + "bitflags", + "cexpr", + "clang-sys", + "itertools", + "lazy_static", + "lazycell", + "log", + "prettyplease", + "proc-macro2", + "quote", + "regex", + "rustc-hash 1.1.0", + "shlex", + "syn 2.0.119", + "which", +] + +[[package]] +name = "bitflags" +version = "2.13.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3ded4057c258ba199e2d26386d3af3780957ecaee6c4ef4041c6b4b8b97c0b06" + +[[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 = "cexpr" +version = "0.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6fac387a98bb7c37292057cffc56d62ecb629900026402633ae9160df93a8766" +dependencies = [ + "nom", +] + +[[package]] +name = "cfg-if" +version = "1.0.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4e7648175b45a9a48536d676f68d918270699102aa8dab5496df06904c914600" + +[[package]] +name = "clang-sys" +version = "1.9.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "157a8ba7b480713b56f4c09fd13fc3e0a22a5dfab8097ba61cbc5feef950788a" +dependencies = [ + "glob", + "libc", + "libloading", +] + +[[package]] +name = "coop-load-store-tests" +version = "0.0.0" +dependencies = [ + "cuda-async", + "cuda-coop", + "cuda-core", + "cuda-device", + "cuda-host", +] + +[[package]] +name = "cpufeatures" +version = "0.2.17" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "59ed5838eebb26a2bb2e58f6d5b5316989ae9d08bab10e0e6d103e656d1b0280" +dependencies = [ + "libc", +] + +[[package]] +name = "crc32fast" +version = "1.5.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0ba106e844045ddc02c20768ac4b0b7cb45f2cc2ecbd833ac9f5dacb183c2055" +dependencies = [ + "cfg-if", +] + +[[package]] +name = "crossbeam-utils" +version = "0.8.23" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a31eee39dddec8330830986fcd7625edb5a24ec90ea038215273bbc3adb08ac6" + +[[package]] +name = "crunchy" +version = "0.2.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "460fbee9c2c2f33933d720630a6a0bac33ba7053db5344fac858d4b8952d77d5" + +[[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-artifact-finalizer" +version = "0.2.1" +source = "git+https://github.com/tpn/cuda-rust.git?rev=be42c33eeb15ea5639fdd7d793a76706f10cb19e#be42c33eeb15ea5639fdd7d793a76706f10cb19e" +dependencies = [ + "libnvvm-sys", + "nvjitlink-sys", + "serde", + "sha2", + "thiserror 2.0.21", +] + +[[package]] +name = "cuda-async" +version = "0.4.0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=be42c33eeb15ea5639fdd7d793a76706f10cb19e#be42c33eeb15ea5639fdd7d793a76706f10cb19e" +dependencies = [ + "anyhow", + "cuda-bindings", + "cuda-core", + "dashmap", + "futures", + "half", + "once_cell", + "rustc-hash 2.1.3", + "thiserror 1.0.69", +] + +[[package]] +name = "cuda-bindings" +version = "0.4.0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=be42c33eeb15ea5639fdd7d793a76706f10cb19e#be42c33eeb15ea5639fdd7d793a76706f10cb19e" +dependencies = [ + "bindgen", + "libloading", + "prettyplease", + "proc-macro2", + "quote", + "syn 2.0.119", +] + +[[package]] +name = "cuda-coop" +version = "0.0.0" +dependencies = [ + "cuda-coop-codegen", + "cuda-device", +] + +[[package]] +name = "cuda-coop-codegen" +version = "0.0.0" +dependencies = [ + "libloading", + "serde", + "serde_json", + "sha2", +] + +[[package]] +name = "cuda-core" +version = "0.4.0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=be42c33eeb15ea5639fdd7d793a76706f10cb19e#be42c33eeb15ea5639fdd7d793a76706f10cb19e" +dependencies = [ + "anyhow", + "cuda-bindings", + "cuda-core-derive", + "half", + "oxide-artifacts", +] + +[[package]] +name = "cuda-core-derive" +version = "0.4.0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=be42c33eeb15ea5639fdd7d793a76706f10cb19e#be42c33eeb15ea5639fdd7d793a76706f10cb19e" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", +] + +[[package]] +name = "cuda-device" +version = "0.2.1" +source = "git+https://github.com/tpn/cuda-rust.git?rev=be42c33eeb15ea5639fdd7d793a76706f10cb19e#be42c33eeb15ea5639fdd7d793a76706f10cb19e" +dependencies = [ + "cuda-macros", +] + +[[package]] +name = "cuda-host" +version = "0.2.1" +source = "git+https://github.com/tpn/cuda-rust.git?rev=be42c33eeb15ea5639fdd7d793a76706f10cb19e#be42c33eeb15ea5639fdd7d793a76706f10cb19e" +dependencies = [ + "cuda-artifact-finalizer", + "cuda-async", + "cuda-core", + "cuda-macros", + "half", + "libc", + "oxide-artifacts", + "ptx-parse", + "sha2", + "thiserror 2.0.21", +] + +[[package]] +name = "cuda-macros" +version = "0.2.1" +source = "git+https://github.com/tpn/cuda-rust.git?rev=be42c33eeb15ea5639fdd7d793a76706f10cb19e#be42c33eeb15ea5639fdd7d793a76706f10cb19e" +dependencies = [ + "oxide-tile-export", + "proc-macro2", + "quote", + "reserved-oxide-symbols", + "serde", + "serde_json", + "syn 2.0.119", +] + +[[package]] +name = "cuda-target-spec" +version = "0.1.0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=be42c33eeb15ea5639fdd7d793a76706f10cb19e#be42c33eeb15ea5639fdd7d793a76706f10cb19e" + +[[package]] +name = "dashmap" +version = "6.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e6361d5c062261c78a176addb82d4c821ae42bed6089de0e12603cd25de2059c" +dependencies = [ + "cfg-if", + "crossbeam-utils", + "hashbrown 0.14.5", + "lock_api", + "once_cell", + "parking_lot_core", +] + +[[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 = "either" +version = "1.19.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0e9c71c2167ca323c882b99918929403426e2373ea17242ff5653e0d5e1058be" + +[[package]] +name = "equivalent" +version = "1.0.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "00d174d5400e5e8fd687ad1049e2f578285fa914201b1af7e8b112a4546bd826" + +[[package]] +name = "errno" +version = "0.3.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "39cab71617ae0d63f51a36d69f866391735b51691dbda63cf6f96d042b63efeb" +dependencies = [ + "libc", + "windows-sys 0.61.2", +] + +[[package]] +name = "foldhash" +version = "0.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d9c4f5dac5e15c24eb999c26181a6ca40b39fe946cbe4c263c7209467bc83af2" + +[[package]] +name = "futures" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9a31d2a3fbaaeb2af2368bbdd904aa8e812d3c04a1ee10d3171f52d556e5d0a3" +dependencies = [ + "futures-channel", + "futures-core", + "futures-executor", + "futures-io", + "futures-sink", + "futures-task", + "futures-util", +] + +[[package]] +name = "futures-channel" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b1f9e3d69d39e4862ffed03ed071a76f9a13ba1d9109d355b0f0aa6b15e393c4" +dependencies = [ + "futures-core", + "futures-sink", +] + +[[package]] +name = "futures-core" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "92d699e522242e69e3003b94ecc1f960f3a5e015aa7c5d7486e65ad01dd94f5e" + +[[package]] +name = "futures-executor" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "031b47cf1a3c6cc8bc2fc76cd437f521619387907d469316e7c0bc278f1f5432" +dependencies = [ + "futures-core", + "futures-task", + "futures-util", +] + +[[package]] +name = "futures-io" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "53c0fa8157de1303bfffdaa1cc2a673bfffb60102f76b0ef4441659124373fed" + +[[package]] +name = "futures-macro" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9fb9654ba8355388abeb8dcb4fc62f511300867002afc858860463bdd9fe0c44" +dependencies = [ + "proc-macro2", + "quote", + "syn 3.0.9", +] + +[[package]] +name = "futures-sink" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1944426bf7d03f1d14f708785e4b33efd750b36d48a157b836b3efc15ede8e1d" + +[[package]] +name = "futures-task" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cd417de3d1d015fc3bfd2b1ea46dfc7bab72ef86f1cc7cc9c78e728b34a6d1fd" + +[[package]] +name = "futures-util" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0d50a92467f8ba5dd6e3ee5d4bd04d73ab2e4e1c44474a0674821dfce14b79bc" +dependencies = [ + "futures-channel", + "futures-core", + "futures-io", + "futures-macro", + "futures-sink", + "futures-task", + "memchr", + "pin-project-lite", + "slab", +] + +[[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 = "glob" +version = "0.3.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e4eba85ea1d0a966a983acd07deee566e67395d2d96b6fb39e62b5a833f1eb0b" + +[[package]] +name = "half" +version = "2.7.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" +dependencies = [ + "cfg-if", + "crunchy", + "zerocopy", +] + +[[package]] +name = "hashbrown" +version = "0.14.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e5274423e17b7c9fc20b6e7e208532f9b19825d82dfd615708b70edd83df41f1" + +[[package]] +name = "hashbrown" +version = "0.15.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9229cfe53dfd69f0609a49f65461bd93001ea1ef889cd5529dd176593f5338a1" +dependencies = [ + "foldhash", +] + +[[package]] +name = "hashbrown" +version = "0.17.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ed5909b6e89a2db4456e54cd5f673791d7eca6732202bbf2a9cc504fe2f9b84a" + +[[package]] +name = "home" +version = "0.5.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cc627f471c528ff0c4a49e1d5e60450c8f6461dd6d10ba9dcd3a61d3dff7728d" +dependencies = [ + "windows-sys 0.61.2", +] + +[[package]] +name = "indexmap" +version = "2.14.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cc4e190f5d26ca7051642629da2c52fc03bde85a03197c99408dcd291734c855" +dependencies = [ + "equivalent", + "hashbrown 0.17.1", +] + +[[package]] +name = "itertools" +version = "0.12.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ba291022dbbd398a455acf126c1e341954079855bc60dfdda641363bd6922569" +dependencies = [ + "either", +] + +[[package]] +name = "itoa" +version = "1.0.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8f42a60cbdf9a97f5d2305f08a87dc4e09308d1276d28c869c684d7777685682" + +[[package]] +name = "lazy_static" +version = "1.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "20870f649af7073d53e38067b2a84312175d56ea15217e1b15bc83506ec50afb" + +[[package]] +name = "lazycell" +version = "1.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "830d08ce1d1d941e6b30645f1a0eb5643013d835ce3779a5fc208261dbe10f55" + +[[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 = "libnvvm-sys" +version = "0.2.1" +source = "git+https://github.com/tpn/cuda-rust.git?rev=be42c33eeb15ea5639fdd7d793a76706f10cb19e#be42c33eeb15ea5639fdd7d793a76706f10cb19e" +dependencies = [ + "cuda-target-spec", + "libloading", + "thiserror 2.0.21", +] + +[[package]] +name = "linux-raw-sys" +version = "0.4.15" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d26c52dbd32dccf2d10cac7725f8eae5296885fb5703b261f7d0a0739ec807ab" + +[[package]] +name = "lock_api" +version = "0.4.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "224399e74b87b5f3557511d98dff8b14089b3dadafcab6bb93eab67d3aace965" +dependencies = [ + "scopeguard", +] + +[[package]] +name = "log" +version = "0.4.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9f8bd3e56ce4dfc153cf470fffbfa98c7620958b312ca5c3a4b8d5181fd13c6" + +[[package]] +name = "memchr" +version = "2.8.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cf8baf1c55e62ffcace7a9f06f4bd9cd3f0c4beb022d3b367256b91b87513d98" + +[[package]] +name = "minimal-lexical" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "68354c5c6bd36d73ff3feceb05efa59b6acb7626617f4962be322a825e61f79a" + +[[package]] +name = "nom" +version = "7.1.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d273983c5a657a70a3e8f2a01329822f3b8c8172b73826411a55751e404a0a4a" +dependencies = [ + "memchr", + "minimal-lexical", +] + +[[package]] +name = "nvjitlink-sys" +version = "0.2.1" +source = "git+https://github.com/tpn/cuda-rust.git?rev=be42c33eeb15ea5639fdd7d793a76706f10cb19e#be42c33eeb15ea5639fdd7d793a76706f10cb19e" +dependencies = [ + "libloading", + "thiserror 2.0.21", +] + +[[package]] +name = "object" +version = "0.36.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "62948e14d923ea95ea2c7c86c71013138b66525b86bdc08d2dcc262bdb497b87" +dependencies = [ + "crc32fast", + "hashbrown 0.15.5", + "indexmap", + "memchr", +] + +[[package]] +name = "once_cell" +version = "1.21.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9f7c3e4beb33f85d45ae3e3a1792185706c8e16d043238c593331cc7cd313b50" + +[[package]] +name = "oxide-artifacts" +version = "0.2.1" +source = "git+https://github.com/tpn/cuda-rust.git?rev=be42c33eeb15ea5639fdd7d793a76706f10cb19e#be42c33eeb15ea5639fdd7d793a76706f10cb19e" +dependencies = [ + "object", +] + +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +source = "git+https://github.com/tpn/cuda-rust.git?rev=be42c33eeb15ea5639fdd7d793a76706f10cb19e#be42c33eeb15ea5639fdd7d793a76706f10cb19e" +dependencies = [ + "serde", + "serde_json", + "thiserror 2.0.21", +] + +[[package]] +name = "parking_lot_core" +version = "0.9.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2621685985a2ebf1c516881c026032ac7deafcda1a2c9b7850dc81e3dfcb64c1" +dependencies = [ + "cfg-if", + "libc", + "redox_syscall", + "smallvec", + "windows-link", +] + +[[package]] +name = "pin-project-lite" +version = "0.2.17" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a89322df9ebe1c1578d689c92318e070967d1042b512afbe49518723f4e6d5cd" + +[[package]] +name = "prettyplease" +version = "0.2.37" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "479ca8adacdd7ce8f1fb39ce9ecccbfe93a3f1344b3d0d97f20bc0196208f62b" +dependencies = [ + "proc-macro2", + "syn 2.0.119", +] + +[[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 = "ptx-parse" +version = "0.1.0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=be42c33eeb15ea5639fdd7d793a76706f10cb19e#be42c33eeb15ea5639fdd7d793a76706f10cb19e" + +[[package]] +name = "quote" +version = "1.0.47" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1fbf4db142a473a8d80c26bbf18454ed458bf8d26c8219c331daecfdbd079001" +dependencies = [ + "proc-macro2", +] + +[[package]] +name = "redox_syscall" +version = "0.5.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ed2bf2547551a7053d6fdfafda3f938979645c44812fbfcda098faae3f1a362d" +dependencies = [ + "bitflags", +] + +[[package]] +name = "regex" +version = "1.13.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f020237b6c8eed93db2e2cb53c00c60a8e1bc73da7d073199a1180401450218d" +dependencies = [ + "aho-corasick", + "memchr", + "regex-automata", + "regex-syntax", +] + +[[package]] +name = "regex-automata" +version = "0.4.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ad8553b9b26413251cbf30e620595c7a41b3887f03da04579c0e6b0d6a06b4b2" +dependencies = [ + "aho-corasick", + "memchr", + "regex-syntax", +] + +[[package]] +name = "regex-syntax" +version = "0.8.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d6f6ff9a378485b298a5286656da665ba74413d36db0979633275d2e708145d4" + +[[package]] +name = "reserved-oxide-symbols" +version = "0.2.1" +source = "git+https://github.com/tpn/cuda-rust.git?rev=be42c33eeb15ea5639fdd7d793a76706f10cb19e#be42c33eeb15ea5639fdd7d793a76706f10cb19e" + +[[package]] +name = "rustc-hash" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "08d43f7aa6b08d49f382cde6a7982047c3426db949b1424bc4b7ec9ae12c6ce2" + +[[package]] +name = "rustc-hash" +version = "2.1.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6b1e7f9a428571be2dc5bc0505c13fb6bf936822b894ec87abf8a08a4e51742d" + +[[package]] +name = "rustix" +version = "0.38.44" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fdb5bc1ae2baa591800df16c9ca78619bf65c0488b41b96ccec5d11220d8c154" +dependencies = [ + "bitflags", + "errno", + "libc", + "linux-raw-sys", + "windows-sys 0.59.0", +] + +[[package]] +name = "scopeguard" +version = "1.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "94143f37725109f92c262ed2cf5e59bce7498c01bcc1502d7b9afe439a4e9f49" + +[[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 3.0.9", +] + +[[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 = "shlex" +version = "1.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0fda2ff0d084019ba4d7c6f371c95d8fd75ce3524c3cb8fb653a3023f6323e64" + +[[package]] +name = "slab" +version = "0.4.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0c790de23124f9ab44544d7ac05d60440adc586479ce501c1d6d7da3cd8c9cf5" + +[[package]] +name = "smallvec" +version = "1.16.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5b3dc8af474f516a851ff4bd12db780f948b9250ad37211e4eec0bccea54e01b" + +[[package]] +name = "syn" +version = "2.0.119" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "872831b642d1a07999a962a351ed35b955ea2cfc8f3862091e2a240a84f17297" +dependencies = [ + "proc-macro2", + "quote", + "unicode-ident", +] + +[[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 = "thiserror" +version = "1.0.69" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b6aaf5339b578ea85b50e080feb250a3e8ae8cfcdff9a461c9ec2904bc923f52" +dependencies = [ + "thiserror-impl 1.0.69", +] + +[[package]] +name = "thiserror" +version = "2.0.21" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "09e52cb86a36cede5cb101bf8908837b3e4c6e5e59fe7fd85c23fb56200d189e" +dependencies = [ + "thiserror-impl 2.0.21", +] + +[[package]] +name = "thiserror-impl" +version = "1.0.69" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4fee6c4efc90059e10f81e6d42c60a18f76588c3d74cb83a0b242a2b6c7504c1" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", +] + +[[package]] +name = "thiserror-impl" +version = "2.0.21" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fe5197923287db20a58125f0bc85c062f7f2c892de97b18c356f9efb14b28524" +dependencies = [ + "proc-macro2", + "quote", + "syn 3.0.9", +] + +[[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 = "which" +version = "4.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "87ba24419a2078cd2b0f2ede2691b6c66d8e47836da3b6db8265ebad47afbfc7" +dependencies = [ + "either", + "home", + "once_cell", + "rustix", +] + +[[package]] +name = "windows-link" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f0805222e57f7521d6a62e36fa9163bc891acd422f971defe97d64e70d0a4fe5" + +[[package]] +name = "windows-sys" +version = "0.59.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1e38bc4d79ed67fd075bcc251a1c39b32a1776bbe92e5bef1f0bf1f8c531853b" +dependencies = [ + "windows-targets", +] + +[[package]] +name = "windows-sys" +version = "0.61.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ae137229bcbd6cdf0f7b80a31df61766145077ddf49416a728b02cb3921ff3fc" +dependencies = [ + "windows-link", +] + +[[package]] +name = "windows-targets" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9b724f72796e036ab90c1021d4780d4d3d648aca59e491e6b98e725b84e99973" +dependencies = [ + "windows_aarch64_gnullvm", + "windows_aarch64_msvc", + "windows_i686_gnu", + "windows_i686_gnullvm", + "windows_i686_msvc", + "windows_x86_64_gnu", + "windows_x86_64_gnullvm", + "windows_x86_64_msvc", +] + +[[package]] +name = "windows_aarch64_gnullvm" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "32a4622180e7a0ec044bb555404c800bc9fd9ec262ec147edd5989ccd0c02cd3" + +[[package]] +name = "windows_aarch64_msvc" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "09ec2a7bb152e2252b53fa7803150007879548bc709c039df7627cabbd05d469" + +[[package]] +name = "windows_i686_gnu" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8e9b5ad5ab802e97eb8e295ac6720e509ee4c243f69d781394014ebfe8bbfa0b" + +[[package]] +name = "windows_i686_gnullvm" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0eee52d38c090b3caa76c563b86c3a4bd71ef1a819287c19d586d7334ae8ed66" + +[[package]] +name = "windows_i686_msvc" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "240948bc05c5e7c6dabba28bf89d89ffce3e303022809e73deaefe4f6ec56c66" + +[[package]] +name = "windows_x86_64_gnu" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "147a5c80aabfbf0c7d901cb5895d1de30ef2907eb21fbbab29ca94c5b08b1a78" + +[[package]] +name = "windows_x86_64_gnullvm" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "24d5b23dc417412679681396f2b49f3de8c1473deb516bd34410872eff51ed0d" + +[[package]] +name = "windows_x86_64_msvc" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "589f6da84c646204747d1270a2a5661ea66ed1cced2631d546fdfb155959f9ec" + +[[package]] +name = "zerocopy" +version = "0.8.63" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e5fe1f8f1b06191a00962174c61aa5005e0bb391a6d80d07e24d115c01a92ed8" +dependencies = [ + "zerocopy-derive", +] + +[[package]] +name = "zerocopy-derive" +version = "0.8.63" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "863ad3ac83293fb4d740aedbfdc9240dd8d1a50c1099acd76ce80ce7c7230c7f" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", +] + +[[package]] +name = "zmij" +version = "1.0.23" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "29666d0abbfad1e3dc4dcf6144730dd3a3ab225bbbdac83319345b1b44ccfc1b" diff --git a/rust/cuda-coop/tests/load_store/Cargo.toml b/rust/cuda-coop/tests/load_store/Cargo.toml new file mode 100644 index 000000000000..6dbdea6cee62 --- /dev/null +++ b/rust/cuda-coop/tests/load_store/Cargo.toml @@ -0,0 +1,17 @@ +# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. +# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +[package] +name = "coop-load-store-tests" +version = "0.0.0" +edition = "2024" +publish = false + +[workspace] + +[dependencies] +cuda-coop = { path = "../.." } +cuda-device = { git = "https://github.com/tpn/cuda-rust.git", rev = "be42c33eeb15ea5639fdd7d793a76706f10cb19e" } +cuda-host = { git = "https://github.com/tpn/cuda-rust.git", rev = "be42c33eeb15ea5639fdd7d793a76706f10cb19e", features = ["async"] } +cuda-core = { git = "https://github.com/tpn/cuda-rust.git", rev = "be42c33eeb15ea5639fdd7d793a76706f10cb19e" } +cuda-async = { git = "https://github.com/tpn/cuda-rust.git", rev = "be42c33eeb15ea5639fdd7d793a76706f10cb19e" } diff --git a/rust/cuda-coop/tests/load_store/src/main.rs b/rust/cuda-coop/tests/load_store/src/main.rs new file mode 100644 index 000000000000..920e04b4c7d5 --- /dev/null +++ b/rust/cuda-coop/tests/load_store/src/main.rs @@ -0,0 +1,626 @@ +// Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +use coop::{Numeric, ThreadData, this_block}; +use cuda_coop as coop; +use cuda_device::{CollectiveTiles, DisjointSlice, kernel, launch_contract, thread}; +use cuda_host::cuda_module; + +#[cuda_module] +mod kernels { + use super::*; + + #[inline(never)] + fn exercise( + input: &[T], + output: &mut DisjointSlice>, + valid: usize, + shift: usize, + default: T, + mode: u32, + ) { + let block = this_block(); + let threads = thread::blockDim_x() as usize; + let lane = thread::threadIdx_x() as usize; + let offset = thread::blockIdx_x() as usize * threads * 4 + shift; + let mut values = ThreadData::<_, ITEMS_PER_THREAD>::new(); + match mode { + 0 => coop::load( + block, + input, + &mut values, + coop::LoadOptions::default().offset(offset), + ), + 1 | 5 | 6 => { + coop::load( + block, + input, + &mut values, + coop::LoadOptions::default().offset(offset), + ); + coop::load( + block, + input, + &mut values, + coop::LoadOptions::default() + .offset(offset) + .valid_items(valid), + ); + } + 2 | 4 => coop::load( + block, + input, + &mut values, + coop::LoadOptions::default() + .offset(offset) + .valid_items(valid) + .oob_default(default), + ), + 9 => {} + 10 => coop::load( + block, + input, + &mut values, + coop::LoadOptions::default() + .offset(offset) + .oob_default(default), + ), + 11 => { + coop::load( + block, + input, + &mut values, + coop::LoadOptions::default().offset(offset).valid_items(0), + ); + let _invalid = values[0]; + } + _ => { + // Real nonadjacent initialized slots beyond the stored prefix. + // The two masks differ, then reconverged lane-local selection + // moves each value together with its validity metadata. + let mut other = ThreadData::<_, ITEMS_PER_THREAD>::new(); + let mut item = 0; + while item < ITEMS_PER_THREAD { + let prefix = lane * ITEMS_PER_THREAD + item < valid; + if prefix || item & 1 == 1 { + values.write(item, default); + } + if prefix || item & 1 == 0 { + other.write(item, default); + } + item += 1; + } + values = if lane & 1 == 0 { values } else { other }; + } + } + let before = values; + if mode == 0 || mode == 2 { + coop::store( + block, + &mut *output, + &values, + coop::StoreOptions::default().offset(offset), + ); + } else { + let written = if mode == 5 { + valid + 1 + } else if mode == 6 { + valid / 2 + } else { + valid + }; + coop::store( + block, + &mut *output, + &values, + coop::StoreOptions::default() + .offset(offset) + .valid_items(written), + ); + } + let mut observations_ok = true; + let mut item = 0; + while item < ITEMS_PER_THREAD { + if let Some(value) = before.get(item) { + observations_ok &= values.get(item) == Some(value); + } else { + observations_ok &= values.get(item).is_none(); + } + if mode == 1 || mode == 6 { + observations_ok &= + values.get(item).is_some() == (lane * ITEMS_PER_THREAD + item < valid); + } + item += 1; + } + if mode == 4 { + thread::sync_threads(); + coop::load( + block, + &mut *output, + &mut values, + coop::LoadOptions::default() + .offset(offset) + .valid_items(valid) + .oob_default(default), + ); + thread::sync_threads(); + coop::store( + block, + &mut *output, + &values, + coop::StoreOptions::default() + .offset(offset) + .valid_items(valid), + ); + } + assert!(observations_ok); + // Final writes complete through the owned async host operation. + } + + #[kernel] + #[launch_contract(domain = 1, block = (128, 1, 1), dynamic_shared = inferred)] + pub fn copy_i8_128( + input: &[i8], + mut output: DisjointSlice>, + valid: usize, + shift: usize, + default: i8, + mode: u32, + ) { + exercise::<_, ITEMS_PER_THREAD>(input, &mut output, valid, shift, default, mode); + } + + #[kernel] + #[launch_contract(domain = 1, block = (128, 1, 1), dynamic_shared = inferred)] + pub fn copy_u8_128( + input: &[u8], + mut output: DisjointSlice>, + valid: usize, + shift: usize, + default: u8, + mode: u32, + ) { + exercise::<_, ITEMS_PER_THREAD>(input, &mut output, valid, shift, default, mode); + } + + #[kernel] + #[launch_contract(domain = 1, block = (128, 1, 1), dynamic_shared = inferred)] + pub fn copy_i16_128( + input: &[i16], + mut output: DisjointSlice>, + valid: usize, + shift: usize, + default: i16, + mode: u32, + ) { + exercise::<_, ITEMS_PER_THREAD>(input, &mut output, valid, shift, default, mode); + } + + #[kernel] + #[launch_contract(domain = 1, block = (128, 1, 1), dynamic_shared = inferred)] + pub fn copy_u16_128( + input: &[u16], + mut output: DisjointSlice>, + valid: usize, + shift: usize, + default: u16, + mode: u32, + ) { + exercise::<_, ITEMS_PER_THREAD>(input, &mut output, valid, shift, default, mode); + } + + #[kernel] + #[launch_contract(domain = 1, block = (128, 1, 1), dynamic_shared = inferred)] + pub fn copy_i32_128( + input: &[i32], + mut output: DisjointSlice>, + valid: usize, + shift: usize, + default: i32, + mode: u32, + ) { + exercise::<_, ITEMS_PER_THREAD>(input, &mut output, valid, shift, default, mode); + } + + #[kernel] + #[launch_contract(domain = 1, block = (128, 1, 1), dynamic_shared = inferred)] + pub fn copy_u32_128( + input: &[u32], + mut output: DisjointSlice>, + valid: usize, + shift: usize, + default: u32, + mode: u32, + ) { + exercise::<_, ITEMS_PER_THREAD>(input, &mut output, valid, shift, default, mode); + } + + #[kernel] + #[launch_contract(domain = 1, block = (128, 1, 1), dynamic_shared = inferred)] + pub fn copy_i64_128( + input: &[i64], + mut output: DisjointSlice>, + valid: usize, + shift: usize, + default: i64, + mode: u32, + ) { + exercise::<_, ITEMS_PER_THREAD>(input, &mut output, valid, shift, default, mode); + } + + #[kernel] + #[launch_contract(domain = 1, block = (128, 1, 1), dynamic_shared = inferred)] + pub fn copy_u64_128( + input: &[u64], + mut output: DisjointSlice>, + valid: usize, + shift: usize, + default: u64, + mode: u32, + ) { + exercise::<_, ITEMS_PER_THREAD>(input, &mut output, valid, shift, default, mode); + } + + #[kernel] + #[launch_contract(domain = 1, block = (128, 1, 1), dynamic_shared = inferred)] + pub fn copy_f32_128( + input: &[f32], + mut output: DisjointSlice>, + valid: usize, + shift: usize, + default: f32, + mode: u32, + ) { + exercise::<_, ITEMS_PER_THREAD>(input, &mut output, valid, shift, default, mode); + } + + #[kernel] + #[launch_contract(domain = 1, block = (128, 1, 1), dynamic_shared = inferred)] + pub fn copy_f64_128( + input: &[f64], + mut output: DisjointSlice>, + valid: usize, + shift: usize, + default: f64, + mode: u32, + ) { + exercise::<_, ITEMS_PER_THREAD>(input, &mut output, valid, shift, default, mode); + } + + #[kernel] + #[launch_contract(domain = 1, block = (64, 1, 1), dynamic_shared = inferred)] + pub fn copy_i32_64( + input: &[i32], + mut output: DisjointSlice>, + valid: usize, + shift: usize, + default: i32, + mode: u32, + ) { + exercise::<_, ITEMS_PER_THREAD>(input, &mut output, valid, shift, default, mode); + } + #[kernel] + #[launch_contract(domain = 1, block = (128, 1, 1), dynamic_shared = inferred)] + pub fn scalar(mut output: DisjointSlice>, value: i16) { + let block = this_block(); + let offset = block.rank_at::() as usize * block.count() as usize; + coop::store( + block, + &mut output, + value, + coop::StoreOptions::default().offset(offset), + ); + thread::sync_threads(); + let mut payload = ThreadData::<_, 1>::new(); + coop::load( + block, + &mut output, + &mut payload, + coop::LoadOptions::default().offset(offset), + ); + let scalar_preserved = payload.get(0) == Some(value); + thread::sync_threads(); + coop::load( + block, + &mut output, + &mut payload, + coop::LoadOptions::default().offset(offset).valid_items(1), + ); + assert!(scalar_preserved && payload.get(0).is_some() == (thread::threadIdx_x() == 0)); + } +} + +fn main() -> Result<(), Box> { + use cuda_async::simt::device_context::{init_device_contexts, with_cuda_context}; + use cuda_async::simt::device_operation::DeviceOperation; + use cuda_core::{DeviceBuffer, LaunchConfig1D}; + init_device_contexts(0, 1)?; + let context = with_cuda_context(0, |context| context.clone())?; + let stream = context.default_stream(); + // SAFETY: embedded kernels and these typed declarations are compiled together. + let module = unsafe { kernels::load_async(0)? }; + let trap = std::env::args().nth(1); + let mut cases = 0; + macro_rules! run { + ($t:ty, $prepare:ident, $launch:ident, $block:literal, $items:literal, $default:expr) => {{ + let default: $t = $default; + for groups in [1_u32, 3] { + let prepared = + module.$prepare::<$items>(LaunchConfig1D::auto_shared(groups, $block))?; + assert_eq!(prepared.__raw_config().shared_mem_bytes, 0); + assert!( + module + .$prepare::<$items>(LaunchConfig1D::auto_shared(groups, $block / 2)) + .is_err() + ); + for mode in [0, 1, 2, 3, 4, 6] { + for valid in [0, 1, $block * $items - 1, $block * $items] { + if (mode == 0) && valid != $block * $items { + continue; + } + let shift = if $items == 1 { 7 } else { 0 }; + let end = (groups as usize - 1) * $block * 4 + shift; + let source_len = end + + if mode <= 1 || mode == 6 { + $block * $items + } else { + valid + }; + let target_len = end + if mode <= 2 { $block * $items } else { valid }; + let input: Vec<$t> = (0..source_len).map(|i| (i % 61 + 1) as $t).collect(); + let source = DeviceBuffer::from_host(&stream, &input)?; + let target = DeviceBuffer::from_host(&stream, &vec![0 as $t; target_len])?; + let (_, target) = module + .$launch::<$items, _, _>( + &prepared, source, target, valid, shift, default, mode, + ) + .sync()?; + let mut result = vec![0 as $t; target_len]; + target.copy_to_host(&stream, &mut result)?; + stream.synchronize()?; + for (index, &value) in result.iter().enumerate() { + let within = index % ($block * 4); + let written = if mode == 0 || mode == 2 { + $block * $items + } else if mode == 6 { + valid / 2 + } else { + valid + }; + let expected = if within >= shift && within - shift < written { + if mode == 3 || (mode == 2 && within - shift >= valid) { + default + } else { + input[index] + } + } else { + 0 as $t + }; + assert_eq!( + value.to_ne_bytes(), + expected.to_ne_bytes(), + "type={} items={} mode={} valid={} index={}", + stringify!($t), + $items, + mode, + valid, + index + ); + } + cases += 1; + } + } + } + }}; + } + if let Some(trap) = trap { + let plan = module.prepare_copy_i32_128::<4>(LaunchConfig1D::auto_shared(1, 128))?; + let len = if trap == "ownership" { 1024 } else { 512 }; + let source = DeviceBuffer::from_host(&stream, &vec![11_i32; len])?; + let target = DeviceBuffer::from_host(&stream, &vec![0_i32; len])?; + let (valid, shift, mode) = match trap.as_str() { + "unwritten" => (512, 0, 9), + "missing-default-count" => (512, 0, 10), + "invalid-read" => (0, 0, 11), + "count" => (513, 0, 1), + "bounds" => (512, 1, 1), + "ownership" => (512, 1, 0), + "different-count" => (127, 0, 5), + _ => panic!("unknown trap mode"), + }; + let _ = module + .copy_i32_128_async_owned::<4, _, _>(&plan, source, target, valid, shift, -7, mode) + .sync()?; + panic!("expected trap returned successfully"); + } + run!( + i8, + prepare_copy_i8_128, + copy_i8_128_async_owned, + 128, + 1, + i8::MIN + ); + run!( + i8, + prepare_copy_i8_128, + copy_i8_128_async_owned, + 128, + 4, + i8::MIN + ); + run!( + u8, + prepare_copy_u8_128, + copy_u8_128_async_owned, + 128, + 1, + u8::MAX + ); + run!( + u8, + prepare_copy_u8_128, + copy_u8_128_async_owned, + 128, + 4, + u8::MAX + ); + run!( + i16, + prepare_copy_i16_128, + copy_i16_128_async_owned, + 128, + 1, + i16::MIN + ); + run!( + i16, + prepare_copy_i16_128, + copy_i16_128_async_owned, + 128, + 4, + i16::MIN + ); + run!( + u16, + prepare_copy_u16_128, + copy_u16_128_async_owned, + 128, + 1, + u16::MAX + ); + run!( + u16, + prepare_copy_u16_128, + copy_u16_128_async_owned, + 128, + 4, + u16::MAX + ); + run!( + i32, + prepare_copy_i32_128, + copy_i32_128_async_owned, + 128, + 1, + i32::MIN + ); + run!( + i32, + prepare_copy_i32_128, + copy_i32_128_async_owned, + 128, + 4, + i32::MIN + ); + run!( + u32, + prepare_copy_u32_128, + copy_u32_128_async_owned, + 128, + 1, + u32::MAX + ); + run!( + u32, + prepare_copy_u32_128, + copy_u32_128_async_owned, + 128, + 4, + u32::MAX + ); + run!( + i64, + prepare_copy_i64_128, + copy_i64_128_async_owned, + 128, + 1, + i64::MIN + 17 + ); + run!( + i64, + prepare_copy_i64_128, + copy_i64_128_async_owned, + 128, + 4, + i64::MIN + 17 + ); + run!( + u64, + prepare_copy_u64_128, + copy_u64_128_async_owned, + 128, + 1, + u64::MAX - 17 + ); + run!( + u64, + prepare_copy_u64_128, + copy_u64_128_async_owned, + 128, + 4, + u64::MAX - 17 + ); + run!( + f32, + prepare_copy_f32_128, + copy_f32_128_async_owned, + 128, + 1, + -0.0_f32 + ); + run!( + f32, + prepare_copy_f32_128, + copy_f32_128_async_owned, + 128, + 4, + -0.0_f32 + ); + run!( + f64, + prepare_copy_f64_128, + copy_f64_128_async_owned, + 128, + 1, + -0.0_f64 + ); + run!( + f64, + prepare_copy_f64_128, + copy_f64_128_async_owned, + 128, + 4, + -0.0_f64 + ); + run!( + i32, + prepare_copy_i32_64, + copy_i32_64_async_owned, + 64, + 1, + -1234567 + ); + run!( + i32, + prepare_copy_i32_64, + copy_i32_64_async_owned, + 64, + 4, + -1234567 + ); + let plan = module.prepare_scalar(LaunchConfig1D::auto_shared(3, 128))?; + let target = DeviceBuffer::from_host(&stream, &vec![0_i16; 384])?; + let target = module + .scalar_async_owned::<_>(&plan, target, -17_i16) + .sync()?; + let mut result = vec![0_i16; 384]; + target.copy_to_host(&stream, &mut result)?; + stream.synchronize()?; + assert!(result.iter().all(|value| *value == -17)); + cases += 1; + println!( + "PASS: {cases} public Direct cases, ten numeric types, scattered validity, short tails, owned async" + ); + Ok(()) +} From a092775dc1a8d963e1995cf4a0b68b6295a49539 Mon Sep 17 00:00:00 2001 From: Trent Nelson Date: Tue, 20 Oct 2026 11:13:12 -0700 Subject: [PATCH 24/33] Check enum controls with production MIR optimization Signed-off-by: Trent Nelson --- .../tests/check_carrier_rejections.py | 10 ++++++++++ .../tests/check_collective_rejections.py | 10 +++++++++- 2 files changed, 19 insertions(+), 1 deletion(-) diff --git a/rust/cuda-coop-codegen/tests/check_carrier_rejections.py b/rust/cuda-coop-codegen/tests/check_carrier_rejections.py index 4a982dc62222..3a6b64af8802 100644 --- a/rust/cuda-coop-codegen/tests/check_carrier_rejections.py +++ b/rust/cuda-coop-codegen/tests/check_carrier_rejections.py @@ -92,6 +92,14 @@ "carrier::load(selected, offset, valid, &mut values);", "mutation through a reference to a resource descriptor", ), + "enum_descriptor_mutation": ( + "let mut selected = Some(output);\n" + "replace_optional(&mut selected, other);\n" + "if let Some(mut target) = selected {\n" + " carrier::store(&mut target, offset, valid, &mut values);\n" + "}", + "mutation through a reference to a resource descriptor", + ), "descriptor_swap": ( "carrier::store(&mut other, offset, valid, &mut values);\n" "core::mem::swap(&mut output, &mut other);\n" @@ -235,6 +243,8 @@ #[inline(never)] fn replace<'a>(target: &mut Buffer<'a>, value: Buffer<'a>) { *target = value; } #[inline(never)] +fn replace_optional<'a>(target: &mut Option>, value: Buffer<'a>) { *target = Some(value); } +#[inline(never)] fn redirect<'a, 'b>(target: &mut (&'a mut Buffer<'b>,), value: &'a mut Buffer<'b>) { target.0 = value; } diff --git a/rust/cuda-coop-codegen/tests/check_collective_rejections.py b/rust/cuda-coop-codegen/tests/check_collective_rejections.py index bb64773b3081..c4fcf2e3ac2b 100644 --- a/rust/cuda-coop-codegen/tests/check_collective_rejections.py +++ b/rust/cuda-coop-codegen/tests/check_collective_rejections.py @@ -411,6 +411,14 @@ def in_helper(code): }, "not proved block-uniform", ), + "enum_projection": ( + in_helper( + " let choice = Some(lane as u32);\n" + " let raw = &choice as *const Option as *const (u32, u32);\n" + " let valid = unsafe { (*raw).1 } as i32;" + ), + "not proved block-uniform", + ), } ) unknown = set(args.case) - cases.keys() @@ -444,7 +452,7 @@ def in_helper(code): case_env = env.copy() if name == "unchecked_bounds": case_env["CUDA_OXIDE_UNCHECKED_INDEXING"] = "1" - if name.startswith("enum_") or name in [ + if name in [ "aggregate_reference", "projected_reborrow", "reborrow_uniform", From a9e6ef379ceac277aa256a31d1cf7d5500dcc1b8 Mon Sep 17 00:00:00 2001 From: Trent Nelson Date: Tue, 20 Oct 2026 11:15:01 -0700 Subject: [PATCH 25/33] Qualify Direct ownership and completion checks Signed-off-by: Trent Nelson --- docs/rust/load_store.rst | 11 ++++++++--- rust/cuda-coop/src/load_store/boundary.rs | 6 +++++- rust/cuda-coop/tests/check_load_store.py | 5 +++++ 3 files changed, 18 insertions(+), 4 deletions(-) diff --git a/docs/rust/load_store.rst b/docs/rust/load_store.rst index 8515808b6241..6de4481c36ac 100644 --- a/docs/rust/load_store.rst +++ b/docs/rust/load_store.rst @@ -86,6 +86,11 @@ collective ownership without giving each lane an overlapping mutable slice. The compiler's generated host launch consumes or borrows the actual device buffers until the operation completes. +For a grid of ``G`` blocks with ownership grain ``C``, block rank ``r`` owns +``[r*C, (r+1)*C)`` in each period of ``G*C`` elements. An offset may start inside +that region, and its nonempty selected range must stay inside the same region. +This also supports grid-stride traversal through later periods. + The descriptor returned by ``this_block()`` is a query value. Safe collective calls additionally require a supported Oxide build, an exact launch contract, uniform control flow, and matching group participation. Compilation diagnoses @@ -93,9 +98,9 @@ lane-varying controls or conflicting resource use. Runtime checks trap invalid bounds, ownership ranges, counts, or unwritten values selected for Store. Failures are not transactional across lanes. -To load from a carrier after storing to it, call -``cuda_device::thread::sync_threads()`` and pass ``&mut carrier`` as the Load -source. The compiler requires this completion before conflicting accesses, +Use ``cuda_device::thread::sync_threads()`` between conflicting accesses to a +carrier: Store followed by Load or Store, or Load followed by Store. Pass +``&mut carrier`` as the Load source. The compiler requires this completion, including reuse in loops. Ending a Rust borrow or dropping a value is not completion. The host example calls ``sync()`` before inspecting its output. diff --git a/rust/cuda-coop/src/load_store/boundary.rs b/rust/cuda-coop/src/load_store/boundary.rs index 9ae7f1e54ab1..4c0afdbc63ef 100644 --- a/rust/cuda-coop/src/load_store/boundary.rs +++ b/rust/cuda-coop/src/load_store/boundary.rs @@ -47,7 +47,11 @@ pub(super) fn containment(offset: usize, valid: usize) { assert!(groups > 0 && groups <= usize::MAX / capacity); let period = groups * capacity; if valid != 0 { - let within = offset % period; + let within = if offset < period { + offset + } else { + offset % period + }; let begin = rank * capacity; assert!(within >= begin && within - begin <= capacity); assert!(valid <= capacity - (within - begin)); diff --git a/rust/cuda-coop/tests/check_load_store.py b/rust/cuda-coop/tests/check_load_store.py index d29d8c0b397b..b881828fd738 100644 --- a/rust/cuda-coop/tests/check_load_store.py +++ b/rust/cuda-coop/tests/check_load_store.py @@ -73,6 +73,11 @@ ), "early_return": ("if thread::threadIdx_x()!=0 { return; }" + load, "lane-varying"), "missing_completion": (load + store + reload, "requires a full-block barrier"), + "write_after_read": ( + reload + + "coop::store(block,&mut output,1_i32,coop::StoreOptions::default().offset(offset));", + "requires a full-block barrier", + ), "loop_without_completion": ( load + "let mut iteration=0;while iteration<2 {" + store + "iteration+=1;}", "requires a full-block barrier", From fed23efcfdfe768ddc5ac90b4f32ed3fac38f0e3 Mon Sep 17 00:00:00 2001 From: Trent Nelson Date: Tue, 20 Oct 2026 11:23:46 -0700 Subject: [PATCH 26/33] Support typed Transpose payloads and validity in native CUB Signed-off-by: Trent Nelson --- rust/cuda-coop-codegen/src/lib.rs | 164 ++++++++++++++++++++----- rust/cuda-coop-codegen/src/oxide.rs | 79 +++++++----- rust/cuda-coop-codegen/tests/native.rs | 118 +++++++++++++----- 3 files changed, 267 insertions(+), 94 deletions(-) diff --git a/rust/cuda-coop-codegen/src/lib.rs b/rust/cuda-coop-codegen/src/lib.rs index 41b532965f6d..a3f4e62bf986 100644 --- a/rust/cuda-coop-codegen/src/lib.rs +++ b/rust/cuda-coop-codegen/src/lib.rs @@ -13,12 +13,12 @@ //! [`Request`] describes a one-dimensional complete block, positive fixed //! items per thread, and [`Algorithm::Direct`] or [`Algorithm::Transpose`]. //! The raw i32 recipes support Direct and Transpose. Their Store payloads must -//! be fully initialized. The validity-aware Direct recipes support all ten +//! be fully initialized. The validity-aware recipes support all ten //! [`Dtype`] variants with separate value storage and packed per-slot validity bits. //! Masked Load resets validity outside its prefix; Direct Store preserves both //! arrays and checks initialization only in its write domain. Transpose Store -//! replaces each lane's blocked payload with striped items across the block; -//! validity-aware Transpose recipes are not yet supported. +//! replaces each lane's blocked payload and validity with striped items across +//! the block, including when its write count is zero. //! //! [`Provider::from_checkout`] snapshots the selected CCCL and CUDA headers; //! [`Provider::compile`] uses those exact bytes and loaded NVRTC to emit an @@ -161,9 +161,11 @@ pub enum Operation { LoadMaskedValidity, /// Counted Load with an exact typed default; marks every slot initialized. LoadFilledValidity, - /// Full Store after checking every validity bit is set; preserves values and bits. + /// Full Store after checking every validity bit is set. + /// Transpose redistributes the payload; all resulting slots remain initialized. StoreFullValidity, - /// Counted Store after checking its write domain; preserves all values and flags. + /// Counted Store after checking its write domain. + /// Transpose redistributes all values and validity bits, even outside that domain. StorePartialValidity, } impl Operation { @@ -256,13 +258,13 @@ impl Request { /// Validate an exact typed block specialization. /// - /// The five validity-aware recipes support every [`Dtype`] with Direct. + /// The five validity-aware recipes support every [`Dtype`] with either algorithm. /// Legacy raw recipes retain their i32 ABI and algorithm support. /// /// # Errors /// /// Rejects the invalid shapes documented by [`Self::new`], a non-i32 raw - /// recipe, or a validity-aware Transpose recipe. + /// recipe. pub fn new_typed( dtype: Dtype, operation: Operation, @@ -275,11 +277,6 @@ impl Request { "legacy raw recipes require signed i32 elements", )); } - if operation.validity() && algorithm != Algorithm::Direct { - return Err(Error::message( - "validity-aware recipes currently require Direct", - )); - } if !(1..=1024).contains(&block_threads) || items_per_thread == 0 || u64::from(block_threads) * u64::from(items_per_thread) > i32::MAX as u64 @@ -349,18 +346,25 @@ pub struct Layout { /// so a null pointer is permitted for that argument only. All providers return /// unit; loads fill payloads in place and stores write memory. /// -/// Validity-aware Direct ABIs add a byte `flags` pointer immediately after +/// Validity-aware ABIs add a byte `flags` pointer immediately after /// scratch. `payload` addresses aligned storage for the requested number of /// elements; it need not contain initialized values. Flags occupy a separate, /// initialized array of `items_per_thread.div_ceil(8)` bytes. Bit `i % 8` of /// byte `i / 8` denotes whether slot `i` is initialized. Both arrays are private /// to the lane, disjoint from memory/scratch and one another. Load exclusively /// borrows both arrays; Direct Store receives const pointers and borrows them -/// read-only. Native code accesses payload bytes without evaluating invalid -/// values. Masked Load clears bits outside its prefix; filled/full Load write +/// read-only. Transpose Store exclusively borrows both arrays and returns their +/// permutation: destination `(lane, item)` receives source +/// `((block_threads * item + lane) / items_per_thread, +/// (block_threads * item + lane) % items_per_thread)`. +/// This applies to every slot, independently of the Store count. Repeated +/// Stores permute the current payload again. Native code accesses payload bytes +/// without evaluating invalid values. Masked Load clears bits outside its prefix; filled/full Load write /// every value before setting its bit. Store checks its entire lane-local -/// write domain before numeric conversion and preserves both arrays. Failure -/// traps; it does not promise rollback of stores by other lanes. +/// write domain before numeric conversion. Its counted Transpose form transports +/// arbitrary validity in a separate shared bitmap included in [`Self::scratch`]. +/// A full Store checks the complete tile and therefore needs no shared bitmap. +/// Failure traps; it does not promise rollback of stores by other lanes. #[derive(Debug)] pub struct Specialization { /// SHA-256 of request/source, target/options, captured headers and native toolchain inputs. @@ -370,12 +374,13 @@ pub struct Specialization { /// 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. + /// Complete shared region, including any validity plane and padding. + /// 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 raw Transpose stores; every slot remains initialized. + /// True for Transpose stores, including counted stores with zero writes. /// This describes Store rearrangement, not whether Load writes its payload. pub mutates_payload: bool, } @@ -450,9 +455,9 @@ impl Provider { /// 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. + /// Separate CUB and complete scratch layout name expressions are registered + /// before compilation; literal template arguments carry sizeof/alignof + /// without another compilation or query kernel. /// /// # Errors /// @@ -501,7 +506,7 @@ impl Provider { let mut batch_identity = inputs.clone(); let mut source = String::from(include_str!("block_io.cuh")); source.push_str("\ntemplate \nextern __device__ unsigned char layout_probe;\n"); - let mut expressions = Vec::with_capacity(requests.len()); + let mut expressions = Vec::with_capacity(requests.len() * 2); let mut specializations = Vec::with_capacity(requests.len()); for &request in requests { let mut hash = inputs.clone(); @@ -515,6 +520,7 @@ impl Provider { let symbol = format!("cuda_coop_{}_{}", request.operation.token(), suffix); source.push_str(&self::source(request, &symbol)); expressions.push(format!("&layout_probe")); + expressions.push(format!("&layout_probe")); specializations.push(Specialization { identity, symbol, @@ -534,10 +540,13 @@ impl Provider { let output = self .compiler .compile(&source, &options, &expressions, &self.headers)?; - for (specialization, layout) in specializations.iter_mut().zip(output.layouts) { - specialization.cub_layout = layout; + for (specialization, layouts) in specializations + .iter_mut() + .zip(output.layouts.chunks_exact(2)) + { + specialization.cub_layout = layouts[0]; if !specialization.scratch_free { - specialization.scratch = layout; + specialization.scratch = layouts[1]; } } Ok(Artifact { @@ -553,6 +562,7 @@ impl Provider { fn source(request: Request, symbol: &str) -> String { let op = request.operation; + let transpose = request.algorithm == Algorithm::Transpose; let (header, primitive, algorithm, qualifier, method) = if op.is_load() { ("load", "BlockLoad", "BLOCK_LOAD", "const ", "Load") } else { @@ -562,6 +572,21 @@ fn source(request: Request, symbol: &str) -> String { let items = request.items_per_thread; let bitmap_bytes = items.div_ceil(8); let tile = request.block_threads * items; + let mask_plane = op == Operation::StorePartialValidity && transpose && items > 1; + let storage = if mask_plane { + let plane_bytes = u64::from(request.block_threads.div_ceil(32)) * u64::from(items) * 4; + format!( + r#"struct scratch_storage {{ + alignas(alignof(primitive::TempStorage)) unsigned char cub[sizeof(primitive::TempStorage)]; + alignas(unsigned int) unsigned char masks[{plane_bytes}]; +}}; +static_assert(sizeof(unsigned int) == 4); +static_assert(offsetof(scratch_storage, masks) == sizeof(primitive::TempStorage)); +static_assert(alignof(scratch_storage) == alignof(primitive::TempStorage));"# + ) + } else { + "using scratch_storage = primitive::TempStorage;".to_owned() + }; let mut controls = String::new(); let mut operands = String::new(); let check = if op.counted() { @@ -588,7 +613,7 @@ fn source(request: Request, symbol: &str) -> String { } else { "memory" }; - let payload_const = if op.validity() && !op.is_load() { + let payload_const = if op.validity() && !op.is_load() && !transpose { "const " } else { "" @@ -622,23 +647,99 @@ fn source(request: Request, symbol: &str) -> String { }}"# ) } else { + // The general Store permutes all source slots, including initialized + // slots outside its write domain. Direct only needs its write domain. + let copy_domain = if transpose { "true" } else { &domain }; + let prepare = if mask_plane { + format!( + r#"unsigned char* const mask_plane = scratch + offsetof(scratch_storage, masks); + const int tid = int(threadIdx.x); + const int warp = tid / 32; + const int remaining = {threads} - warp * 32; + const unsigned int active = remaining >= 32 ? 0xffffffffu : (1u << remaining) - 1; +#pragma unroll + for (int i = 0; i < {items}; ++i) {{ + const unsigned int bits = __ballot_sync(active, (flags[i / 8] & (1u << (i % 8))) != 0); + if ((tid & 31) == 0) {{ + void* const address = mask_plane + sizeof(unsigned int) * (warp * {items} + i); + ::new (address) unsigned int(bits); + }} + }}"#, + threads = request.block_threads + ) + } else { + String::new() + }; + let restore = if transpose { + let validity = if op == Operation::StoreFullValidity { + // The full-domain preflight proves every permuted slot + // initialized; no collective mask exchange is necessary. + format!( + r#"#pragma unroll + for (int byte = 0; byte < {bitmap_bytes}; ++byte) {{ + unsigned int bits = 0; +#pragma unroll + for (int bit = 0; bit < 8; ++bit) {{ + if (byte * 8 + bit < {items}) bits |= 1u << bit; + }} + flags[byte] = static_cast(bits); + }}"# + ) + } else if mask_plane { + format!( + r#"// CUB's exchange barrier orders the initialized plane writes above. +#pragma unroll + for (int byte = 0; byte < {bitmap_bytes}; ++byte) {{ + unsigned int bits = 0; +#pragma unroll + for (int bit = 0; bit < 8; ++bit) {{ + const int i = byte * 8 + bit; + if (i < {items}) {{ + const int q = {threads} * i + tid; + const int source_thread = q / {items}; + const int source_item = q % {items}; + const auto* const address = reinterpret_cast( + mask_plane + sizeof(unsigned int) * ((source_thread / 32) * {items} + source_item)); + if (((*address >> (source_thread % 32)) & 1u) != 0) bits |= 1u << bit; + }} + }} + flags[byte] = static_cast(bits); + }}"#, + threads = request.block_threads + ) + } else { + // With one item per lane, BlockedToStriped is identity. + String::new() + }; + format!( + r#" +#pragma unroll + for (int i = 0; i < {items}; ++i) {{ + cuda::std::memcpy(bytes + sizeof(value_type) * i, values[i].bytes, sizeof(value_type)); + }} + {validity}"# + ) + } else { + String::new() + }; format!( r#"// Complete this lane's preflight before any numeric conversion or store. #pragma unroll for (int i = 0; i < {items}; ++i) {{ if (({domain}) && (flags[i / 8] & (1u << (i % 8))) == 0) __trap(); }} + {prepare} cell values[{items}]; #pragma unroll for (int i = 0; i < {items}; ++i) {{ - if ({domain}) cuda::std::memcpy(values[i].bytes, bytes + sizeof(value_type) * i, sizeof(value_type)); + if ({copy_domain}) cuda::std::memcpy(values[i].bytes, bytes + sizeof(value_type) * i, sizeof(value_type)); }} - operation.Store({memory}, values{operands});"# + operation.Store({memory}, values{operands});{restore}"# ) }; ( "cuda_coop_detail::opaque_cell", - if op.is_load() { + if op.is_load() || transpose { ", unsigned char* flags" } else { ", const unsigned char* flags" @@ -660,10 +761,13 @@ fn source(request: Request, symbol: &str) -> String { r#"// Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception #include +#include +#include namespace {symbol}_detail {{ using value_type = {dtype}; using cell = {cell}; using primitive = cub::{primitive}; +{storage} static_assert(sizeof(value_type) == {size} && alignof(value_type) == {size}); static_assert(sizeof(cell) == sizeof(value_type) && alignof(cell) == alignof(value_type)); extern "C" __device__ void {symbol}({qualifier}value_type* memory, {payload_const}value_type* payload, unsigned char* scratch{flags}{controls}) {{ diff --git a/rust/cuda-coop-codegen/src/oxide.rs b/rust/cuda-coop-codegen/src/oxide.rs index 1d2f10a04c2d..c12d5b7e34ce 100644 --- a/rust/cuda-coop-codegen/src/oxide.rs +++ b/rust/cuda-coop-codegen/src/oxide.rs @@ -10,7 +10,7 @@ //! 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 raw i32 and typed Direct validity recipes +//! The protocol supports the raw i32 and typed Direct/Transpose validity recipes //! documented by [`Operation`]. //! Device calls retain all participation, validity and scratch obligations of //! [`crate::Specialization`]. No allocation or synchronization policy is supplied. @@ -197,23 +197,25 @@ struct BatchResponse { fn arguments(request: Request) -> Vec { let operation = request.operation(); let dtype = Scalar::from(request.dtype()); + let shared_payload = + operation.validity() && !operation.is_load() && request.algorithm() == Algorithm::Direct; let mut arguments = vec![ if operation.is_load() { Argument::ConstPointer(dtype) } else { Argument::MutPointer(dtype) }, - if operation.validity() && !operation.is_load() { + if shared_payload { Argument::ConstPointer(dtype) } else { Argument::MutPointer(dtype) }, ]; if operation.validity() { - arguments.push(if operation.is_load() { - Argument::MutPointer(Scalar::Unsigned(8)) - } else { + arguments.push(if shared_payload { Argument::ConstPointer(Scalar::Unsigned(8)) + } else { + Argument::MutPointer(Scalar::Unsigned(8)) }); } if operation.counted() { @@ -487,41 +489,52 @@ mod tests { ] { let scalar = Scalar::from(dtype); assert_eq!(scalar.dtype().unwrap(), dtype); - let request = SpecializationRequest { - operation: "load_filled_validity_bits_v1".into(), - algorithm: "direct".into(), - dtype: scalar, - block: [128, 1, 1], - items_per_thread: 4, + for algorithm in ["direct", "transpose"] { + let request = SpecializationRequest { + operation: "load_filled_validity_bits_v1".into(), + algorithm: algorithm.into(), + dtype: scalar, + block: [128, 1, 1], + items_per_thread: 4, + } + .native() + .unwrap(); + assert_eq!( + arguments(request), + vec![ + Argument::ConstPointer(scalar), + Argument::MutPointer(scalar), + Argument::MutPointer(Scalar::Unsigned(8)), + Argument::Value(Scalar::Signed(32)), + Argument::Value(scalar), + ] + ); } - .native() - .unwrap(); - assert_eq!( - arguments(request), - vec![ - Argument::ConstPointer(scalar), - Argument::MutPointer(scalar), - Argument::MutPointer(Scalar::Unsigned(8)), - Argument::Value(Scalar::Signed(32)), - Argument::Value(scalar), - ] - ); } for operation in [ Operation::StoreFullValidity, Operation::StorePartialValidity, ] { - let request = - Request::new_typed(Dtype::I16, operation, Algorithm::Direct, 128, 4).unwrap(); - let mut expected = vec![ - Argument::MutPointer(Scalar::Signed(16)), - Argument::ConstPointer(Scalar::Signed(16)), - Argument::ConstPointer(Scalar::Unsigned(8)), - ]; - if operation.counted() { - expected.push(Argument::Value(Scalar::Signed(32))); + for algorithm in [Algorithm::Direct, Algorithm::Transpose] { + let request = Request::new_typed(Dtype::I16, operation, algorithm, 128, 4).unwrap(); + let mut expected = vec![ + Argument::MutPointer(Scalar::Signed(16)), + if algorithm == Algorithm::Direct { + Argument::ConstPointer(Scalar::Signed(16)) + } else { + Argument::MutPointer(Scalar::Signed(16)) + }, + if algorithm == Algorithm::Direct { + Argument::ConstPointer(Scalar::Unsigned(8)) + } else { + Argument::MutPointer(Scalar::Unsigned(8)) + }, + ]; + if operation.counted() { + expected.push(Argument::Value(Scalar::Signed(32))); + } + assert_eq!(arguments(request), expected); } - assert_eq!(arguments(request), expected); } for scalar in [Scalar::Signed(1), Scalar::Unsigned(128), Scalar::Float(16)] { assert!(scalar.dtype().is_err()); diff --git a/rust/cuda-coop-codegen/tests/native.rs b/rust/cuda-coop-codegen/tests/native.rs index e61d2c2609cf..08bde41edc5c 100644 --- a/rust/cuda-coop-codegen/tests/native.rs +++ b/rust/cuda-coop-codegen/tests/native.rs @@ -14,7 +14,6 @@ fn rejects_invalid_shapes() { assert!( Request::new_typed(Dtype::U64, Operation::LoadFull, Algorithm::Direct, 128, 4).is_err() ); - assert!(Request::new(Operation::LoadMaskedValidity, Algorithm::Transpose, 128, 4).is_err()); } #[test] @@ -37,48 +36,105 @@ fn native_typed_validity_and_determinism() -> Result<(), Box 1 + { + let plane = 4 + * request.block_threads().div_ceil(32) as usize + * request.items_per_thread() as usize; + let aligned_total = (entry.cub_layout.size + plane) + .div_ceil(entry.cub_layout.alignment) + * entry.cub_layout.alignment; + assert_eq!(entry.scratch.size, aligned_total); + assert!(entry.scratch.size > entry.cub_layout.size); + if request.block_threads() == 128 { + assert_eq!(entry.cub_layout.size, 2064); + assert_eq!(entry.scratch.size, 2128); + } + } else { + assert_eq!(entry.scratch, entry.cub_layout); } - ); - assert!(entry.scratch_free); - assert!(!entry.mutates_payload); + } } Ok(()) } From fb9b4c49fd40a8dc0e3dbc1aac44db73993049a8 Mon Sep 17 00:00:00 2001 From: Trent Nelson Date: Tue, 20 Oct 2026 11:38:23 -0700 Subject: [PATCH 27/33] Pin public Load and Store to the reviewed collective compiler Signed-off-by: Trent Nelson --- rust/Cargo.lock | 8 ++--- rust/cuda-coop/Cargo.toml | 2 +- rust/cuda-coop/examples/load_store/Cargo.lock | 30 +++++++++---------- rust/cuda-coop/examples/load_store/Cargo.toml | 8 ++--- rust/cuda-coop/tests/load_store/Cargo.lock | 30 +++++++++---------- rust/cuda-coop/tests/load_store/Cargo.toml | 8 ++--- 6 files changed, 43 insertions(+), 43 deletions(-) diff --git a/rust/Cargo.lock b/rust/Cargo.lock index e8b7b6d4af62..03bddcbec468 100644 --- a/rust/Cargo.lock +++ b/rust/Cargo.lock @@ -57,7 +57,7 @@ dependencies = [ [[package]] name = "cuda-device" version = "0.2.1" -source = "git+https://github.com/tpn/cuda-rust.git?rev=be42c33eeb15ea5639fdd7d793a76706f10cb19e#be42c33eeb15ea5639fdd7d793a76706f10cb19e" +source = "git+https://github.com/tpn/cuda-rust.git?rev=344d4c683ab4e4d9a8802aac636060d6b627a38b#344d4c683ab4e4d9a8802aac636060d6b627a38b" dependencies = [ "cuda-macros", ] @@ -65,7 +65,7 @@ dependencies = [ [[package]] name = "cuda-macros" version = "0.2.1" -source = "git+https://github.com/tpn/cuda-rust.git?rev=be42c33eeb15ea5639fdd7d793a76706f10cb19e#be42c33eeb15ea5639fdd7d793a76706f10cb19e" +source = "git+https://github.com/tpn/cuda-rust.git?rev=344d4c683ab4e4d9a8802aac636060d6b627a38b#344d4c683ab4e4d9a8802aac636060d6b627a38b" dependencies = [ "oxide-tile-export", "proc-macro2", @@ -127,7 +127,7 @@ checksum = "cf8baf1c55e62ffcace7a9f06f4bd9cd3f0c4beb022d3b367256b91b87513d98" [[package]] name = "oxide-tile-export" version = "0.2.1" -source = "git+https://github.com/tpn/cuda-rust.git?rev=be42c33eeb15ea5639fdd7d793a76706f10cb19e#be42c33eeb15ea5639fdd7d793a76706f10cb19e" +source = "git+https://github.com/tpn/cuda-rust.git?rev=344d4c683ab4e4d9a8802aac636060d6b627a38b#344d4c683ab4e4d9a8802aac636060d6b627a38b" dependencies = [ "serde", "serde_json", @@ -155,7 +155,7 @@ dependencies = [ [[package]] name = "reserved-oxide-symbols" version = "0.2.1" -source = "git+https://github.com/tpn/cuda-rust.git?rev=be42c33eeb15ea5639fdd7d793a76706f10cb19e#be42c33eeb15ea5639fdd7d793a76706f10cb19e" +source = "git+https://github.com/tpn/cuda-rust.git?rev=344d4c683ab4e4d9a8802aac636060d6b627a38b#344d4c683ab4e4d9a8802aac636060d6b627a38b" [[package]] name = "serde" diff --git a/rust/cuda-coop/Cargo.toml b/rust/cuda-coop/Cargo.toml index 0a903f2e1ed2..8bd4ee382a20 100644 --- a/rust/cuda-coop/Cargo.toml +++ b/rust/cuda-coop/Cargo.toml @@ -9,7 +9,7 @@ license.workspace = true publish = false [dependencies] -cuda-device = { git = "https://github.com/tpn/cuda-rust.git", rev = "be42c33eeb15ea5639fdd7d793a76706f10cb19e" } +cuda-device = { git = "https://github.com/tpn/cuda-rust.git", rev = "344d4c683ab4e4d9a8802aac636060d6b627a38b" } [build-dependencies] cuda-coop-codegen = { path = "../cuda-coop-codegen" } diff --git a/rust/cuda-coop/examples/load_store/Cargo.lock b/rust/cuda-coop/examples/load_store/Cargo.lock index 7e39af8b87c4..f7068217c0a7 100644 --- a/rust/cuda-coop/examples/load_store/Cargo.lock +++ b/rust/cuda-coop/examples/load_store/Cargo.lock @@ -143,7 +143,7 @@ dependencies = [ [[package]] name = "cuda-artifact-finalizer" version = "0.2.1" -source = "git+https://github.com/tpn/cuda-rust.git?rev=be42c33eeb15ea5639fdd7d793a76706f10cb19e#be42c33eeb15ea5639fdd7d793a76706f10cb19e" +source = "git+https://github.com/tpn/cuda-rust.git?rev=344d4c683ab4e4d9a8802aac636060d6b627a38b#344d4c683ab4e4d9a8802aac636060d6b627a38b" dependencies = [ "libnvvm-sys", "nvjitlink-sys", @@ -155,7 +155,7 @@ dependencies = [ [[package]] name = "cuda-async" version = "0.4.0" -source = "git+https://github.com/tpn/cuda-rust.git?rev=be42c33eeb15ea5639fdd7d793a76706f10cb19e#be42c33eeb15ea5639fdd7d793a76706f10cb19e" +source = "git+https://github.com/tpn/cuda-rust.git?rev=344d4c683ab4e4d9a8802aac636060d6b627a38b#344d4c683ab4e4d9a8802aac636060d6b627a38b" dependencies = [ "anyhow", "cuda-bindings", @@ -171,7 +171,7 @@ dependencies = [ [[package]] name = "cuda-bindings" version = "0.4.0" -source = "git+https://github.com/tpn/cuda-rust.git?rev=be42c33eeb15ea5639fdd7d793a76706f10cb19e#be42c33eeb15ea5639fdd7d793a76706f10cb19e" +source = "git+https://github.com/tpn/cuda-rust.git?rev=344d4c683ab4e4d9a8802aac636060d6b627a38b#344d4c683ab4e4d9a8802aac636060d6b627a38b" dependencies = [ "bindgen", "libloading", @@ -202,7 +202,7 @@ dependencies = [ [[package]] name = "cuda-core" version = "0.4.0" -source = "git+https://github.com/tpn/cuda-rust.git?rev=be42c33eeb15ea5639fdd7d793a76706f10cb19e#be42c33eeb15ea5639fdd7d793a76706f10cb19e" +source = "git+https://github.com/tpn/cuda-rust.git?rev=344d4c683ab4e4d9a8802aac636060d6b627a38b#344d4c683ab4e4d9a8802aac636060d6b627a38b" dependencies = [ "anyhow", "cuda-bindings", @@ -214,7 +214,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" version = "0.4.0" -source = "git+https://github.com/tpn/cuda-rust.git?rev=be42c33eeb15ea5639fdd7d793a76706f10cb19e#be42c33eeb15ea5639fdd7d793a76706f10cb19e" +source = "git+https://github.com/tpn/cuda-rust.git?rev=344d4c683ab4e4d9a8802aac636060d6b627a38b#344d4c683ab4e4d9a8802aac636060d6b627a38b" dependencies = [ "proc-macro2", "quote", @@ -224,7 +224,7 @@ dependencies = [ [[package]] name = "cuda-device" version = "0.2.1" -source = "git+https://github.com/tpn/cuda-rust.git?rev=be42c33eeb15ea5639fdd7d793a76706f10cb19e#be42c33eeb15ea5639fdd7d793a76706f10cb19e" +source = "git+https://github.com/tpn/cuda-rust.git?rev=344d4c683ab4e4d9a8802aac636060d6b627a38b#344d4c683ab4e4d9a8802aac636060d6b627a38b" dependencies = [ "cuda-macros", ] @@ -232,7 +232,7 @@ dependencies = [ [[package]] name = "cuda-host" version = "0.2.1" -source = "git+https://github.com/tpn/cuda-rust.git?rev=be42c33eeb15ea5639fdd7d793a76706f10cb19e#be42c33eeb15ea5639fdd7d793a76706f10cb19e" +source = "git+https://github.com/tpn/cuda-rust.git?rev=344d4c683ab4e4d9a8802aac636060d6b627a38b#344d4c683ab4e4d9a8802aac636060d6b627a38b" dependencies = [ "cuda-artifact-finalizer", "cuda-async", @@ -249,7 +249,7 @@ dependencies = [ [[package]] name = "cuda-macros" version = "0.2.1" -source = "git+https://github.com/tpn/cuda-rust.git?rev=be42c33eeb15ea5639fdd7d793a76706f10cb19e#be42c33eeb15ea5639fdd7d793a76706f10cb19e" +source = "git+https://github.com/tpn/cuda-rust.git?rev=344d4c683ab4e4d9a8802aac636060d6b627a38b#344d4c683ab4e4d9a8802aac636060d6b627a38b" dependencies = [ "oxide-tile-export", "proc-macro2", @@ -263,7 +263,7 @@ dependencies = [ [[package]] name = "cuda-target-spec" version = "0.1.0" -source = "git+https://github.com/tpn/cuda-rust.git?rev=be42c33eeb15ea5639fdd7d793a76706f10cb19e#be42c33eeb15ea5639fdd7d793a76706f10cb19e" +source = "git+https://github.com/tpn/cuda-rust.git?rev=344d4c683ab4e4d9a8802aac636060d6b627a38b#344d4c683ab4e4d9a8802aac636060d6b627a38b" [[package]] name = "dashmap" @@ -518,7 +518,7 @@ dependencies = [ [[package]] name = "libnvvm-sys" version = "0.2.1" -source = "git+https://github.com/tpn/cuda-rust.git?rev=be42c33eeb15ea5639fdd7d793a76706f10cb19e#be42c33eeb15ea5639fdd7d793a76706f10cb19e" +source = "git+https://github.com/tpn/cuda-rust.git?rev=344d4c683ab4e4d9a8802aac636060d6b627a38b#344d4c683ab4e4d9a8802aac636060d6b627a38b" dependencies = [ "cuda-target-spec", "libloading", @@ -571,7 +571,7 @@ dependencies = [ [[package]] name = "nvjitlink-sys" version = "0.2.1" -source = "git+https://github.com/tpn/cuda-rust.git?rev=be42c33eeb15ea5639fdd7d793a76706f10cb19e#be42c33eeb15ea5639fdd7d793a76706f10cb19e" +source = "git+https://github.com/tpn/cuda-rust.git?rev=344d4c683ab4e4d9a8802aac636060d6b627a38b#344d4c683ab4e4d9a8802aac636060d6b627a38b" dependencies = [ "libloading", "thiserror 2.0.21", @@ -598,7 +598,7 @@ checksum = "9f7c3e4beb33f85d45ae3e3a1792185706c8e16d043238c593331cc7cd313b50" [[package]] name = "oxide-artifacts" version = "0.2.1" -source = "git+https://github.com/tpn/cuda-rust.git?rev=be42c33eeb15ea5639fdd7d793a76706f10cb19e#be42c33eeb15ea5639fdd7d793a76706f10cb19e" +source = "git+https://github.com/tpn/cuda-rust.git?rev=344d4c683ab4e4d9a8802aac636060d6b627a38b#344d4c683ab4e4d9a8802aac636060d6b627a38b" dependencies = [ "object", ] @@ -606,7 +606,7 @@ dependencies = [ [[package]] name = "oxide-tile-export" version = "0.2.1" -source = "git+https://github.com/tpn/cuda-rust.git?rev=be42c33eeb15ea5639fdd7d793a76706f10cb19e#be42c33eeb15ea5639fdd7d793a76706f10cb19e" +source = "git+https://github.com/tpn/cuda-rust.git?rev=344d4c683ab4e4d9a8802aac636060d6b627a38b#344d4c683ab4e4d9a8802aac636060d6b627a38b" dependencies = [ "serde", "serde_json", @@ -654,7 +654,7 @@ dependencies = [ [[package]] name = "ptx-parse" version = "0.1.0" -source = "git+https://github.com/tpn/cuda-rust.git?rev=be42c33eeb15ea5639fdd7d793a76706f10cb19e#be42c33eeb15ea5639fdd7d793a76706f10cb19e" +source = "git+https://github.com/tpn/cuda-rust.git?rev=344d4c683ab4e4d9a8802aac636060d6b627a38b#344d4c683ab4e4d9a8802aac636060d6b627a38b" [[package]] name = "quote" @@ -706,7 +706,7 @@ checksum = "d6f6ff9a378485b298a5286656da665ba74413d36db0979633275d2e708145d4" [[package]] name = "reserved-oxide-symbols" version = "0.2.1" -source = "git+https://github.com/tpn/cuda-rust.git?rev=be42c33eeb15ea5639fdd7d793a76706f10cb19e#be42c33eeb15ea5639fdd7d793a76706f10cb19e" +source = "git+https://github.com/tpn/cuda-rust.git?rev=344d4c683ab4e4d9a8802aac636060d6b627a38b#344d4c683ab4e4d9a8802aac636060d6b627a38b" [[package]] name = "rustc-hash" diff --git a/rust/cuda-coop/examples/load_store/Cargo.toml b/rust/cuda-coop/examples/load_store/Cargo.toml index 0cf131efc562..442eaf871c8b 100644 --- a/rust/cuda-coop/examples/load_store/Cargo.toml +++ b/rust/cuda-coop/examples/load_store/Cargo.toml @@ -13,7 +13,7 @@ members = ["helper"] [dependencies] cuda-coop = { path = "../.." } coop-load-store-helper = { path = "helper" } -cuda-device = { git = "https://github.com/tpn/cuda-rust.git", rev = "be42c33eeb15ea5639fdd7d793a76706f10cb19e" } -cuda-host = { git = "https://github.com/tpn/cuda-rust.git", rev = "be42c33eeb15ea5639fdd7d793a76706f10cb19e", features = ["async"] } -cuda-core = { git = "https://github.com/tpn/cuda-rust.git", rev = "be42c33eeb15ea5639fdd7d793a76706f10cb19e" } -cuda-async = { git = "https://github.com/tpn/cuda-rust.git", rev = "be42c33eeb15ea5639fdd7d793a76706f10cb19e" } +cuda-device = { git = "https://github.com/tpn/cuda-rust.git", rev = "344d4c683ab4e4d9a8802aac636060d6b627a38b" } +cuda-host = { git = "https://github.com/tpn/cuda-rust.git", rev = "344d4c683ab4e4d9a8802aac636060d6b627a38b", features = ["async"] } +cuda-core = { git = "https://github.com/tpn/cuda-rust.git", rev = "344d4c683ab4e4d9a8802aac636060d6b627a38b" } +cuda-async = { git = "https://github.com/tpn/cuda-rust.git", rev = "344d4c683ab4e4d9a8802aac636060d6b627a38b" } diff --git a/rust/cuda-coop/tests/load_store/Cargo.lock b/rust/cuda-coop/tests/load_store/Cargo.lock index 133f595cdf96..5bcb0ea6a68a 100644 --- a/rust/cuda-coop/tests/load_store/Cargo.lock +++ b/rust/cuda-coop/tests/load_store/Cargo.lock @@ -135,7 +135,7 @@ dependencies = [ [[package]] name = "cuda-artifact-finalizer" version = "0.2.1" -source = "git+https://github.com/tpn/cuda-rust.git?rev=be42c33eeb15ea5639fdd7d793a76706f10cb19e#be42c33eeb15ea5639fdd7d793a76706f10cb19e" +source = "git+https://github.com/tpn/cuda-rust.git?rev=344d4c683ab4e4d9a8802aac636060d6b627a38b#344d4c683ab4e4d9a8802aac636060d6b627a38b" dependencies = [ "libnvvm-sys", "nvjitlink-sys", @@ -147,7 +147,7 @@ dependencies = [ [[package]] name = "cuda-async" version = "0.4.0" -source = "git+https://github.com/tpn/cuda-rust.git?rev=be42c33eeb15ea5639fdd7d793a76706f10cb19e#be42c33eeb15ea5639fdd7d793a76706f10cb19e" +source = "git+https://github.com/tpn/cuda-rust.git?rev=344d4c683ab4e4d9a8802aac636060d6b627a38b#344d4c683ab4e4d9a8802aac636060d6b627a38b" dependencies = [ "anyhow", "cuda-bindings", @@ -163,7 +163,7 @@ dependencies = [ [[package]] name = "cuda-bindings" version = "0.4.0" -source = "git+https://github.com/tpn/cuda-rust.git?rev=be42c33eeb15ea5639fdd7d793a76706f10cb19e#be42c33eeb15ea5639fdd7d793a76706f10cb19e" +source = "git+https://github.com/tpn/cuda-rust.git?rev=344d4c683ab4e4d9a8802aac636060d6b627a38b#344d4c683ab4e4d9a8802aac636060d6b627a38b" dependencies = [ "bindgen", "libloading", @@ -194,7 +194,7 @@ dependencies = [ [[package]] name = "cuda-core" version = "0.4.0" -source = "git+https://github.com/tpn/cuda-rust.git?rev=be42c33eeb15ea5639fdd7d793a76706f10cb19e#be42c33eeb15ea5639fdd7d793a76706f10cb19e" +source = "git+https://github.com/tpn/cuda-rust.git?rev=344d4c683ab4e4d9a8802aac636060d6b627a38b#344d4c683ab4e4d9a8802aac636060d6b627a38b" dependencies = [ "anyhow", "cuda-bindings", @@ -206,7 +206,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" version = "0.4.0" -source = "git+https://github.com/tpn/cuda-rust.git?rev=be42c33eeb15ea5639fdd7d793a76706f10cb19e#be42c33eeb15ea5639fdd7d793a76706f10cb19e" +source = "git+https://github.com/tpn/cuda-rust.git?rev=344d4c683ab4e4d9a8802aac636060d6b627a38b#344d4c683ab4e4d9a8802aac636060d6b627a38b" dependencies = [ "proc-macro2", "quote", @@ -216,7 +216,7 @@ dependencies = [ [[package]] name = "cuda-device" version = "0.2.1" -source = "git+https://github.com/tpn/cuda-rust.git?rev=be42c33eeb15ea5639fdd7d793a76706f10cb19e#be42c33eeb15ea5639fdd7d793a76706f10cb19e" +source = "git+https://github.com/tpn/cuda-rust.git?rev=344d4c683ab4e4d9a8802aac636060d6b627a38b#344d4c683ab4e4d9a8802aac636060d6b627a38b" dependencies = [ "cuda-macros", ] @@ -224,7 +224,7 @@ dependencies = [ [[package]] name = "cuda-host" version = "0.2.1" -source = "git+https://github.com/tpn/cuda-rust.git?rev=be42c33eeb15ea5639fdd7d793a76706f10cb19e#be42c33eeb15ea5639fdd7d793a76706f10cb19e" +source = "git+https://github.com/tpn/cuda-rust.git?rev=344d4c683ab4e4d9a8802aac636060d6b627a38b#344d4c683ab4e4d9a8802aac636060d6b627a38b" dependencies = [ "cuda-artifact-finalizer", "cuda-async", @@ -241,7 +241,7 @@ dependencies = [ [[package]] name = "cuda-macros" version = "0.2.1" -source = "git+https://github.com/tpn/cuda-rust.git?rev=be42c33eeb15ea5639fdd7d793a76706f10cb19e#be42c33eeb15ea5639fdd7d793a76706f10cb19e" +source = "git+https://github.com/tpn/cuda-rust.git?rev=344d4c683ab4e4d9a8802aac636060d6b627a38b#344d4c683ab4e4d9a8802aac636060d6b627a38b" dependencies = [ "oxide-tile-export", "proc-macro2", @@ -255,7 +255,7 @@ dependencies = [ [[package]] name = "cuda-target-spec" version = "0.1.0" -source = "git+https://github.com/tpn/cuda-rust.git?rev=be42c33eeb15ea5639fdd7d793a76706f10cb19e#be42c33eeb15ea5639fdd7d793a76706f10cb19e" +source = "git+https://github.com/tpn/cuda-rust.git?rev=344d4c683ab4e4d9a8802aac636060d6b627a38b#344d4c683ab4e4d9a8802aac636060d6b627a38b" [[package]] name = "dashmap" @@ -510,7 +510,7 @@ dependencies = [ [[package]] name = "libnvvm-sys" version = "0.2.1" -source = "git+https://github.com/tpn/cuda-rust.git?rev=be42c33eeb15ea5639fdd7d793a76706f10cb19e#be42c33eeb15ea5639fdd7d793a76706f10cb19e" +source = "git+https://github.com/tpn/cuda-rust.git?rev=344d4c683ab4e4d9a8802aac636060d6b627a38b#344d4c683ab4e4d9a8802aac636060d6b627a38b" dependencies = [ "cuda-target-spec", "libloading", @@ -563,7 +563,7 @@ dependencies = [ [[package]] name = "nvjitlink-sys" version = "0.2.1" -source = "git+https://github.com/tpn/cuda-rust.git?rev=be42c33eeb15ea5639fdd7d793a76706f10cb19e#be42c33eeb15ea5639fdd7d793a76706f10cb19e" +source = "git+https://github.com/tpn/cuda-rust.git?rev=344d4c683ab4e4d9a8802aac636060d6b627a38b#344d4c683ab4e4d9a8802aac636060d6b627a38b" dependencies = [ "libloading", "thiserror 2.0.21", @@ -590,7 +590,7 @@ checksum = "9f7c3e4beb33f85d45ae3e3a1792185706c8e16d043238c593331cc7cd313b50" [[package]] name = "oxide-artifacts" version = "0.2.1" -source = "git+https://github.com/tpn/cuda-rust.git?rev=be42c33eeb15ea5639fdd7d793a76706f10cb19e#be42c33eeb15ea5639fdd7d793a76706f10cb19e" +source = "git+https://github.com/tpn/cuda-rust.git?rev=344d4c683ab4e4d9a8802aac636060d6b627a38b#344d4c683ab4e4d9a8802aac636060d6b627a38b" dependencies = [ "object", ] @@ -598,7 +598,7 @@ dependencies = [ [[package]] name = "oxide-tile-export" version = "0.2.1" -source = "git+https://github.com/tpn/cuda-rust.git?rev=be42c33eeb15ea5639fdd7d793a76706f10cb19e#be42c33eeb15ea5639fdd7d793a76706f10cb19e" +source = "git+https://github.com/tpn/cuda-rust.git?rev=344d4c683ab4e4d9a8802aac636060d6b627a38b#344d4c683ab4e4d9a8802aac636060d6b627a38b" dependencies = [ "serde", "serde_json", @@ -646,7 +646,7 @@ dependencies = [ [[package]] name = "ptx-parse" version = "0.1.0" -source = "git+https://github.com/tpn/cuda-rust.git?rev=be42c33eeb15ea5639fdd7d793a76706f10cb19e#be42c33eeb15ea5639fdd7d793a76706f10cb19e" +source = "git+https://github.com/tpn/cuda-rust.git?rev=344d4c683ab4e4d9a8802aac636060d6b627a38b#344d4c683ab4e4d9a8802aac636060d6b627a38b" [[package]] name = "quote" @@ -698,7 +698,7 @@ checksum = "d6f6ff9a378485b298a5286656da665ba74413d36db0979633275d2e708145d4" [[package]] name = "reserved-oxide-symbols" version = "0.2.1" -source = "git+https://github.com/tpn/cuda-rust.git?rev=be42c33eeb15ea5639fdd7d793a76706f10cb19e#be42c33eeb15ea5639fdd7d793a76706f10cb19e" +source = "git+https://github.com/tpn/cuda-rust.git?rev=344d4c683ab4e4d9a8802aac636060d6b627a38b#344d4c683ab4e4d9a8802aac636060d6b627a38b" [[package]] name = "rustc-hash" diff --git a/rust/cuda-coop/tests/load_store/Cargo.toml b/rust/cuda-coop/tests/load_store/Cargo.toml index 6dbdea6cee62..235f9d76fdc6 100644 --- a/rust/cuda-coop/tests/load_store/Cargo.toml +++ b/rust/cuda-coop/tests/load_store/Cargo.toml @@ -11,7 +11,7 @@ publish = false [dependencies] cuda-coop = { path = "../.." } -cuda-device = { git = "https://github.com/tpn/cuda-rust.git", rev = "be42c33eeb15ea5639fdd7d793a76706f10cb19e" } -cuda-host = { git = "https://github.com/tpn/cuda-rust.git", rev = "be42c33eeb15ea5639fdd7d793a76706f10cb19e", features = ["async"] } -cuda-core = { git = "https://github.com/tpn/cuda-rust.git", rev = "be42c33eeb15ea5639fdd7d793a76706f10cb19e" } -cuda-async = { git = "https://github.com/tpn/cuda-rust.git", rev = "be42c33eeb15ea5639fdd7d793a76706f10cb19e" } +cuda-device = { git = "https://github.com/tpn/cuda-rust.git", rev = "344d4c683ab4e4d9a8802aac636060d6b627a38b" } +cuda-host = { git = "https://github.com/tpn/cuda-rust.git", rev = "344d4c683ab4e4d9a8802aac636060d6b627a38b", features = ["async"] } +cuda-core = { git = "https://github.com/tpn/cuda-rust.git", rev = "344d4c683ab4e4d9a8802aac636060d6b627a38b" } +cuda-async = { git = "https://github.com/tpn/cuda-rust.git", rev = "344d4c683ab4e4d9a8802aac636060d6b627a38b" } From 512456ad5e9cd1111d44e014614723980259c443 Mon Sep 17 00:00:00 2001 From: Trent Nelson Date: Tue, 20 Oct 2026 11:51:30 -0700 Subject: [PATCH 28/33] Add typed Transpose Load and Store with managed storage Signed-off-by: Trent Nelson --- docs/rust/load_store.rst | 51 +- rust/cuda-coop/build.rs | 23 +- .../examples/load_store/helper/src/lib.rs | 29 +- .../cuda-coop/examples/load_store/src/main.rs | 37 +- rust/cuda-coop/src/lib.rs | 11 +- rust/cuda-coop/src/load_store.rs | 171 +++-- rust/cuda-coop/src/load_store/boundary.rs | 117 ++- rust/cuda-coop/tests/check_load_store.py | 49 +- rust/cuda-coop/tests/load_store/src/main.rs | 11 + .../tests/load_store/src/transpose.rs | 710 ++++++++++++++++++ 10 files changed, 1110 insertions(+), 99 deletions(-) create mode 100644 rust/cuda-coop/tests/load_store/src/transpose.rs diff --git a/docs/rust/load_store.rst b/docs/rust/load_store.rst index 6de4481c36ac..b0264afe105e 100644 --- a/docs/rust/load_store.rst +++ b/docs/rust/load_store.rst @@ -1,15 +1,15 @@ .. Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. .. SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception -Direct Load and Store -===================== +Block Load and Store +==================== ``cuda_coop::load`` fills an existing ``ThreadData`` and returns unit. ``cuda_coop::store`` writes its values and returns unit. Both take the group first. This initial implementation supports complete, one-dimensional blocks -and the Direct algorithm: each thread owns consecutive elements in blocked -order. Other algorithms, warp groups, and public temporary storage remain -under development. +with Direct and Transpose algorithms. Both Load algorithms produce blocked +payloads, where each thread owns consecutive elements. Other algorithms, warp +groups, and public temporary storage remain under development. Use ``use cuda_coop as coop;`` and keep the per-thread item count a const generic named ``ITEMS_PER_THREAD``. The payload element type is inferred from its source. @@ -70,10 +70,36 @@ missing tail. Partial Store writes only its selected prefix and requires only those slots to be initialized. Direct Store accepts ``&ThreadData``, a shared reborrow of ``&mut ThreadData``, -or a scalar numeric value. A scalar contributes one item per thread. Store -preserves every payload value and validity flag, including initialized slots -outside its write domain. Algorithms that mutate payloads will require mutable -borrows when those algorithms are supported. +or a scalar numeric value. A scalar contributes one item per thread. Direct Store preserves every payload value and validity flag, including +initialized slots outside its write domain. + +Selecting Transpose +------------------- + +Select the algorithm through the trailing options. Selection is a Rust type, +so each call compiles one algorithm; ordinary defaults continue to select Direct. + +.. code-block:: rust + + coop::load(block, input, &mut values, + coop::LoadOptions::default().algorithm(coop::Transpose)); + coop::store(block, output, &mut values, + coop::StoreOptions::default().valid_items(count).algorithm(coop::Transpose)); + +Transpose uses CUB's coalesced accesses and shared block exchange. Its Store +requires a mutable payload borrow and permutes every value and initialized bit +from blocked into striped order. For a block of ``B`` threads and ``I`` items +per thread, destination lane/item ``(t, i)`` receives the old blocked slot at +linear index ``B*i+t``. This mutation occurs even for a zero-count Store; a +partial count only controls the global writes. Initialized slots outside the +written prefix remain initialized at their permuted positions. Scalar Store +uses a one-item temporary and accepts either algorithm. + +The compiler allocates the provider's complete shared storage, including the +validity plane when needed, and inserts completion before reusing managed +scratch. This does not add an unconditional trailing barrier. ``auto_shared`` +uses the inferred total for the kernel. Explicit temporary storage controls +remain under development. Collective ownership and completion ----------------------------------- @@ -100,8 +126,9 @@ Failures are not transactional across lanes. Use ``cuda_device::thread::sync_threads()`` between conflicting accesses to a carrier: Store followed by Load or Store, or Load followed by Store. Pass -``&mut carrier`` as the Load source. The compiler requires this completion, -including reuse in loops. Ending a Rust borrow or dropping a value is not +``&mut carrier`` as the Load source. The compiler requires completion, including reuse in loops. A genuine planned +scratch-reuse barrier can complete earlier accesses; code should not rely on an +algorithm merely being named Transpose to provide a trailing barrier. Ending a Rust borrow or dropping a value is not completion. The host example calls ``sync()`` before inspecting its output. Host development @@ -124,7 +151,7 @@ replacing them with unchanged or older timestamps may require an explicit rebuild. Crate publication is not supported yet. The rustdoc entries for ``load``, ``store``, ``LoadOptions``, ``StoreOptions``, -``LoadSource`` and ``StoreValue`` specify their full generic signatures and +``LoadSource``, ``StoreValue``, ``Algorithm``, ``Direct`` and ``Transpose`` specify their full generic signatures and supported argument forms. Native correctness and artifact checks are separate from host compilation. Performance qualification for masked payload access is ongoing; the current implementation makes no negligible-overhead claim. diff --git a/rust/cuda-coop/build.rs b/rust/cuda-coop/build.rs index 858fae4c0249..64adc8bb40e8 100644 --- a/rust/cuda-coop/build.rs +++ b/rust/cuda-coop/build.rs @@ -1,8 +1,27 @@ // Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception -//! Register the owning host provider for native Direct operations. +//! Register the owning host provider for native block Load and Store. +use cuda_coop_codegen::{Algorithm, Operation, oxide}; fn main() { - cuda_coop_codegen::oxide::build_script().expect("cuda.coop native provider registration"); + oxide::build_script().expect("cuda.coop native provider registration"); + if matches!( + std::env::args().nth(1).as_deref(), + Some("--oxide-library-request-v1" | "--oxide-library-inputs-v1") + ) { + return; + } + let executable = std::env::current_exe().expect("cuda.coop provider executable"); + for (operation, name) in [ + (Operation::LoadFullValidity, "LOAD_FULL"), + (Operation::LoadMaskedValidity, "LOAD_MASKED"), + (Operation::LoadFilledValidity, "LOAD_FILLED"), + (Operation::StoreFullValidity, "STORE_FULL"), + (Operation::StorePartialValidity, "STORE_PARTIAL"), + ] { + let value = oxide::registration(&executable, operation, Algorithm::Transpose) + .expect("cuda.coop Transpose registration"); + println!("cargo::rustc-env=COOP_TRANSPOSE_VALIDITY_{name}={value}"); + } } diff --git a/rust/cuda-coop/examples/load_store/helper/src/lib.rs b/rust/cuda-coop/examples/load_store/helper/src/lib.rs index 986ac5ec32bc..b587972831d8 100644 --- a/rust/cuda-coop/examples/load_store/helper/src/lib.rs +++ b/rust/cuda-coop/examples/load_store/helper/src/lib.rs @@ -13,6 +13,27 @@ pub fn copy( offset: usize, valid: usize, default: T, +) { + copy_with_algorithm::( + block, + input, + output, + offset, + valid, + default, + coop::Direct, + ); +} + +/// Copy a partial tile using a compile-time algorithm chosen by the caller. +pub fn copy_with_algorithm( + block: coop::ThreadGroup, + input: &[T], + output: &mut coop::DisjointSlice<'_, T, coop::CollectiveTiles>, + offset: usize, + valid: usize, + default: T, + algorithm: A, ) { let mut values = coop::ThreadData::<_, ITEMS_PER_THREAD>::new(); coop::load( @@ -22,14 +43,16 @@ pub fn copy( coop::LoadOptions::default() .valid_items(valid) .oob_default(default) - .offset(offset), + .offset(offset) + .algorithm(algorithm), ); coop::store( block, output, - &values, + &mut values, coop::StoreOptions::default() .valid_items(valid) - .offset(offset), + .offset(offset) + .algorithm(algorithm), ); } diff --git a/rust/cuda-coop/examples/load_store/src/main.rs b/rust/cuda-coop/examples/load_store/src/main.rs index 3f534fcc3ccb..0f41cc7fdfa7 100644 --- a/rust/cuda-coop/examples/load_store/src/main.rs +++ b/rust/cuda-coop/examples/load_store/src/main.rs @@ -33,6 +33,31 @@ mod kernels { -7_i16, ); } + #[kernel] + #[launch_contract(domain = 1, block = (128, 1, 1), dynamic_shared = inferred)] + pub fn transpose( + input: &[i16], + mut output: coop::DisjointSlice>, + ) { + let block = coop::this_block(); + let tile = block.count() as usize * ITEMS_PER_THREAD; + let offset = block.rank_at::() as usize * tile; + let remaining = if offset <= input.len() { + input.len() - offset + } else { + 0 + }; + let valid = remaining.min(tile); + coop_load_store_helper::copy_with_algorithm::<_, ITEMS_PER_THREAD, _>( + block, + input, + &mut output, + offset, + valid, + -7_i16, + coop::Transpose, + ); + } } fn main() -> Result<(), Box> { @@ -46,15 +71,15 @@ fn main() -> Result<(), Box> { // SAFETY: these typed declarations and embedded kernels compile together. let module = unsafe { kernels::load_async(0)? }; macro_rules! run { - ($items:literal) => {{ + ($items:literal, $prepare:ident, $launch:ident) => {{ let input: Vec = (0..3 * 128 * $items - 7) .map(|i| (i % 251) as i16) .collect(); let source = DeviceBuffer::from_host(&stream, &input)?; let target = DeviceBuffer::from_host(&stream, &vec![-1_i16; input.len()])?; - let plan = module.prepare_copy::<$items>(LaunchConfig1D::auto_shared(3, 128))?; + let plan = module.$prepare::<$items>(LaunchConfig1D::auto_shared(3, 128))?; let (_, target) = module - .copy_async_owned::<$items, _, _>(&plan, source, target) + .$launch::<$items, _, _>(&plan, source, target) .sync()?; let mut output = vec![0_i16; input.len()]; target.copy_to_host(&stream, &mut output)?; @@ -63,7 +88,9 @@ fn main() -> Result<(), Box> { println!("items/thread={}: {} values copied", $items, output.len()); }}; } - run!(1); - run!(4); + run!(1, prepare_copy, copy_async_owned); + run!(4, prepare_copy, copy_async_owned); + run!(1, prepare_transpose, transpose_async_owned); + run!(4, prepare_transpose, transpose_async_owned); Ok(()) } diff --git a/rust/cuda-coop/src/lib.rs b/rust/cuda-coop/src/lib.rs index a73e14c85920..c7536f88cf7e 100644 --- a/rust/cuda-coop/src/lib.rs +++ b/rust/cuda-coop/src/lib.rs @@ -29,9 +29,10 @@ //! Device queries use Oxide and the pinned `nightly-2026-10-11` toolchain. //! [`this_block`] and the other dimension-free group constructors describe //! the executing thread's current launch. Queries grant no collective -//! participation or synchronization guarantee. This crate supplies no -//! scratch storage. Direct [`load`] and [`store`] use complete blocks with -//! compiler-checked participation and collective memory ownership. Host payload tests need no GPU; +//! participation or synchronization guarantee. [`load`] and [`store`] support +//! Direct and Transpose on complete blocks with compiler-checked participation +//! and collective memory ownership. Transpose uses compiler-planned shared +//! scratch. Host payload tests need no GPU; //! executing hardware queries requires an Oxide-compiled CUDA kernel. #![no_std] @@ -43,7 +44,9 @@ mod thread_data; pub use alignment::{Alignment, SupportedAlignment}; pub use cuda_device::{CollectiveTiles, DisjointSlice}; -pub use load_store::{LoadOptions, LoadSource, StoreOptions, StoreValue, load, store}; +pub use load_store::{ + Algorithm, Direct, LoadOptions, LoadSource, StoreOptions, StoreValue, Transpose, load, store, +}; pub use thread_data::{Numeric, ThreadData, ThreadDataLike}; pub use group::{ diff --git a/rust/cuda-coop/src/load_store.rs b/rust/cuda-coop/src/load_store.rs index b73135d3c4ff..094c83b7356b 100644 --- a/rust/cuda-coop/src/load_store.rs +++ b/rust/cuda-coop/src/load_store.rs @@ -2,11 +2,34 @@ // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception use crate::{Alignment, Block, Numeric, SupportedAlignment, ThreadData, ThreadGroup}; +use core::marker::PhantomData; use cuda_device::{CollectiveTiles, DisjointSlice}; mod boundary; -/// Controls for a Direct [`load`]. +/// Direct blocked memory access without shared scratch. +#[derive(Clone, Copy, Debug, Default)] +pub struct Direct; + +/// Coalesced access through CUB's shared-memory block exchange. +/// +/// Load produces blocked payloads. Store permutes payload values and validity +/// into striped order, including slots outside its selected write prefix. +#[derive(Clone, Copy, Debug, Default)] +pub struct Transpose; + +/// A sealed, compile-time algorithm selected through Load/Store options. +pub trait Algorithm: sealed::Algorithm + Copy {} +impl Algorithm for Direct {} +impl Algorithm for Transpose {} +impl sealed::Algorithm for Direct { + const TRANSPOSE: bool = false; +} +impl sealed::Algorithm for Transpose { + const TRANSPOSE: bool = true; +} + +/// Controls for a [`load`], using [`Direct`] by default. /// /// `T` is inferred from the source and output. Defaults load the entire group /// tile, leave no unspecified slots, and use element offset zero. A partial @@ -16,15 +39,17 @@ mod boundary; /// Controls, including whether a count/default is present, must be uniform /// across the block. They are checked when the operation executes. #[derive(Clone, Copy)] -pub struct LoadOptions { +pub struct LoadOptions { + algorithm: PhantomData, valid_items: Option, oob_default: Option, offset: usize, } -impl Default for LoadOptions { +impl Default for LoadOptions { fn default() -> Self { Self { + algorithm: PhantomData, valid_items: None, oob_default: None, offset: 0, @@ -32,7 +57,17 @@ impl Default for LoadOptions { } } -impl LoadOptions { +impl LoadOptions { + /// Select the compile-time algorithm while preserving the other controls. + pub fn algorithm(self, _algorithm: B) -> LoadOptions { + LoadOptions { + algorithm: PhantomData, + valid_items: self.valid_items, + oob_default: self.oob_default, + offset: self.offset, + } + } + /// Load this many elements of the whole group tile, starting at its origin. /// /// The count must not exceed the tile size or `i32::MAX`. Zero is permitted. @@ -60,18 +95,38 @@ impl LoadOptions { } } -/// Controls for a Direct [`store`]. +/// Controls for a [`store`], using [`Direct`] by default. /// /// Defaults write the entire group tile at element offset zero. A partial /// store leaves destination elements outside its prefix unchanged. Controls, /// including whether a count is present, must be uniform across the block. -#[derive(Clone, Copy, Default)] -pub struct StoreOptions { +#[derive(Clone, Copy)] +pub struct StoreOptions { + algorithm: PhantomData, valid_items: Option, offset: usize, } -impl StoreOptions { +impl Default for StoreOptions { + fn default() -> Self { + Self { + algorithm: PhantomData, + valid_items: None, + offset: 0, + } + } +} + +impl StoreOptions { + /// Select the compile-time algorithm while preserving the other controls. + pub fn algorithm(self, _algorithm: B) -> StoreOptions { + StoreOptions { + algorithm: PhantomData, + valid_items: self.valid_items, + offset: self.offset, + } + } + /// Write this many elements of the whole group tile. /// /// The count must not exceed the tile size or `i32::MAX`. A scalar payload @@ -98,9 +153,10 @@ impl StoreOptions { /// `group.count() * ITEMS` elements. Each thread receives consecutive items. /// [`LoadSource`] accepts an immutable slice or a mutable borrow of a collective /// carrier. [`LoadOptions`] supplies the count, typed default and element offset. -/// The operation returns unit and uses no shared scratch. +/// The operation returns unit. [`Direct`] uses no shared scratch; [`Transpose`] +/// uses compiler-planned shared storage and returns the same blocked layout. /// -/// Only complete one-dimensional blocks and the Direct algorithm are supported. +/// Complete one-dimensional blocks support [`Direct`] and [`Transpose`]. /// Compile with the supported Oxide backend and an exact block launch contract. /// Every block thread must reach matching calls with uniform source, controls /// and participation. The compiler checks these requirements and resource @@ -126,31 +182,38 @@ impl StoreOptions { /// } /// ``` #[inline(always)] -pub fn load, const ITEMS: usize, const ALIGN: usize>( +pub fn load, const ITEMS: usize, const ALIGN: usize, A: Algorithm>( group: ThreadGroup, source: S, output: &mut ThreadData, - options: LoadOptions, + options: LoadOptions, ) where Alignment: SupportedAlignment, { source.load(group, output, options); } -/// Store a block tile in blocked order, preserving the per-thread payload. +/// Store a block tile from a blocked per-thread payload. /// /// `T` is inferred from the destination and value. [`StoreValue`] accepts a -/// shared or mutable borrow of [`ThreadData`], or a numeric scalar (one item per -/// thread). `CAPACITY` describes the carrier's ownership grain; it need not equal -/// the payload extent. This Direct operation borrows payload values and flags -/// read-only, returns unit, and uses no shared scratch. +/// shared or mutable borrow of [`ThreadData`] for [`Direct`], a mutable borrow +/// for [`Transpose`], or a numeric scalar (one item per thread) for either. +/// `CAPACITY` describes the carrier's ownership grain; it need not equal the +/// payload extent. Direct preserves values and flags and uses no shared scratch. +/// Transpose uses compiler-planned shared storage and permutes all payload slots +/// into striped order, including a zero-count Store. Destination lane/item +/// `(t, i)` receives the prior blocked slot at linear index `group.count()*i+t`. +/// Its initialized bit follows the value, including unwritten slots and slots +/// outside the write prefix. Scalars use a mutable one-item temporary. /// /// The supported groups, compiler/launch and uniformity requirements are those /// of [`load`]. The destination is the canonical collective carrier: it grants /// disjoint block regions without exposing duplicate mutable lane references. /// Every written payload slot must be initialized. Slots outside the written -/// prefix may be unwritten and are preserved. A following conflicting access -/// to this carrier needs explicit block completion; Drop is not completion. +/// prefix may be unwritten. A following conflicting access to this carrier +/// needs block completion; Drop is not completion. Managed scratch reuse may +/// insert a barrier before a later scratchful call, but is not a trailing +/// completion guarantee. Use an explicit block barrier when completion is needed. /// /// # Panics /// @@ -165,11 +228,11 @@ pub fn load, const ITEMS: usize, const ALIGN: usize /// } /// ``` #[inline(always)] -pub fn store, const CAPACITY: usize>( +pub fn store, const CAPACITY: usize, A: Algorithm>( group: ThreadGroup, destination: &mut DisjointSlice<'_, T, CollectiveTiles>, value: V, - options: StoreOptions, + options: StoreOptions, ) { value.store(group, destination, options); } @@ -182,11 +245,11 @@ pub fn store, const CAPACITY: usize>( /// this trait does not expose pointers or grant lane references. pub trait LoadSource: sealed::LoadSource { #[doc(hidden)] - fn load( + fn load( self, group: ThreadGroup, output: &mut ThreadData, - options: LoadOptions, + options: LoadOptions, ) where Alignment: SupportedAlignment; } @@ -194,11 +257,11 @@ pub trait LoadSource: sealed::LoadSource { impl sealed::LoadSource for &[T] {} impl LoadSource for &[T] { #[inline(always)] - fn load( + fn load( self, _group: ThreadGroup, output: &mut ThreadData, - options: LoadOptions, + options: LoadOptions, ) where Alignment: SupportedAlignment, { @@ -215,11 +278,11 @@ impl sealed::LoadSource } impl LoadSource for &mut DisjointSlice<'_, T, CollectiveTiles> { #[inline(always)] - fn load( + fn load( self, _group: ThreadGroup, output: &mut ThreadData, - options: LoadOptions, + options: LoadOptions, ) where Alignment: SupportedAlignment, { @@ -238,23 +301,26 @@ impl LoadSource for &mut DisjointSlice<'_, T, Col } } -/// A sealed per-thread value accepted by Direct [`store`] for numeric type `T`. +/// A sealed per-thread value accepted by [`store`] for type `T` and algorithm `A`. /// -/// Implemented for numeric scalars and shared/mutable borrows of [`ThreadData`]. -/// A scalar contributes one item per thread. Payload borrows preserve every -/// value and flag, including slots outside the selected write domain. -pub trait StoreValue: sealed::StoreValue { +/// Numeric scalars and mutable [`ThreadData`] borrows support both algorithms. +/// Shared payload borrows support only [`Direct`]. A scalar contributes one item +/// per thread. Direct preserves every value and bit; [`Transpose`] mutates their +/// positions as described by [`store`]. +pub trait StoreValue: sealed::StoreValue { #[doc(hidden)] fn store( self, group: ThreadGroup, destination: &mut DisjointSlice<'_, T, CollectiveTiles>, - options: StoreOptions, + options: StoreOptions, ); } -impl sealed::StoreValue for &ThreadData where - Alignment: SupportedAlignment +impl sealed::StoreValue + for &ThreadData +where + Alignment: SupportedAlignment, { } impl StoreValue for &ThreadData @@ -271,40 +337,51 @@ where boundary::store(destination, self, options); } } -impl sealed::StoreValue for &mut ThreadData where - Alignment: SupportedAlignment +impl sealed::StoreValue + for &mut ThreadData +where + Alignment: SupportedAlignment, { } -impl StoreValue for &mut ThreadData + +impl StoreValue + for &mut ThreadData where - Alignment: SupportedAlignment, + Alignment: SupportedAlignment, { #[inline(always)] fn store( self, - group: ThreadGroup, + _group: ThreadGroup, destination: &mut DisjointSlice<'_, T, CollectiveTiles>, - options: StoreOptions, + options: StoreOptions, ) { - (&*self).store(group, destination, options); + if A::TRANSPOSE { + boundary::store_transpose(destination, self, options.algorithm(Transpose)); + } else { + boundary::store(destination, self, options.algorithm(Direct)); + } } } -impl sealed::StoreValue for T {} -impl StoreValue for T { +impl sealed::StoreValue for T {} +impl StoreValue for T { #[inline(always)] fn store( self, group: ThreadGroup, destination: &mut DisjointSlice<'_, T, CollectiveTiles>, - options: StoreOptions, + options: StoreOptions, ) { let mut value = ThreadData::<_, 1>::new(); value.write(0, self); - (&value).store(group, destination, options); + (&mut value).store(group, destination, options); } } mod sealed { pub trait LoadSource {} - pub trait StoreValue {} + pub trait StoreValue {} + pub trait Algorithm { + const TRANSPOSE: bool; + } } diff --git a/rust/cuda-coop/src/load_store/boundary.rs b/rust/cuda-coop/src/load_store/boundary.rs index 4c0afdbc63ef..88e8a239e439 100644 --- a/rust/cuda-coop/src/load_store/boundary.rs +++ b/rust/cuda-coop/src/load_store/boundary.rs @@ -1,7 +1,7 @@ // Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception -use super::{LoadOptions, StoreOptions}; +use super::{Algorithm, LoadOptions, StoreOptions, Transpose}; use crate::{Alignment, Numeric, SupportedAlignment, ThreadData}; use cuda_device::{CollectiveTiles, DisjointSlice, library, thread}; @@ -18,6 +18,15 @@ registration!(load_filled, "COOP_VALIDITY_LOAD_FILLED"); registration!(store_full, "COOP_VALIDITY_STORE_FULL"); registration!(store_partial, "COOP_VALIDITY_STORE_PARTIAL"); +registration!(transpose_load_full, "COOP_TRANSPOSE_VALIDITY_LOAD_FULL"); +registration!(transpose_load_masked, "COOP_TRANSPOSE_VALIDITY_LOAD_MASKED"); +registration!(transpose_load_filled, "COOP_TRANSPOSE_VALIDITY_LOAD_FILLED"); +registration!(transpose_store_full, "COOP_TRANSPOSE_VALIDITY_STORE_FULL"); +registration!( + transpose_store_partial, + "COOP_TRANSPOSE_VALIDITY_STORE_PARTIAL" +); + #[inline(always)] pub(super) fn tile() -> usize { let threads = thread::blockDim_x() as usize @@ -61,12 +70,18 @@ pub(super) fn containment(offset: usize, valid: usize) { // The caller binds pointer/length to this immutable allocation witness and // checks collective containment when the source is a writable carrier. #[inline(always)] -pub(super) unsafe fn load( +pub(super) unsafe fn load< + T: Numeric, + Memory, + const ITEMS: usize, + const ALIGN: usize, + A: Algorithm, +>( memory: Memory, pointer: *const T, len: usize, output: &mut ThreadData, - options: LoadOptions, + options: LoadOptions, ) where Alignment: SupportedAlignment, { @@ -79,34 +94,40 @@ pub(super) unsafe fn load library::checked_block_resource_call_v1::( - load_full, - memory, - (pointer,), - payload, - ), - (Some(count), None) => { - library::checked_block_resource_call_v1::( - load_masked, + // Each static algorithm keeps its registration function-item identity. The + // storage marker plans the complete provider layout and synchronizes reuse. + macro_rules! call { + ($direct:ident, $transpose:ident, $controls:expr) => { + if A::TRANSPOSE { + library::checked_block_storage_call_v1::( + $transpose, memory, - (pointer, count as i32), + (), + $controls, payload, ) - } - (Some(count), Some(default)) => { + } else { library::checked_block_resource_call_v1::( - load_filled, - memory, - (pointer, count as i32, default), - payload, + $direct, memory, $controls, payload, ) } + }; + } + // SAFETY: the checked marker enforces participation, uniform controls and + // resource epochs. The exact typed provider initializes values before flags + // and transports masked invalid cells without reading them as T. The + // compiler allocates Transpose scratch and completes its previous use. + unsafe { + match (options.valid_items, options.oob_default) { + (None, None) => call!(load_full, transpose_load_full, (pointer,)), + (Some(count), None) => { + call!(load_masked, transpose_load_masked, (pointer, count as i32)) + } + (Some(count), Some(default)) => call!( + load_filled, + transpose_load_filled, + (pointer, count as i32, default) + ), (None, Some(_)) => panic!(), } } @@ -150,3 +171,49 @@ pub(super) fn store( + destination: &mut DisjointSlice<'_, T, CollectiveTiles>, + value: &mut ThreadData, + options: StoreOptions, +) where + Alignment: SupportedAlignment, +{ + let valid = match options.valid_items { + Some(count) => count, + None => tile::(), + }; + extent::(destination.len(), options.offset, valid); + containment::(options.offset, valid); + let pointer = destination.as_mut_ptr(); + let payload = value.block_io_parts(); + // SAFETY: this carrier owns the checked range. The provider checks the + // source write domain before numeric reads, then permutes every opaque cell + // together with its initialized bit. The exclusive payload borrow permits + // that mutation. The compiler plans total scratch and completion before reuse. + unsafe { + let pointer = pointer.add(options.offset); + match options.valid_items { + None => library::checked_block_storage_call_v1::( + transpose_store_full, + destination, + (), + (pointer,), + payload, + ), + Some(count) => library::checked_block_storage_call_v1::( + transpose_store_partial, + destination, + (), + (pointer, count as i32), + payload, + ), + } + } +} diff --git a/rust/cuda-coop/tests/check_load_store.py b/rust/cuda-coop/tests/check_load_store.py index b881828fd738..e83c6dd74914 100644 --- a/rust/cuda-coop/tests/check_load_store.py +++ b/rust/cuda-coop/tests/check_load_store.py @@ -20,6 +20,9 @@ parser.add_argument("--backend", required=True, type=Path) parser.add_argument("--output", required=True, type=Path) parser.add_argument("--cargo-config", type=Path) +parser.add_argument( + "--cargo-lock", type=Path, help="Compatible lockfile for local source patches" +) parser.add_argument("--cuda-root", type=Path, default=Path("/usr/local/cuda")) parser.add_argument("--arch", default="sm_120") args = parser.parse_args() @@ -30,7 +33,9 @@ manifest = (package / "tests/load_store/Cargo.toml").read_text() manifest = manifest.replace('path = "../.."', f"path = {json.dumps(str(package))}") (output / "Cargo.toml").write_text(manifest) -shutil.copyfile(package / "tests/load_store/Cargo.lock", output / "Cargo.lock") +shutil.copyfile( + args.cargo_lock or package / "tests/load_store/Cargo.lock", output / "Cargo.lock" +) source = output / "src/main.rs" prefix = """// Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception @@ -101,6 +106,48 @@ None, ), } +transpose_load = load.replace( + ".oob_default(0)", ".oob_default(0).algorithm(coop::Transpose)" +) +transpose_store = store.replace("&values", "&mut values").replace( + ".valid_items(valid)", ".valid_items(valid).algorithm(coop::Transpose)" +) +cases.update( + shared_transpose=( + load + transpose_store.replace("&mut values", "&values"), + "expected `Direct`, found `Transpose`", + ), + transpose_managed_reuse=( + transpose_load + transpose_store + transpose_store, + None, + ), + transpose_no_trailing_completion=( + transpose_load + transpose_store + reload, + "requires a full-block barrier", + ), + transpose_varying_count=( + transpose_load.replace( + ".valid_items(valid)", ".valid_items(thread::threadIdx_x() as usize)" + ), + "not proved block-uniform", + ), + transpose_loop=( + transpose_load + + "let mut iteration=0;while iteration<2 {" + + transpose_store + + "iteration+=1;}", + None, + ), + generic_algorithm=( + "values.write(0,0_i32);fn forward(block:coop::ThreadGroup,input:&[i32],output:&mut DisjointSlice>,algorithm:A) {" + "let mut values=coop::ThreadData::<_,4>::new();" + "coop::load(block,input,&mut values,coop::LoadOptions::default().algorithm(algorithm));" + "coop::store(block,output,&mut values,coop::StoreOptions::default().algorithm(algorithm));}" + "forward(block,input,&mut output,coop::Transpose);", + None, + ), +) + command = [ str(args.cargo_oxide.resolve()), "oxide", diff --git a/rust/cuda-coop/tests/load_store/src/main.rs b/rust/cuda-coop/tests/load_store/src/main.rs index 920e04b4c7d5..142cfbe7bde9 100644 --- a/rust/cuda-coop/tests/load_store/src/main.rs +++ b/rust/cuda-coop/tests/load_store/src/main.rs @@ -2,6 +2,8 @@ // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception use coop::{Numeric, ThreadData, this_block}; +mod transpose; + use cuda_coop as coop; use cuda_device::{CollectiveTiles, DisjointSlice, kernel, launch_contract, thread}; use cuda_host::cuda_module; @@ -331,6 +333,15 @@ mod kernels { } fn main() -> Result<(), Box> { + if let Some(argument) = std::env::args().nth(1) { + if argument == "transpose" { + return transpose::run(None); + } + if argument.starts_with("transpose-") { + return transpose::run(Some(&argument)); + } + } + use cuda_async::simt::device_context::{init_device_contexts, with_cuda_context}; use cuda_async::simt::device_operation::DeviceOperation; use cuda_core::{DeviceBuffer, LaunchConfig1D}; diff --git a/rust/cuda-coop/tests/load_store/src/transpose.rs b/rust/cuda-coop/tests/load_store/src/transpose.rs new file mode 100644 index 000000000000..8323f2223ed1 --- /dev/null +++ b/rust/cuda-coop/tests/load_store/src/transpose.rs @@ -0,0 +1,710 @@ +// Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +//! Public Transpose checks observe values and arbitrary validity after each Store. +use cuda_coop as coop; +use cuda_device::{CollectiveTiles, DisjointSlice, kernel, launch_contract, thread}; +use cuda_host::cuda_module; + +#[cuda_module] +mod kernels { + use super::*; + + fn expected( + input: &[T], + offset: usize, + index: usize, + items: usize, + valid: usize, + default: T, + mode: u32, + ) -> Option { + if mode == 1 || mode == 5 || mode == 6 { + if index < valid { + Some(input[offset + index]) + } else { + None + } + } else if mode == 2 || mode == 4 { + if index < valid { + Some(input[offset + index]) + } else { + Some(default) + } + } else if mode == 3 { + let lane = index / items; + let item = index % items; + if index < valid || (lane & 1 == 0 && item & 1 == 1) || (lane & 1 == 1 && item & 1 == 0) + { + Some(input[offset + index]) + } else { + None + } + } else { + Some(input[offset + index]) + } + } + + #[inline(never)] + fn exercise( + input: &[T], + output: &mut DisjointSlice>, + valid: usize, + default: T, + mode: u32, + ) { + let block = coop::this_block(); + let threads = thread::blockDim_x() as usize; + let lane = thread::threadIdx_x() as usize; + let offset = thread::blockIdx_x() as usize * threads * CAPACITY; + let mut values = coop::ThreadData::<_, ITEMS>::new(); + if mode != 9 { + coop::load( + block, + input, + &mut values, + coop::LoadOptions::default() + .offset(offset) + .algorithm(coop::Transpose), + ); + } + if mode == 1 || mode == 5 || mode == 6 { + coop::load( + block, + input, + &mut values, + coop::LoadOptions::default() + .offset(offset) + .valid_items(valid) + .algorithm(coop::Transpose), + ); + } else if mode == 2 || mode == 4 { + coop::load( + block, + input, + &mut values, + coop::LoadOptions::default() + .offset(offset) + .valid_items(valid) + .oob_default(default) + .algorithm(coop::Transpose), + ); + } else if mode == 3 { + let incoming = values; + values = coop::ThreadData::new(); + let mut item = 0; + while item < ITEMS { + let index = lane * ITEMS + item; + let initialize = index < valid + || (lane & 1 == 0 && item & 1 == 1) + || (lane & 1 == 1 && item & 1 == 0); + if let (true, Some(value)) = (initialize, incoming.get(item)) { + values.write(item, value); + } + item += 1; + } + } + if mode == 0 || mode == 2 { + coop::store( + block, + &mut *output, + &mut values, + coop::StoreOptions::default() + .offset(offset) + .algorithm(coop::Transpose), + ); + } else { + let written = if mode == 5 { + valid + 1 + } else if mode == 6 { + valid / 2 + } else { + valid + }; + coop::store( + block, + &mut *output, + &mut values, + coop::StoreOptions::default() + .offset(offset) + .valid_items(written) + .algorithm(coop::Transpose), + ); + } + let mut correct = true; + let mut item = 0; + while item < ITEMS { + let source = threads * item + lane; + correct &= + values.get(item) == expected(input, offset, source, ITEMS, valid, default, mode); + item += 1; + } + // Managed scratch reuse inserts completion before this second operation. + // A zero count still permutes all values and bits, with no global writes. + coop::store( + block, + &mut *output, + &mut values, + coop::StoreOptions::default() + .offset(offset) + .valid_items(0) + .algorithm(coop::Transpose), + ); + item = 0; + while item < ITEMS { + let first = threads * item + lane; + let source = threads * (first % ITEMS) + first / ITEMS; + correct &= + values.get(item) == expected(input, offset, source, ITEMS, valid, default, mode); + item += 1; + } + if mode == 4 { + thread::sync_threads(); + coop::load( + block, + &mut *output, + &mut values, + coop::LoadOptions::default() + .offset(offset) + .valid_items(valid) + .oob_default(default) + .algorithm(coop::Transpose), + ); + coop::store( + block, + &mut *output, + &mut values, + coop::StoreOptions::default() + .offset(offset) + .valid_items(valid) + .algorithm(coop::Transpose), + ); + item = 0; + while item < ITEMS { + correct &= values.get(item) + == expected( + input, + offset, + threads * item + lane, + ITEMS, + valid, + default, + mode, + ); + item += 1; + } + } + assert!(correct); + } + #[kernel] + #[launch_contract(domain=1,block=(33,1,1),dynamic_shared=inferred)] + pub fn copy_i8_33_4( + input: &[i8], + mut output: DisjointSlice>, + valid: usize, + default: i8, + mode: u32, + ) { + exercise::<_, ITEMS_PER_THREAD, 4>(input, &mut output, valid, default, mode); + } + #[kernel] + #[launch_contract(domain=1,block=(33,1,1),dynamic_shared=inferred)] + pub fn copy_u8_33_4( + input: &[u8], + mut output: DisjointSlice>, + valid: usize, + default: u8, + mode: u32, + ) { + exercise::<_, ITEMS_PER_THREAD, 4>(input, &mut output, valid, default, mode); + } + #[kernel] + #[launch_contract(domain=1,block=(33,1,1),dynamic_shared=inferred)] + pub fn copy_i16_33_4( + input: &[i16], + mut output: DisjointSlice>, + valid: usize, + default: i16, + mode: u32, + ) { + exercise::<_, ITEMS_PER_THREAD, 4>(input, &mut output, valid, default, mode); + } + #[kernel] + #[launch_contract(domain=1,block=(33,1,1),dynamic_shared=inferred)] + pub fn copy_u16_33_4( + input: &[u16], + mut output: DisjointSlice>, + valid: usize, + default: u16, + mode: u32, + ) { + exercise::<_, ITEMS_PER_THREAD, 4>(input, &mut output, valid, default, mode); + } + #[kernel] + #[launch_contract(domain=1,block=(33,1,1),dynamic_shared=inferred)] + pub fn copy_i32_33_4( + input: &[i32], + mut output: DisjointSlice>, + valid: usize, + default: i32, + mode: u32, + ) { + exercise::<_, ITEMS_PER_THREAD, 4>(input, &mut output, valid, default, mode); + } + #[kernel] + #[launch_contract(domain=1,block=(33,1,1),dynamic_shared=inferred)] + pub fn copy_u32_33_4( + input: &[u32], + mut output: DisjointSlice>, + valid: usize, + default: u32, + mode: u32, + ) { + exercise::<_, ITEMS_PER_THREAD, 4>(input, &mut output, valid, default, mode); + } + #[kernel] + #[launch_contract(domain=1,block=(33,1,1),dynamic_shared=inferred)] + pub fn copy_i64_33_4( + input: &[i64], + mut output: DisjointSlice>, + valid: usize, + default: i64, + mode: u32, + ) { + exercise::<_, ITEMS_PER_THREAD, 4>(input, &mut output, valid, default, mode); + } + #[kernel] + #[launch_contract(domain=1,block=(33,1,1),dynamic_shared=inferred)] + pub fn copy_u64_33_4( + input: &[u64], + mut output: DisjointSlice>, + valid: usize, + default: u64, + mode: u32, + ) { + exercise::<_, ITEMS_PER_THREAD, 4>(input, &mut output, valid, default, mode); + } + #[kernel] + #[launch_contract(domain=1,block=(33,1,1),dynamic_shared=inferred)] + pub fn copy_f32_33_4( + input: &[f32], + mut output: DisjointSlice>, + valid: usize, + default: f32, + mode: u32, + ) { + exercise::<_, ITEMS_PER_THREAD, 4>(input, &mut output, valid, default, mode); + } + #[kernel] + #[launch_contract(domain=1,block=(33,1,1),dynamic_shared=inferred)] + pub fn copy_f64_33_4( + input: &[f64], + mut output: DisjointSlice>, + valid: usize, + default: f64, + mode: u32, + ) { + exercise::<_, ITEMS_PER_THREAD, 4>(input, &mut output, valid, default, mode); + } + #[kernel] + #[launch_contract(domain=1,block=(1,1,1),dynamic_shared=inferred)] + pub fn copy_i32_1_4( + input: &[i32], + mut output: DisjointSlice>, + valid: usize, + default: i32, + mode: u32, + ) { + exercise::<_, ITEMS_PER_THREAD, 4>(input, &mut output, valid, default, mode); + } + #[kernel] + #[launch_contract(domain=1,block=(128,1,1),dynamic_shared=inferred)] + pub fn copy_i32_128_4( + input: &[i32], + mut output: DisjointSlice>, + valid: usize, + default: i32, + mode: u32, + ) { + exercise::<_, ITEMS_PER_THREAD, 4>(input, &mut output, valid, default, mode); + } + #[kernel] + #[launch_contract(domain=1,block=(130,1,1),dynamic_shared=inferred)] + pub fn copy_i32_130_4( + input: &[i32], + mut output: DisjointSlice>, + valid: usize, + default: i32, + mode: u32, + ) { + exercise::<_, ITEMS_PER_THREAD, 4>(input, &mut output, valid, default, mode); + } + #[kernel] + #[launch_contract(domain=1,block=(33,1,1),dynamic_shared=inferred)] + pub fn copy_i32_33_9( + input: &[i32], + mut output: DisjointSlice>, + valid: usize, + default: i32, + mode: u32, + ) { + exercise::<_, ITEMS_PER_THREAD, 9>(input, &mut output, valid, default, mode); + } + #[kernel] + #[launch_contract(domain=1,block=(33,1,1),dynamic_shared=inferred)] + pub fn transpose_scalar(mut output: DisjointSlice>) { + coop::store( + coop::this_block(), + &mut output, + -17_i16, + coop::StoreOptions::default().algorithm(coop::Transpose), + ); + } +} + +pub fn run(trap: Option<&str>) -> Result<(), Box> { + use cuda_async::simt::device_context::{init_device_contexts, with_cuda_context}; + use cuda_async::simt::device_operation::DeviceOperation; + use cuda_core::{DeviceBuffer, LaunchConfig1D}; + init_device_contexts(0, 1)?; + let context = with_cuda_context(0, |context| context.clone())?; + let stream = context.default_stream(); + // SAFETY: these typed declarations and native kernels are compiled together. + let module = unsafe { kernels::load_async(0)? }; + if let Some(trap) = trap { + let plan = module.prepare_copy_i32_33_4::<4>(LaunchConfig1D::auto_shared(1, 33))?; + let source = DeviceBuffer::from_host(&stream, &vec![11_i32; 132])?; + let target = DeviceBuffer::from_host(&stream, &vec![0_i32; 132])?; + let (valid, mode) = match trap { + "transpose-count" => (133, 1), + "transpose-unwritten" => (132, 9), + "transpose-different-count" => (31, 5), + _ => panic!("unknown Transpose trap"), + }; + let _ = module + .copy_i32_33_4_async_owned::<4, _, _>(&plan, source, target, valid, -7, mode) + .sync()?; + panic!("expected Transpose trap returned successfully"); + } + let mut cases = 0; + macro_rules! run { + ($t:ty,$prepare:ident,$launch:ident,$block:literal,$capacity:literal,$items:literal,$default:expr) => {{ + let default: $t = $default; + for groups in [1_u32, 3] { + let plan = + module.$prepare::<$items>(LaunchConfig1D::auto_shared(groups, $block))?; + assert!(plan.__raw_config().shared_mem_bytes > 0); + assert!( + module + .$prepare::<$items>(LaunchConfig1D::auto_shared(groups, $block + 1)) + .is_err() + ); + for mode in [0, 1, 2, 3, 4, 6] { + for valid in [0, 1, $block * $items - 1, $block * $items] { + if mode == 0 && valid != $block * $items { + continue; + } + let size = groups as usize * $block * $capacity; + let input: Vec<$t> = (0..size).map(|i| (i % 61 + 1) as $t).collect(); + let source = DeviceBuffer::from_host(&stream, &input)?; + let target = DeviceBuffer::from_host(&stream, &vec![0 as $t; size])?; + let (_, target) = module + .$launch::<$items, _, _>(&plan, source, target, valid, default, mode) + .sync()?; + let mut result = vec![0 as $t; size]; + target.copy_to_host(&stream, &mut result)?; + stream.synchronize()?; + for (index, &value) in result.iter().enumerate() { + let within = index % ($block * $capacity); + let written = if mode == 0 || mode == 2 { + $block * $items + } else if mode == 6 { + valid / 2 + } else { + valid + }; + let expected = if within < written { + if mode == 2 && within >= valid { + default + } else { + input[index] + } + } else { + 0 as $t + }; + assert_eq!( + value.to_ne_bytes(), + expected.to_ne_bytes(), + "Transpose type={} B={} I={} mode={} valid={} index={}", + stringify!($t), + $block, + $items, + mode, + valid, + index + ); + } + cases += 1; + } + } + } + }}; + } + run!( + i8, + prepare_copy_i8_33_4, + copy_i8_33_4_async_owned, + 33, + 4, + 1, + i8::MIN + ); + run!( + i8, + prepare_copy_i8_33_4, + copy_i8_33_4_async_owned, + 33, + 4, + 4, + i8::MIN + ); + run!( + u8, + prepare_copy_u8_33_4, + copy_u8_33_4_async_owned, + 33, + 4, + 1, + u8::MAX + ); + run!( + u8, + prepare_copy_u8_33_4, + copy_u8_33_4_async_owned, + 33, + 4, + 4, + u8::MAX + ); + run!( + i16, + prepare_copy_i16_33_4, + copy_i16_33_4_async_owned, + 33, + 4, + 1, + i16::MIN + ); + run!( + i16, + prepare_copy_i16_33_4, + copy_i16_33_4_async_owned, + 33, + 4, + 4, + i16::MIN + ); + run!( + u16, + prepare_copy_u16_33_4, + copy_u16_33_4_async_owned, + 33, + 4, + 1, + u16::MAX + ); + run!( + u16, + prepare_copy_u16_33_4, + copy_u16_33_4_async_owned, + 33, + 4, + 4, + u16::MAX + ); + run!( + i32, + prepare_copy_i32_33_4, + copy_i32_33_4_async_owned, + 33, + 4, + 1, + i32::MIN + ); + run!( + i32, + prepare_copy_i32_33_4, + copy_i32_33_4_async_owned, + 33, + 4, + 4, + i32::MIN + ); + run!( + u32, + prepare_copy_u32_33_4, + copy_u32_33_4_async_owned, + 33, + 4, + 1, + u32::MAX + ); + run!( + u32, + prepare_copy_u32_33_4, + copy_u32_33_4_async_owned, + 33, + 4, + 4, + u32::MAX + ); + run!( + i64, + prepare_copy_i64_33_4, + copy_i64_33_4_async_owned, + 33, + 4, + 1, + i64::MIN + 17 + ); + run!( + i64, + prepare_copy_i64_33_4, + copy_i64_33_4_async_owned, + 33, + 4, + 4, + i64::MIN + 17 + ); + run!( + u64, + prepare_copy_u64_33_4, + copy_u64_33_4_async_owned, + 33, + 4, + 1, + u64::MAX - 17 + ); + run!( + u64, + prepare_copy_u64_33_4, + copy_u64_33_4_async_owned, + 33, + 4, + 4, + u64::MAX - 17 + ); + run!( + f32, + prepare_copy_f32_33_4, + copy_f32_33_4_async_owned, + 33, + 4, + 1, + -0.0_f32 + ); + run!( + f32, + prepare_copy_f32_33_4, + copy_f32_33_4_async_owned, + 33, + 4, + 4, + -0.0_f32 + ); + run!( + f64, + prepare_copy_f64_33_4, + copy_f64_33_4_async_owned, + 33, + 4, + 1, + -0.0_f64 + ); + run!( + f64, + prepare_copy_f64_33_4, + copy_f64_33_4_async_owned, + 33, + 4, + 4, + -0.0_f64 + ); + run!( + i32, + prepare_copy_i32_1_4, + copy_i32_1_4_async_owned, + 1, + 4, + 1, + i32::MIN + ); + run!( + i32, + prepare_copy_i32_1_4, + copy_i32_1_4_async_owned, + 1, + 4, + 4, + i32::MIN + ); + run!( + i32, + prepare_copy_i32_128_4, + copy_i32_128_4_async_owned, + 128, + 4, + 1, + i32::MIN + ); + run!( + i32, + prepare_copy_i32_128_4, + copy_i32_128_4_async_owned, + 128, + 4, + 4, + i32::MIN + ); + run!( + i32, + prepare_copy_i32_130_4, + copy_i32_130_4_async_owned, + 130, + 4, + 1, + i32::MIN + ); + run!( + i32, + prepare_copy_i32_130_4, + copy_i32_130_4_async_owned, + 130, + 4, + 4, + i32::MIN + ); + run!( + i32, + prepare_copy_i32_33_9, + copy_i32_33_9_async_owned, + 33, + 9, + 9, + i32::MIN + ); + let plan = module.prepare_transpose_scalar(LaunchConfig1D::auto_shared(1, 33))?; + let target = DeviceBuffer::from_host(&stream, &[0_i16; 33])?; + let target = module + .transpose_scalar_async_owned::<_>(&plan, target) + .sync()?; + let mut result = vec![0_i16; 33]; + target.copy_to_host(&stream, &mut result)?; + stream.synchronize()?; + assert!(result.iter().all(|&value| value == -17)); + cases += 1; + println!( + "PASS: {cases} public Transpose cases, all ten types, arbitrary masks, repeated Store, partial warps and N9" + ); + Ok(()) +} From ef81dc5e1ff08354594fcaad4b9be8e9196a384e Mon Sep 17 00:00:00 2001 From: Trent Nelson Date: Tue, 20 Oct 2026 11:54:22 -0700 Subject: [PATCH 29/33] Pin Transpose operations to the published storage compiler Signed-off-by: Trent Nelson --- rust/Cargo.lock | 8 ++--- rust/cuda-coop/Cargo.toml | 2 +- rust/cuda-coop/examples/load_store/Cargo.lock | 30 +++++++++---------- rust/cuda-coop/examples/load_store/Cargo.toml | 8 ++--- rust/cuda-coop/tests/load_store/Cargo.lock | 30 +++++++++---------- rust/cuda-coop/tests/load_store/Cargo.toml | 8 ++--- 6 files changed, 43 insertions(+), 43 deletions(-) diff --git a/rust/Cargo.lock b/rust/Cargo.lock index 03bddcbec468..b092c7b62893 100644 --- a/rust/Cargo.lock +++ b/rust/Cargo.lock @@ -57,7 +57,7 @@ dependencies = [ [[package]] name = "cuda-device" version = "0.2.1" -source = "git+https://github.com/tpn/cuda-rust.git?rev=344d4c683ab4e4d9a8802aac636060d6b627a38b#344d4c683ab4e4d9a8802aac636060d6b627a38b" +source = "git+https://github.com/tpn/cuda-rust.git?rev=204e4e6f721bd111bf5e0e00caa69dc672a4d378#204e4e6f721bd111bf5e0e00caa69dc672a4d378" dependencies = [ "cuda-macros", ] @@ -65,7 +65,7 @@ dependencies = [ [[package]] name = "cuda-macros" version = "0.2.1" -source = "git+https://github.com/tpn/cuda-rust.git?rev=344d4c683ab4e4d9a8802aac636060d6b627a38b#344d4c683ab4e4d9a8802aac636060d6b627a38b" +source = "git+https://github.com/tpn/cuda-rust.git?rev=204e4e6f721bd111bf5e0e00caa69dc672a4d378#204e4e6f721bd111bf5e0e00caa69dc672a4d378" dependencies = [ "oxide-tile-export", "proc-macro2", @@ -127,7 +127,7 @@ checksum = "cf8baf1c55e62ffcace7a9f06f4bd9cd3f0c4beb022d3b367256b91b87513d98" [[package]] name = "oxide-tile-export" version = "0.2.1" -source = "git+https://github.com/tpn/cuda-rust.git?rev=344d4c683ab4e4d9a8802aac636060d6b627a38b#344d4c683ab4e4d9a8802aac636060d6b627a38b" +source = "git+https://github.com/tpn/cuda-rust.git?rev=204e4e6f721bd111bf5e0e00caa69dc672a4d378#204e4e6f721bd111bf5e0e00caa69dc672a4d378" dependencies = [ "serde", "serde_json", @@ -155,7 +155,7 @@ dependencies = [ [[package]] name = "reserved-oxide-symbols" version = "0.2.1" -source = "git+https://github.com/tpn/cuda-rust.git?rev=344d4c683ab4e4d9a8802aac636060d6b627a38b#344d4c683ab4e4d9a8802aac636060d6b627a38b" +source = "git+https://github.com/tpn/cuda-rust.git?rev=204e4e6f721bd111bf5e0e00caa69dc672a4d378#204e4e6f721bd111bf5e0e00caa69dc672a4d378" [[package]] name = "serde" diff --git a/rust/cuda-coop/Cargo.toml b/rust/cuda-coop/Cargo.toml index 8bd4ee382a20..0f3c80862caf 100644 --- a/rust/cuda-coop/Cargo.toml +++ b/rust/cuda-coop/Cargo.toml @@ -9,7 +9,7 @@ license.workspace = true publish = false [dependencies] -cuda-device = { git = "https://github.com/tpn/cuda-rust.git", rev = "344d4c683ab4e4d9a8802aac636060d6b627a38b" } +cuda-device = { git = "https://github.com/tpn/cuda-rust.git", rev = "204e4e6f721bd111bf5e0e00caa69dc672a4d378" } [build-dependencies] cuda-coop-codegen = { path = "../cuda-coop-codegen" } diff --git a/rust/cuda-coop/examples/load_store/Cargo.lock b/rust/cuda-coop/examples/load_store/Cargo.lock index f7068217c0a7..161a9f9544e0 100644 --- a/rust/cuda-coop/examples/load_store/Cargo.lock +++ b/rust/cuda-coop/examples/load_store/Cargo.lock @@ -143,7 +143,7 @@ dependencies = [ [[package]] name = "cuda-artifact-finalizer" version = "0.2.1" -source = "git+https://github.com/tpn/cuda-rust.git?rev=344d4c683ab4e4d9a8802aac636060d6b627a38b#344d4c683ab4e4d9a8802aac636060d6b627a38b" +source = "git+https://github.com/tpn/cuda-rust.git?rev=204e4e6f721bd111bf5e0e00caa69dc672a4d378#204e4e6f721bd111bf5e0e00caa69dc672a4d378" dependencies = [ "libnvvm-sys", "nvjitlink-sys", @@ -155,7 +155,7 @@ dependencies = [ [[package]] name = "cuda-async" version = "0.4.0" -source = "git+https://github.com/tpn/cuda-rust.git?rev=344d4c683ab4e4d9a8802aac636060d6b627a38b#344d4c683ab4e4d9a8802aac636060d6b627a38b" +source = "git+https://github.com/tpn/cuda-rust.git?rev=204e4e6f721bd111bf5e0e00caa69dc672a4d378#204e4e6f721bd111bf5e0e00caa69dc672a4d378" dependencies = [ "anyhow", "cuda-bindings", @@ -171,7 +171,7 @@ dependencies = [ [[package]] name = "cuda-bindings" version = "0.4.0" -source = "git+https://github.com/tpn/cuda-rust.git?rev=344d4c683ab4e4d9a8802aac636060d6b627a38b#344d4c683ab4e4d9a8802aac636060d6b627a38b" +source = "git+https://github.com/tpn/cuda-rust.git?rev=204e4e6f721bd111bf5e0e00caa69dc672a4d378#204e4e6f721bd111bf5e0e00caa69dc672a4d378" dependencies = [ "bindgen", "libloading", @@ -202,7 +202,7 @@ dependencies = [ [[package]] name = "cuda-core" version = "0.4.0" -source = "git+https://github.com/tpn/cuda-rust.git?rev=344d4c683ab4e4d9a8802aac636060d6b627a38b#344d4c683ab4e4d9a8802aac636060d6b627a38b" +source = "git+https://github.com/tpn/cuda-rust.git?rev=204e4e6f721bd111bf5e0e00caa69dc672a4d378#204e4e6f721bd111bf5e0e00caa69dc672a4d378" dependencies = [ "anyhow", "cuda-bindings", @@ -214,7 +214,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" version = "0.4.0" -source = "git+https://github.com/tpn/cuda-rust.git?rev=344d4c683ab4e4d9a8802aac636060d6b627a38b#344d4c683ab4e4d9a8802aac636060d6b627a38b" +source = "git+https://github.com/tpn/cuda-rust.git?rev=204e4e6f721bd111bf5e0e00caa69dc672a4d378#204e4e6f721bd111bf5e0e00caa69dc672a4d378" dependencies = [ "proc-macro2", "quote", @@ -224,7 +224,7 @@ dependencies = [ [[package]] name = "cuda-device" version = "0.2.1" -source = "git+https://github.com/tpn/cuda-rust.git?rev=344d4c683ab4e4d9a8802aac636060d6b627a38b#344d4c683ab4e4d9a8802aac636060d6b627a38b" +source = "git+https://github.com/tpn/cuda-rust.git?rev=204e4e6f721bd111bf5e0e00caa69dc672a4d378#204e4e6f721bd111bf5e0e00caa69dc672a4d378" dependencies = [ "cuda-macros", ] @@ -232,7 +232,7 @@ dependencies = [ [[package]] name = "cuda-host" version = "0.2.1" -source = "git+https://github.com/tpn/cuda-rust.git?rev=344d4c683ab4e4d9a8802aac636060d6b627a38b#344d4c683ab4e4d9a8802aac636060d6b627a38b" +source = "git+https://github.com/tpn/cuda-rust.git?rev=204e4e6f721bd111bf5e0e00caa69dc672a4d378#204e4e6f721bd111bf5e0e00caa69dc672a4d378" dependencies = [ "cuda-artifact-finalizer", "cuda-async", @@ -249,7 +249,7 @@ dependencies = [ [[package]] name = "cuda-macros" version = "0.2.1" -source = "git+https://github.com/tpn/cuda-rust.git?rev=344d4c683ab4e4d9a8802aac636060d6b627a38b#344d4c683ab4e4d9a8802aac636060d6b627a38b" +source = "git+https://github.com/tpn/cuda-rust.git?rev=204e4e6f721bd111bf5e0e00caa69dc672a4d378#204e4e6f721bd111bf5e0e00caa69dc672a4d378" dependencies = [ "oxide-tile-export", "proc-macro2", @@ -263,7 +263,7 @@ dependencies = [ [[package]] name = "cuda-target-spec" version = "0.1.0" -source = "git+https://github.com/tpn/cuda-rust.git?rev=344d4c683ab4e4d9a8802aac636060d6b627a38b#344d4c683ab4e4d9a8802aac636060d6b627a38b" +source = "git+https://github.com/tpn/cuda-rust.git?rev=204e4e6f721bd111bf5e0e00caa69dc672a4d378#204e4e6f721bd111bf5e0e00caa69dc672a4d378" [[package]] name = "dashmap" @@ -518,7 +518,7 @@ dependencies = [ [[package]] name = "libnvvm-sys" version = "0.2.1" -source = "git+https://github.com/tpn/cuda-rust.git?rev=344d4c683ab4e4d9a8802aac636060d6b627a38b#344d4c683ab4e4d9a8802aac636060d6b627a38b" +source = "git+https://github.com/tpn/cuda-rust.git?rev=204e4e6f721bd111bf5e0e00caa69dc672a4d378#204e4e6f721bd111bf5e0e00caa69dc672a4d378" dependencies = [ "cuda-target-spec", "libloading", @@ -571,7 +571,7 @@ dependencies = [ [[package]] name = "nvjitlink-sys" version = "0.2.1" -source = "git+https://github.com/tpn/cuda-rust.git?rev=344d4c683ab4e4d9a8802aac636060d6b627a38b#344d4c683ab4e4d9a8802aac636060d6b627a38b" +source = "git+https://github.com/tpn/cuda-rust.git?rev=204e4e6f721bd111bf5e0e00caa69dc672a4d378#204e4e6f721bd111bf5e0e00caa69dc672a4d378" dependencies = [ "libloading", "thiserror 2.0.21", @@ -598,7 +598,7 @@ checksum = "9f7c3e4beb33f85d45ae3e3a1792185706c8e16d043238c593331cc7cd313b50" [[package]] name = "oxide-artifacts" version = "0.2.1" -source = "git+https://github.com/tpn/cuda-rust.git?rev=344d4c683ab4e4d9a8802aac636060d6b627a38b#344d4c683ab4e4d9a8802aac636060d6b627a38b" +source = "git+https://github.com/tpn/cuda-rust.git?rev=204e4e6f721bd111bf5e0e00caa69dc672a4d378#204e4e6f721bd111bf5e0e00caa69dc672a4d378" dependencies = [ "object", ] @@ -606,7 +606,7 @@ dependencies = [ [[package]] name = "oxide-tile-export" version = "0.2.1" -source = "git+https://github.com/tpn/cuda-rust.git?rev=344d4c683ab4e4d9a8802aac636060d6b627a38b#344d4c683ab4e4d9a8802aac636060d6b627a38b" +source = "git+https://github.com/tpn/cuda-rust.git?rev=204e4e6f721bd111bf5e0e00caa69dc672a4d378#204e4e6f721bd111bf5e0e00caa69dc672a4d378" dependencies = [ "serde", "serde_json", @@ -654,7 +654,7 @@ dependencies = [ [[package]] name = "ptx-parse" version = "0.1.0" -source = "git+https://github.com/tpn/cuda-rust.git?rev=344d4c683ab4e4d9a8802aac636060d6b627a38b#344d4c683ab4e4d9a8802aac636060d6b627a38b" +source = "git+https://github.com/tpn/cuda-rust.git?rev=204e4e6f721bd111bf5e0e00caa69dc672a4d378#204e4e6f721bd111bf5e0e00caa69dc672a4d378" [[package]] name = "quote" @@ -706,7 +706,7 @@ checksum = "d6f6ff9a378485b298a5286656da665ba74413d36db0979633275d2e708145d4" [[package]] name = "reserved-oxide-symbols" version = "0.2.1" -source = "git+https://github.com/tpn/cuda-rust.git?rev=344d4c683ab4e4d9a8802aac636060d6b627a38b#344d4c683ab4e4d9a8802aac636060d6b627a38b" +source = "git+https://github.com/tpn/cuda-rust.git?rev=204e4e6f721bd111bf5e0e00caa69dc672a4d378#204e4e6f721bd111bf5e0e00caa69dc672a4d378" [[package]] name = "rustc-hash" diff --git a/rust/cuda-coop/examples/load_store/Cargo.toml b/rust/cuda-coop/examples/load_store/Cargo.toml index 442eaf871c8b..258c8f4142bd 100644 --- a/rust/cuda-coop/examples/load_store/Cargo.toml +++ b/rust/cuda-coop/examples/load_store/Cargo.toml @@ -13,7 +13,7 @@ members = ["helper"] [dependencies] cuda-coop = { path = "../.." } coop-load-store-helper = { path = "helper" } -cuda-device = { git = "https://github.com/tpn/cuda-rust.git", rev = "344d4c683ab4e4d9a8802aac636060d6b627a38b" } -cuda-host = { git = "https://github.com/tpn/cuda-rust.git", rev = "344d4c683ab4e4d9a8802aac636060d6b627a38b", features = ["async"] } -cuda-core = { git = "https://github.com/tpn/cuda-rust.git", rev = "344d4c683ab4e4d9a8802aac636060d6b627a38b" } -cuda-async = { git = "https://github.com/tpn/cuda-rust.git", rev = "344d4c683ab4e4d9a8802aac636060d6b627a38b" } +cuda-device = { git = "https://github.com/tpn/cuda-rust.git", rev = "204e4e6f721bd111bf5e0e00caa69dc672a4d378" } +cuda-host = { git = "https://github.com/tpn/cuda-rust.git", rev = "204e4e6f721bd111bf5e0e00caa69dc672a4d378", features = ["async"] } +cuda-core = { git = "https://github.com/tpn/cuda-rust.git", rev = "204e4e6f721bd111bf5e0e00caa69dc672a4d378" } +cuda-async = { git = "https://github.com/tpn/cuda-rust.git", rev = "204e4e6f721bd111bf5e0e00caa69dc672a4d378" } diff --git a/rust/cuda-coop/tests/load_store/Cargo.lock b/rust/cuda-coop/tests/load_store/Cargo.lock index 5bcb0ea6a68a..4e95c613b326 100644 --- a/rust/cuda-coop/tests/load_store/Cargo.lock +++ b/rust/cuda-coop/tests/load_store/Cargo.lock @@ -135,7 +135,7 @@ dependencies = [ [[package]] name = "cuda-artifact-finalizer" version = "0.2.1" -source = "git+https://github.com/tpn/cuda-rust.git?rev=344d4c683ab4e4d9a8802aac636060d6b627a38b#344d4c683ab4e4d9a8802aac636060d6b627a38b" +source = "git+https://github.com/tpn/cuda-rust.git?rev=204e4e6f721bd111bf5e0e00caa69dc672a4d378#204e4e6f721bd111bf5e0e00caa69dc672a4d378" dependencies = [ "libnvvm-sys", "nvjitlink-sys", @@ -147,7 +147,7 @@ dependencies = [ [[package]] name = "cuda-async" version = "0.4.0" -source = "git+https://github.com/tpn/cuda-rust.git?rev=344d4c683ab4e4d9a8802aac636060d6b627a38b#344d4c683ab4e4d9a8802aac636060d6b627a38b" +source = "git+https://github.com/tpn/cuda-rust.git?rev=204e4e6f721bd111bf5e0e00caa69dc672a4d378#204e4e6f721bd111bf5e0e00caa69dc672a4d378" dependencies = [ "anyhow", "cuda-bindings", @@ -163,7 +163,7 @@ dependencies = [ [[package]] name = "cuda-bindings" version = "0.4.0" -source = "git+https://github.com/tpn/cuda-rust.git?rev=344d4c683ab4e4d9a8802aac636060d6b627a38b#344d4c683ab4e4d9a8802aac636060d6b627a38b" +source = "git+https://github.com/tpn/cuda-rust.git?rev=204e4e6f721bd111bf5e0e00caa69dc672a4d378#204e4e6f721bd111bf5e0e00caa69dc672a4d378" dependencies = [ "bindgen", "libloading", @@ -194,7 +194,7 @@ dependencies = [ [[package]] name = "cuda-core" version = "0.4.0" -source = "git+https://github.com/tpn/cuda-rust.git?rev=344d4c683ab4e4d9a8802aac636060d6b627a38b#344d4c683ab4e4d9a8802aac636060d6b627a38b" +source = "git+https://github.com/tpn/cuda-rust.git?rev=204e4e6f721bd111bf5e0e00caa69dc672a4d378#204e4e6f721bd111bf5e0e00caa69dc672a4d378" dependencies = [ "anyhow", "cuda-bindings", @@ -206,7 +206,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" version = "0.4.0" -source = "git+https://github.com/tpn/cuda-rust.git?rev=344d4c683ab4e4d9a8802aac636060d6b627a38b#344d4c683ab4e4d9a8802aac636060d6b627a38b" +source = "git+https://github.com/tpn/cuda-rust.git?rev=204e4e6f721bd111bf5e0e00caa69dc672a4d378#204e4e6f721bd111bf5e0e00caa69dc672a4d378" dependencies = [ "proc-macro2", "quote", @@ -216,7 +216,7 @@ dependencies = [ [[package]] name = "cuda-device" version = "0.2.1" -source = "git+https://github.com/tpn/cuda-rust.git?rev=344d4c683ab4e4d9a8802aac636060d6b627a38b#344d4c683ab4e4d9a8802aac636060d6b627a38b" +source = "git+https://github.com/tpn/cuda-rust.git?rev=204e4e6f721bd111bf5e0e00caa69dc672a4d378#204e4e6f721bd111bf5e0e00caa69dc672a4d378" dependencies = [ "cuda-macros", ] @@ -224,7 +224,7 @@ dependencies = [ [[package]] name = "cuda-host" version = "0.2.1" -source = "git+https://github.com/tpn/cuda-rust.git?rev=344d4c683ab4e4d9a8802aac636060d6b627a38b#344d4c683ab4e4d9a8802aac636060d6b627a38b" +source = "git+https://github.com/tpn/cuda-rust.git?rev=204e4e6f721bd111bf5e0e00caa69dc672a4d378#204e4e6f721bd111bf5e0e00caa69dc672a4d378" dependencies = [ "cuda-artifact-finalizer", "cuda-async", @@ -241,7 +241,7 @@ dependencies = [ [[package]] name = "cuda-macros" version = "0.2.1" -source = "git+https://github.com/tpn/cuda-rust.git?rev=344d4c683ab4e4d9a8802aac636060d6b627a38b#344d4c683ab4e4d9a8802aac636060d6b627a38b" +source = "git+https://github.com/tpn/cuda-rust.git?rev=204e4e6f721bd111bf5e0e00caa69dc672a4d378#204e4e6f721bd111bf5e0e00caa69dc672a4d378" dependencies = [ "oxide-tile-export", "proc-macro2", @@ -255,7 +255,7 @@ dependencies = [ [[package]] name = "cuda-target-spec" version = "0.1.0" -source = "git+https://github.com/tpn/cuda-rust.git?rev=344d4c683ab4e4d9a8802aac636060d6b627a38b#344d4c683ab4e4d9a8802aac636060d6b627a38b" +source = "git+https://github.com/tpn/cuda-rust.git?rev=204e4e6f721bd111bf5e0e00caa69dc672a4d378#204e4e6f721bd111bf5e0e00caa69dc672a4d378" [[package]] name = "dashmap" @@ -510,7 +510,7 @@ dependencies = [ [[package]] name = "libnvvm-sys" version = "0.2.1" -source = "git+https://github.com/tpn/cuda-rust.git?rev=344d4c683ab4e4d9a8802aac636060d6b627a38b#344d4c683ab4e4d9a8802aac636060d6b627a38b" +source = "git+https://github.com/tpn/cuda-rust.git?rev=204e4e6f721bd111bf5e0e00caa69dc672a4d378#204e4e6f721bd111bf5e0e00caa69dc672a4d378" dependencies = [ "cuda-target-spec", "libloading", @@ -563,7 +563,7 @@ dependencies = [ [[package]] name = "nvjitlink-sys" version = "0.2.1" -source = "git+https://github.com/tpn/cuda-rust.git?rev=344d4c683ab4e4d9a8802aac636060d6b627a38b#344d4c683ab4e4d9a8802aac636060d6b627a38b" +source = "git+https://github.com/tpn/cuda-rust.git?rev=204e4e6f721bd111bf5e0e00caa69dc672a4d378#204e4e6f721bd111bf5e0e00caa69dc672a4d378" dependencies = [ "libloading", "thiserror 2.0.21", @@ -590,7 +590,7 @@ checksum = "9f7c3e4beb33f85d45ae3e3a1792185706c8e16d043238c593331cc7cd313b50" [[package]] name = "oxide-artifacts" version = "0.2.1" -source = "git+https://github.com/tpn/cuda-rust.git?rev=344d4c683ab4e4d9a8802aac636060d6b627a38b#344d4c683ab4e4d9a8802aac636060d6b627a38b" +source = "git+https://github.com/tpn/cuda-rust.git?rev=204e4e6f721bd111bf5e0e00caa69dc672a4d378#204e4e6f721bd111bf5e0e00caa69dc672a4d378" dependencies = [ "object", ] @@ -598,7 +598,7 @@ dependencies = [ [[package]] name = "oxide-tile-export" version = "0.2.1" -source = "git+https://github.com/tpn/cuda-rust.git?rev=344d4c683ab4e4d9a8802aac636060d6b627a38b#344d4c683ab4e4d9a8802aac636060d6b627a38b" +source = "git+https://github.com/tpn/cuda-rust.git?rev=204e4e6f721bd111bf5e0e00caa69dc672a4d378#204e4e6f721bd111bf5e0e00caa69dc672a4d378" dependencies = [ "serde", "serde_json", @@ -646,7 +646,7 @@ dependencies = [ [[package]] name = "ptx-parse" version = "0.1.0" -source = "git+https://github.com/tpn/cuda-rust.git?rev=344d4c683ab4e4d9a8802aac636060d6b627a38b#344d4c683ab4e4d9a8802aac636060d6b627a38b" +source = "git+https://github.com/tpn/cuda-rust.git?rev=204e4e6f721bd111bf5e0e00caa69dc672a4d378#204e4e6f721bd111bf5e0e00caa69dc672a4d378" [[package]] name = "quote" @@ -698,7 +698,7 @@ checksum = "d6f6ff9a378485b298a5286656da665ba74413d36db0979633275d2e708145d4" [[package]] name = "reserved-oxide-symbols" version = "0.2.1" -source = "git+https://github.com/tpn/cuda-rust.git?rev=344d4c683ab4e4d9a8802aac636060d6b627a38b#344d4c683ab4e4d9a8802aac636060d6b627a38b" +source = "git+https://github.com/tpn/cuda-rust.git?rev=204e4e6f721bd111bf5e0e00caa69dc672a4d378#204e4e6f721bd111bf5e0e00caa69dc672a4d378" [[package]] name = "rustc-hash" diff --git a/rust/cuda-coop/tests/load_store/Cargo.toml b/rust/cuda-coop/tests/load_store/Cargo.toml index 235f9d76fdc6..2b2d125d2aa3 100644 --- a/rust/cuda-coop/tests/load_store/Cargo.toml +++ b/rust/cuda-coop/tests/load_store/Cargo.toml @@ -11,7 +11,7 @@ publish = false [dependencies] cuda-coop = { path = "../.." } -cuda-device = { git = "https://github.com/tpn/cuda-rust.git", rev = "344d4c683ab4e4d9a8802aac636060d6b627a38b" } -cuda-host = { git = "https://github.com/tpn/cuda-rust.git", rev = "344d4c683ab4e4d9a8802aac636060d6b627a38b", features = ["async"] } -cuda-core = { git = "https://github.com/tpn/cuda-rust.git", rev = "344d4c683ab4e4d9a8802aac636060d6b627a38b" } -cuda-async = { git = "https://github.com/tpn/cuda-rust.git", rev = "344d4c683ab4e4d9a8802aac636060d6b627a38b" } +cuda-device = { git = "https://github.com/tpn/cuda-rust.git", rev = "204e4e6f721bd111bf5e0e00caa69dc672a4d378" } +cuda-host = { git = "https://github.com/tpn/cuda-rust.git", rev = "204e4e6f721bd111bf5e0e00caa69dc672a4d378", features = ["async"] } +cuda-core = { git = "https://github.com/tpn/cuda-rust.git", rev = "204e4e6f721bd111bf5e0e00caa69dc672a4d378" } +cuda-async = { git = "https://github.com/tpn/cuda-rust.git", rev = "204e4e6f721bd111bf5e0e00caa69dc672a4d378" } From a2553c98b9049be7cebddec5981e492e7b8bae99 Mon Sep 17 00:00:00 2001 From: Trent Nelson Date: Tue, 20 Oct 2026 12:16:47 -0700 Subject: [PATCH 30/33] Clarify Transpose examples and check participation Signed-off-by: Trent Nelson --- docs/rust/load_store.rst | 12 +++++++----- rust/cuda-coop-codegen/src/lib.rs | 2 +- rust/cuda-coop/src/load_store.rs | 15 +++++++++++++++ rust/cuda-coop/tests/check_load_store.py | 4 ++++ 4 files changed, 27 insertions(+), 6 deletions(-) diff --git a/docs/rust/load_store.rst b/docs/rust/load_store.rst index b0264afe105e..be7ce6d16142 100644 --- a/docs/rust/load_store.rst +++ b/docs/rust/load_store.rst @@ -20,8 +20,9 @@ A complete program ------------------ The example is an independent Cargo workspace with a separate generic helper -crate. It copies three block tiles at item counts 1 and 4, including a short last -tile. The consumer has no build script, native shim, provider, or scratch argument. +crate. It copies three block tiles with Direct and Transpose at item counts 1 +and 4, including a short last tile. The consumer has no build script, native +shim, provider, or scratch argument. Its ``cuda-coop`` dependency supplies the native specialization during the normal Oxide build. @@ -39,14 +40,15 @@ Choose the architecture for your GPU. Materialization produces the native device artifact during the Cargo build. No separate compilation or linking step is required. The exact block launch contract is checked when preparing the launch; ``auto_shared`` uses the compiler's inferred requirement, which is zero for -Direct Load/Store. +Direct Load/Store and includes the planned shared storage for Transpose. .. literalinclude:: ../../rust/cuda-coop/examples/load_store/src/main.rs :language: rust :start-after: // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception -The reusable helper preserves the source element type and borrows its payload -read-only for Direct Store: +The reusable helper preserves the source element type. Its algorithm-generic +form borrows the payload mutably because Transpose Store permutes it. +Direct-only code may pass a shared payload borrow: .. literalinclude:: ../../rust/cuda-coop/examples/load_store/helper/src/lib.rs :language: rust diff --git a/rust/cuda-coop-codegen/src/lib.rs b/rust/cuda-coop-codegen/src/lib.rs index a3f4e62bf986..83851bb6cd2e 100644 --- a/rust/cuda-coop-codegen/src/lib.rs +++ b/rust/cuda-coop-codegen/src/lib.rs @@ -263,7 +263,7 @@ impl Request { /// /// # Errors /// - /// Rejects the invalid shapes documented by [`Self::new`], a non-i32 raw + /// Rejects the invalid shapes documented by [`Self::new`], or a non-i32 raw /// recipe. pub fn new_typed( dtype: Dtype, diff --git a/rust/cuda-coop/src/load_store.rs b/rust/cuda-coop/src/load_store.rs index 094c83b7356b..2b9b770f288d 100644 --- a/rust/cuda-coop/src/load_store.rs +++ b/rust/cuda-coop/src/load_store.rs @@ -15,6 +15,21 @@ pub struct Direct; /// /// Load produces blocked payloads. Store permutes payload values and validity /// into striped order, including slots outside its selected write prefix. +/// +/// ```no_run +/// use cuda_coop as coop; +/// fn copy_tile( +/// input: &[i16], +/// output: &mut coop::DisjointSlice<'_, i16, coop::CollectiveTiles>, +/// ) { +/// let block = coop::this_block(); +/// let mut values = coop::ThreadData::<_, ITEMS_PER_THREAD>::new(); +/// coop::load(block, input, &mut values, +/// coop::LoadOptions::default().algorithm(coop::Transpose)); +/// coop::store(block, output, &mut values, +/// coop::StoreOptions::default().algorithm(coop::Transpose)); +/// } +/// ``` #[derive(Clone, Copy, Debug, Default)] pub struct Transpose; diff --git a/rust/cuda-coop/tests/check_load_store.py b/rust/cuda-coop/tests/check_load_store.py index e83c6dd74914..e10ab4d21d19 100644 --- a/rust/cuda-coop/tests/check_load_store.py +++ b/rust/cuda-coop/tests/check_load_store.py @@ -131,6 +131,10 @@ ), "not proved block-uniform", ), + transpose_varying_participation=( + "if thread::threadIdx_x()&1==0 {" + transpose_load + "}", + "lane-varying", + ), transpose_loop=( transpose_load + "let mut iteration=0;while iteration<2 {" From 41a9f11c38935fdbbee4908e8f8c66d9039e629e Mon Sep 17 00:00:00 2001 From: Trent Nelson Date: Tue, 20 Oct 2026 12:09:59 -0700 Subject: [PATCH 31/33] Preserve exact block shapes in native Load and Store Signed-off-by: Trent Nelson --- docs/rust/index.rst | 7 +- docs/rust/load_store.rst | 12 +- rust/cuda-coop-codegen/src/lib.rs | 91 +++++++++++--- rust/cuda-coop-codegen/src/oxide.rs | 35 ++++-- .../tests/check_collective_rejections.py | 8 -- rust/cuda-coop-codegen/tests/native.rs | 71 ++++++++++- rust/cuda-coop/src/load_store.rs | 4 +- rust/cuda-coop/tests/load_store/src/main.rs | 109 +++++++++++++++-- .../tests/load_store/src/transpose.rs | 112 ++++++++++++++++-- 9 files changed, 379 insertions(+), 70 deletions(-) diff --git a/docs/rust/index.rst b/docs/rust/index.rst index 82d6bd1a140d..d6e7802cc427 100644 --- a/docs/rust/index.rst +++ b/docs/rust/index.rst @@ -23,8 +23,11 @@ libraries does not require Rust, and the provider does not depend on Python. Supported provider forms ------------------------ -Each request describes a complete, one-dimensional block and a positive -compile-time item count per thread. The original raw ``i32`` recipes support +Each request describes a complete block with exact ``[X,Y,Z]`` dimensions and a +positive compile-time item count per thread. ``Request::new`` and +``Request::new_typed`` retain their one-dimensional ``[B,1,1]`` shape; +``new_with_shape`` and ``new_typed_with_shape`` accept all three dimensions. +Thread rank is x-fastest. The original raw ``i32`` recipes support these forms: .. list-table:: diff --git a/docs/rust/load_store.rst b/docs/rust/load_store.rst index be7ce6d16142..35e407a7c9b1 100644 --- a/docs/rust/load_store.rst +++ b/docs/rust/load_store.rst @@ -6,8 +6,8 @@ Block Load and Store ``cuda_coop::load`` fills an existing ``ThreadData`` and returns unit. ``cuda_coop::store`` writes its values and returns unit. Both take the group -first. This initial implementation supports complete, one-dimensional blocks -with Direct and Transpose algorithms. Both Load algorithms produce blocked +first. Complete one-, two- and three-dimensional blocks support Direct and +Transpose algorithms. Both Load algorithms produce blocked payloads, where each thread owns consecutive elements. Other algorithms, warp groups, and public temporary storage remain under development. @@ -64,6 +64,14 @@ per-thread count. It must fit the tile and ``i32::MAX``. Offsets are element indices from the start of the allocation; supply each block's origin explicitly. Counts, defaults, offsets, and their presence must be uniform across the block. +For block shape ``[X,Y,Z]``, thread rank is x-fastest: +``t = x + X * (y + Y * z)``. A blocked payload slot ``i`` addresses +``offset + t * ITEMS_PER_THREAD + i``. The exact shape is part of the kernel +launch contract; another shape with the same thread count is rejected. +Use ``LaunchConfig2D`` or ``LaunchConfig3D`` for the corresponding launch domain. +The block origin in a multidimensional grid uses x-fastest block rank as well; +include that rank times the block tile extent in the element offset. + A masked Load without ``oob_default`` resets every out-of-prefix slot to unwritten, even if initialized before the Load. ``get`` returns ``None`` for an unwritten slot; indexing it traps. A default requires an explicit count and has diff --git a/rust/cuda-coop-codegen/src/lib.rs b/rust/cuda-coop-codegen/src/lib.rs index 83851bb6cd2e..dfc24b818e3d 100644 --- a/rust/cuda-coop-codegen/src/lib.rs +++ b/rust/cuda-coop-codegen/src/lib.rs @@ -10,7 +10,7 @@ //! //! # Supported specializations //! -//! [`Request`] describes a one-dimensional complete block, positive fixed +//! [`Request`] describes an exact complete block shape, positive fixed //! items per thread, and [`Algorithm::Direct`] or [`Algorithm::Transpose`]. //! The raw i32 recipes support Direct and Transpose. Their Store payloads must //! be fully initialized. The validity-aware recipes support all ten @@ -217,7 +217,7 @@ impl Operation { } } -/// A validated, one-dimensional block specialization. +/// A validated specialization for an exact one-, two- or three-dimensional block. /// /// This is an in-process implementation interface, not a serialized compiler /// protocol. All fields are static specialization inputs. The generated ABI @@ -227,11 +227,11 @@ pub struct Request { operation: Operation, algorithm: Algorithm, dtype: Dtype, - block_threads: u32, + block: [u32; 3], items_per_thread: u32, } impl Request { - /// Validate an exact block shape and payload extent. + /// Validate a one-dimensional 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 @@ -271,15 +271,64 @@ impl Request { algorithm: Algorithm, block_threads: u32, items_per_thread: u32, + ) -> Result { + Self::new_typed_with_shape( + dtype, + operation, + algorithm, + [block_threads, 1, 1], + items_per_thread, + ) + } + + /// Validate an exact `[X, Y, Z]` block shape with i32 elements. + /// + /// Each axis must be positive, their product at most 1024, and the positive + /// per-thread item count times that product must fit i32. Threads are ranked + /// x-fastest: `x + X * (y + Y * z)`. Device-specific axis and launch limits + /// must also be checked by the caller when preparing a launch. + /// + /// # Errors + /// + /// Rejects zero axes, an oversized or overflowing block product, zero items, + /// or a tile extent that cannot be represented by CUB's signed prefix count. + pub fn new_with_shape( + operation: Operation, + algorithm: Algorithm, + block: [u32; 3], + items_per_thread: u32, + ) -> Result { + Self::new_typed_with_shape(Dtype::I32, operation, algorithm, block, items_per_thread) + } + + /// Validate an exact typed `[X, Y, Z]` block specialization. + /// + /// Shape and launch requirements match [`Self::new_with_shape`]. The five + /// validity-aware recipes support every [`Dtype`]; raw recipes require i32. + /// + /// # Errors + /// + /// Rejects the shapes documented by [`Self::new_with_shape`] and a non-i32 + /// raw recipe. + pub fn new_typed_with_shape( + dtype: Dtype, + operation: Operation, + algorithm: Algorithm, + block: [u32; 3], + items_per_thread: u32, ) -> Result { if !operation.validity() && dtype != Dtype::I32 { return Err(Error::message( "legacy raw recipes require signed i32 elements", )); } - if !(1..=1024).contains(&block_threads) + let block_threads = block + .into_iter() + .try_fold(1_u32, u32::checked_mul) + .filter(|threads| (1..=1024).contains(threads)); + if block_threads.is_none() || items_per_thread == 0 - || u64::from(block_threads) * u64::from(items_per_thread) > i32::MAX as u64 + || u64::from(block_threads.unwrap_or(0)) * 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", @@ -289,7 +338,7 @@ impl Request { operation, algorithm, dtype, - block_threads, + block, items_per_thread, }) } @@ -305,9 +354,13 @@ impl Request { pub fn dtype(self) -> Dtype { self.dtype } - /// Exact number of threads in the one-dimensional block. + /// Exact block shape `[X, Y, Z]`, with x-fastest thread rank. + pub fn block_shape(self) -> [u32; 3] { + self.block + } + /// Total number of threads in the validated block. pub fn block_threads(self) -> u32 { - self.block_threads + self.block[0] * self.block[1] * self.block[2] } /// Number of payload elements per thread. pub fn items_per_thread(self) -> u32 { @@ -570,11 +623,13 @@ fn source(request: Request, symbol: &str) -> String { }; let (dtype, size) = request.dtype.native(); let items = request.items_per_thread; + let [block_x, block_y, block_z] = request.block; + let threads = request.block_threads(); let bitmap_bytes = items.div_ceil(8); - let tile = request.block_threads * items; + let tile = threads * items; let mask_plane = op == Operation::StorePartialValidity && transpose && items > 1; let storage = if mask_plane { - let plane_bytes = u64::from(request.block_threads.div_ceil(32)) * u64::from(items) * 4; + let plane_bytes = u64::from(threads.div_ceil(32)) * u64::from(items) * 4; format!( r#"struct scratch_storage {{ alignas(alignof(primitive::TempStorage)) unsigned char cub[sizeof(primitive::TempStorage)]; @@ -623,7 +678,7 @@ static_assert(alignof(scratch_storage) == alignof(primitive::TempStorage));"# op, Operation::LoadMaskedValidity | Operation::StorePartialValidity ) { - format!("int(threadIdx.x) * {items} + i < valid") + format!("tid * {items} + i < valid") } else { "true".to_owned() }; @@ -653,7 +708,6 @@ static_assert(alignof(scratch_storage) == alignof(primitive::TempStorage));"# let prepare = if mask_plane { format!( r#"unsigned char* const mask_plane = scratch + offsetof(scratch_storage, masks); - const int tid = int(threadIdx.x); const int warp = tid / 32; const int remaining = {threads} - warp * 32; const unsigned int active = remaining >= 32 ? 0xffffffffu : (1u << remaining) - 1; @@ -664,8 +718,7 @@ static_assert(alignof(scratch_storage) == alignof(primitive::TempStorage));"# void* const address = mask_plane + sizeof(unsigned int) * (warp * {items} + i); ::new (address) unsigned int(bits); }} - }}"#, - threads = request.block_threads + }}"# ) } else { String::new() @@ -704,8 +757,7 @@ static_assert(alignof(scratch_storage) == alignof(primitive::TempStorage));"# }} }} flags[byte] = static_cast(bits); - }}"#, - threads = request.block_threads + }}"# ) } else { // With one item per lane, BlockedToStriped is identity. @@ -761,23 +813,24 @@ static_assert(alignof(scratch_storage) == alignof(primitive::TempStorage));"# r#"// Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception #include +#include #include #include namespace {symbol}_detail {{ using value_type = {dtype}; using cell = {cell}; -using primitive = cub::{primitive}; +using primitive = cub::{primitive}; {storage} static_assert(sizeof(value_type) == {size} && alignof(value_type) == {size}); static_assert(sizeof(cell) == sizeof(value_type) && alignof(cell) == alignof(value_type)); extern "C" __device__ void {symbol}({qualifier}value_type* memory, {payload_const}value_type* payload, unsigned char* scratch{flags}{controls}) {{ + const int tid = cub::RowMajorTid({block_x}, {block_y}, {block_z}); {check} {constructor} {body} }} }} "#, - threads = request.block_threads, choice = request.algorithm.token() ) } diff --git a/rust/cuda-coop-codegen/src/oxide.rs b/rust/cuda-coop-codegen/src/oxide.rs index c12d5b7e34ce..08972f32f85e 100644 --- a/rust/cuda-coop-codegen/src/oxide.rs +++ b/rust/cuda-coop-codegen/src/oxide.rs @@ -115,11 +115,6 @@ struct SpecializationRequest { impl SpecializationRequest { fn native(&self) -> Result { - 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, @@ -147,11 +142,11 @@ impl SpecializationRequest { ))); } }; - Request::new_typed( + Request::new_typed_with_shape( self.dtype.dtype()?, operation, algorithm, - self.block[0], + self.block, self.items_per_thread, ) } @@ -593,6 +588,26 @@ mod tests { ); } + #[test] + fn preserves_multidimensional_raw_and_typed_requests() { + for (operation, dtype) in [ + ("load_full", Scalar::Signed(32)), + ("store_partial_validity_bits_v1", Scalar::Unsigned(64)), + ] { + let request = SpecializationRequest { + operation: operation.into(), + algorithm: "transpose".into(), + dtype, + block: [8, 4, 4], + items_per_thread: 4, + }; + let native = request.native().unwrap(); + assert_eq!(native.block_shape(), [8, 4, 4]); + assert_eq!(native.block_threads(), 128); + assert_eq!(native.dtype(), request.dtype.dtype().unwrap()); + } + } + #[test] fn rejects_request_before_opening_toolkit() { let base = serde_json::json!({"version":1,"target":"compute_120","cuda_root":"/absent", "requests":[{ @@ -609,11 +624,7 @@ mod tests { serde_json::json!(64), "signed i32", ), - ( - "/requests/0/block/1", - serde_json::json!(2), - "exact block shape", - ), + ("/requests/0/block/1", serde_json::json!(0), "block threads"), ( "/requests/0/operation", serde_json::json!("unknown"), diff --git a/rust/cuda-coop-codegen/tests/check_collective_rejections.py b/rust/cuda-coop-codegen/tests/check_collective_rejections.py index c4fcf2e3ac2b..7f422c751fca 100644 --- a/rust/cuda-coop-codegen/tests/check_collective_rejections.py +++ b/rust/cuda-coop-codegen/tests/check_collective_rejections.py @@ -179,14 +179,6 @@ def in_helper(code): {"root": original["root"].replace("block = (128, 1, 1), ", "")}, "requires #[launch_contract(block", ), - "two_dimensional_shape": ( - { - "root": original["root"].replace( - "block = (128, 1, 1)", "block = (64, 2, 1)", 1 - ) - }, - "exact block shape [B,1,1]", - ), "shared_shape": ( { "root": "block = (64, 1, 1)".join( diff --git a/rust/cuda-coop-codegen/tests/native.rs b/rust/cuda-coop-codegen/tests/native.rs index 08bde41edc5c..a2abff03a481 100644 --- a/rust/cuda-coop-codegen/tests/native.rs +++ b/rust/cuda-coop-codegen/tests/native.rs @@ -11,11 +11,50 @@ 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()); } + for block in [[0, 4, 4], [8, 0, 4], [8, 4, 0], [8, 4, 33], [u32::MAX; 3]] { + assert!(Request::new_with_shape(Operation::LoadFull, Algorithm::Direct, block, 4).is_err()); + } assert!( Request::new_typed(Dtype::U64, Operation::LoadFull, Algorithm::Direct, 128, 4).is_err() ); } +#[test] +fn preserves_exact_shape_and_legacy_constructors() { + let old = Request::new_typed( + Dtype::U8, + Operation::LoadFullValidity, + Algorithm::Direct, + 128, + 4, + ) + .unwrap(); + let shaped = Request::new_typed_with_shape( + Dtype::U8, + Operation::LoadFullValidity, + Algorithm::Direct, + [128, 1, 1], + 4, + ) + .unwrap(); + assert_eq!(old, shaped); + assert_eq!(shaped.block_shape(), [128, 1, 1]); + let three = Request::new_typed_with_shape( + Dtype::U8, + Operation::LoadFullValidity, + Algorithm::Direct, + [8, 4, 4], + 4, + ) + .unwrap(); + assert_eq!(three.block_threads(), old.block_threads()); + assert_ne!(three, old); + assert_eq!( + Request::new(Operation::LoadFull, Algorithm::Direct, 128, 4).unwrap(), + Request::new_with_shape(Operation::LoadFull, Algorithm::Direct, [128, 1, 1], 4).unwrap(), + ); +} + #[test] #[ignore = "requires CUDA_HOME and a Linux CUDA toolkit with NVRTC LTO support"] fn native_typed_validity_and_determinism() -> Result<(), Box> { @@ -77,6 +116,23 @@ fn native_typed_validity_and_determinism() -> Result<(), Box Result<(), Box entry.cub_layout.size); - if request.block_threads() == 128 { + if request.block_threads() == 128 && request.items_per_thread() == 4 { assert_eq!(entry.cub_layout.size, 2064); assert_eq!(entry.scratch.size, 2128); } @@ -161,10 +217,21 @@ fn native_specializations_and_determinism() -> Result<(), Box StoreOptions { /// The operation returns unit. [`Direct`] uses no shared scratch; [`Transpose`] /// uses compiler-planned shared storage and returns the same blocked layout. /// -/// Complete one-dimensional blocks support [`Direct`] and [`Transpose`]. +/// Complete one-, two- and three-dimensional blocks support [`Direct`] and +/// [`Transpose`]. Thread rank is x-fastest: `x + X * (y + Y * z)` for shape +/// `[X, Y, Z]`. Each ranked thread receives `ITEMS` consecutive elements. /// Compile with the supported Oxide backend and an exact block launch contract. /// Every block thread must reach matching calls with uniform source, controls /// and participation. The compiler checks these requirements and resource diff --git a/rust/cuda-coop/tests/load_store/src/main.rs b/rust/cuda-coop/tests/load_store/src/main.rs index 142cfbe7bde9..f93056d411d1 100644 --- a/rust/cuda-coop/tests/load_store/src/main.rs +++ b/rust/cuda-coop/tests/load_store/src/main.rs @@ -22,9 +22,16 @@ mod kernels { mode: u32, ) { let block = this_block(); - let threads = thread::blockDim_x() as usize; - let lane = thread::threadIdx_x() as usize; - let offset = thread::blockIdx_x() as usize * threads * 4 + shift; + let x = thread::blockDim_x() as usize; + let y = thread::blockDim_y() as usize; + let threads = x * y * thread::blockDim_z() as usize; + let lane = thread::threadIdx_x() as usize + + x * (thread::threadIdx_y() as usize + y * thread::threadIdx_z() as usize); + let group = thread::blockIdx_x() as usize + + thread::gridDim_x() as usize + * (thread::blockIdx_y() as usize + + thread::gridDim_y() as usize * thread::blockIdx_z() as usize); + let offset = group * threads * 4 + shift; let mut values = ThreadData::<_, ITEMS_PER_THREAD>::new(); match mode { 0 => coop::load( @@ -302,6 +309,30 @@ mod kernels { exercise::<_, ITEMS_PER_THREAD>(input, &mut output, valid, shift, default, mode); } #[kernel] + #[launch_contract(domain = 2, block = (11, 3, 1), dynamic_shared = inferred)] + pub fn copy_i32_2d( + input: &[i32], + mut output: DisjointSlice>, + valid: usize, + shift: usize, + default: i32, + mode: u32, + ) { + exercise::<_, ITEMS_PER_THREAD>(input, &mut output, valid, shift, default, mode); + } + #[kernel] + #[launch_contract(domain = 3, block = (8, 4, 4), dynamic_shared = inferred)] + pub fn copy_i32_3d( + input: &[i32], + mut output: DisjointSlice>, + valid: usize, + shift: usize, + default: i32, + mode: u32, + ) { + exercise::<_, ITEMS_PER_THREAD>(input, &mut output, valid, shift, default, mode); + } + #[kernel] #[launch_contract(domain = 1, block = (128, 1, 1), dynamic_shared = inferred)] pub fn scalar(mut output: DisjointSlice>, value: i16) { let block = this_block(); @@ -344,7 +375,7 @@ fn main() -> Result<(), Box> { use cuda_async::simt::device_context::{init_device_contexts, with_cuda_context}; use cuda_async::simt::device_operation::DeviceOperation; - use cuda_core::{DeviceBuffer, LaunchConfig1D}; + use cuda_core::{DeviceBuffer, LaunchConfig1D, LaunchConfig2D, LaunchConfig3D}; init_device_contexts(0, 1)?; let context = with_cuda_context(0, |context| context.clone())?; let stream = context.default_stream(); @@ -353,17 +384,31 @@ fn main() -> Result<(), Box> { let trap = std::env::args().nth(1); let mut cases = 0; macro_rules! run { - ($t:ty, $prepare:ident, $launch:ident, $block:literal, $items:literal, $default:expr) => {{ + ($t:ty, $prepare:ident, $launch:ident, $block:literal, $items:literal, $default:expr) => { + run!( + $t, + $prepare, + $launch, + $block, + $items, + $default, + |groups| LaunchConfig1D::auto_shared(groups, $block), + |groups| LaunchConfig1D::auto_shared(groups, $block / 2) + ); + }; + ($t:ty, $prepare:ident, $launch:ident, $block:literal, $items:literal, $default:expr, + $config:expr, $wrong_shape:expr) => {{ let default: $t = $default; for groups in [1_u32, 3] { - let prepared = - module.$prepare::<$items>(LaunchConfig1D::auto_shared(groups, $block))?; + let prepared = module.$prepare::<$items>(($config)(groups))?; assert_eq!(prepared.__raw_config().shared_mem_bytes, 0); - assert!( - module - .$prepare::<$items>(LaunchConfig1D::auto_shared(groups, $block / 2)) - .is_err() - ); + let error = module + .$prepare::<$items>(($wrong_shape)(groups)) + .err() + .expect("a different exact block shape must reject"); + assert!(error.to_string().contains("does not match required block")); + let (gx, gy, gz) = prepared.__raw_config().grid_dim; + let groups = gx * gy * gz; for mode in [0, 1, 2, 3, 4, 6] { for valid in [0, 1, $block * $items - 1, $block * $items] { if (mode == 0) && valid != $block * $items { @@ -620,6 +665,46 @@ fn main() -> Result<(), Box> { 4, -1234567 ); + run!( + i32, + prepare_copy_i32_2d, + copy_i32_2d_async_owned, + 33, + 1, + -7, + |groups| LaunchConfig2D::auto_shared((groups, 2), (11, 3)), + |groups| LaunchConfig2D::auto_shared((groups, 2), (3, 11)) + ); + run!( + i32, + prepare_copy_i32_3d, + copy_i32_3d_async_owned, + 128, + 1, + -7, + |groups| LaunchConfig3D::auto_shared((groups, 2, 2), (8, 4, 4)), + |groups| LaunchConfig3D::auto_shared((groups, 2, 2), (16, 2, 4)) + ); + run!( + i32, + prepare_copy_i32_2d, + copy_i32_2d_async_owned, + 33, + 4, + -7, + |groups| LaunchConfig2D::auto_shared((groups, 2), (11, 3)), + |groups| LaunchConfig2D::auto_shared((groups, 2), (3, 11)) + ); + run!( + i32, + prepare_copy_i32_3d, + copy_i32_3d_async_owned, + 128, + 4, + -7, + |groups| LaunchConfig3D::auto_shared((groups, 2, 2), (8, 4, 4)), + |groups| LaunchConfig3D::auto_shared((groups, 2, 2), (16, 2, 4)) + ); let plan = module.prepare_scalar(LaunchConfig1D::auto_shared(3, 128))?; let target = DeviceBuffer::from_host(&stream, &vec![0_i16; 384])?; let target = module diff --git a/rust/cuda-coop/tests/load_store/src/transpose.rs b/rust/cuda-coop/tests/load_store/src/transpose.rs index 8323f2223ed1..6fc2cb9fbadb 100644 --- a/rust/cuda-coop/tests/load_store/src/transpose.rs +++ b/rust/cuda-coop/tests/load_store/src/transpose.rs @@ -54,9 +54,16 @@ mod kernels { mode: u32, ) { let block = coop::this_block(); - let threads = thread::blockDim_x() as usize; - let lane = thread::threadIdx_x() as usize; - let offset = thread::blockIdx_x() as usize * threads * CAPACITY; + let x = thread::blockDim_x() as usize; + let y = thread::blockDim_y() as usize; + let threads = x * y * thread::blockDim_z() as usize; + let lane = thread::threadIdx_x() as usize + + x * (thread::threadIdx_y() as usize + y * thread::threadIdx_z() as usize); + let group = thread::blockIdx_x() as usize + + thread::gridDim_x() as usize + * (thread::blockIdx_y() as usize + + thread::gridDim_y() as usize * thread::blockIdx_z() as usize); + let offset = group * threads * CAPACITY; let mut values = coop::ThreadData::<_, ITEMS>::new(); if mode != 9 { coop::load( @@ -351,6 +358,28 @@ mod kernels { exercise::<_, ITEMS_PER_THREAD, 9>(input, &mut output, valid, default, mode); } #[kernel] + #[launch_contract(domain=2,block=(11,3,1),dynamic_shared=inferred)] + pub fn copy_i32_2d( + input: &[i32], + mut output: DisjointSlice>, + valid: usize, + default: i32, + mode: u32, + ) { + exercise::<_, ITEMS_PER_THREAD, 4>(input, &mut output, valid, default, mode); + } + #[kernel] + #[launch_contract(domain=3,block=(8,4,4),dynamic_shared=inferred)] + pub fn copy_i32_3d( + input: &[i32], + mut output: DisjointSlice>, + valid: usize, + default: i32, + mode: u32, + ) { + exercise::<_, ITEMS_PER_THREAD, 4>(input, &mut output, valid, default, mode); + } + #[kernel] #[launch_contract(domain=1,block=(33,1,1),dynamic_shared=inferred)] pub fn transpose_scalar(mut output: DisjointSlice>) { coop::store( @@ -365,7 +394,7 @@ mod kernels { pub fn run(trap: Option<&str>) -> Result<(), Box> { use cuda_async::simt::device_context::{init_device_contexts, with_cuda_context}; use cuda_async::simt::device_operation::DeviceOperation; - use cuda_core::{DeviceBuffer, LaunchConfig1D}; + use cuda_core::{DeviceBuffer, LaunchConfig1D, LaunchConfig2D, LaunchConfig3D}; init_device_contexts(0, 1)?; let context = with_cuda_context(0, |context| context.clone())?; let stream = context.default_stream(); @@ -388,17 +417,32 @@ pub fn run(trap: Option<&str>) -> Result<(), Box> { } let mut cases = 0; macro_rules! run { - ($t:ty,$prepare:ident,$launch:ident,$block:literal,$capacity:literal,$items:literal,$default:expr) => {{ + ($t:ty,$prepare:ident,$launch:ident,$block:literal,$capacity:literal,$items:literal,$default:expr) => { + run!( + $t, + $prepare, + $launch, + $block, + $capacity, + $items, + $default, + |groups| LaunchConfig1D::auto_shared(groups, $block), + |groups| LaunchConfig1D::auto_shared(groups, $block + 1) + ); + }; + ($t:ty,$prepare:ident,$launch:ident,$block:literal,$capacity:literal,$items:literal,$default:expr, + $config:expr,$wrong_shape:expr) => {{ let default: $t = $default; for groups in [1_u32, 3] { - let plan = - module.$prepare::<$items>(LaunchConfig1D::auto_shared(groups, $block))?; + let plan = module.$prepare::<$items>(($config)(groups))?; assert!(plan.__raw_config().shared_mem_bytes > 0); - assert!( - module - .$prepare::<$items>(LaunchConfig1D::auto_shared(groups, $block + 1)) - .is_err() - ); + let error = module + .$prepare::<$items>(($wrong_shape)(groups)) + .err() + .expect("a different exact block shape must reject"); + assert!(error.to_string().contains("does not match required block")); + let (gx, gy, gz) = plan.__raw_config().grid_dim; + let groups = gx * gy * gz; for mode in [0, 1, 2, 3, 4, 6] { for valid in [0, 1, $block * $items - 1, $block * $items] { if mode == 0 && valid != $block * $items { @@ -693,6 +737,50 @@ pub fn run(trap: Option<&str>) -> Result<(), Box> { 9, i32::MIN ); + run!( + i32, + prepare_copy_i32_2d, + copy_i32_2d_async_owned, + 33, + 4, + 1, + -7, + |groups| LaunchConfig2D::auto_shared((groups, 2), (11, 3)), + |groups| LaunchConfig2D::auto_shared((groups, 2), (3, 11)) + ); + run!( + i32, + prepare_copy_i32_3d, + copy_i32_3d_async_owned, + 128, + 4, + 1, + -7, + |groups| LaunchConfig3D::auto_shared((groups, 2, 2), (8, 4, 4)), + |groups| LaunchConfig3D::auto_shared((groups, 2, 2), (16, 2, 4)) + ); + run!( + i32, + prepare_copy_i32_2d, + copy_i32_2d_async_owned, + 33, + 4, + 4, + -7, + |groups| LaunchConfig2D::auto_shared((groups, 2), (11, 3)), + |groups| LaunchConfig2D::auto_shared((groups, 2), (3, 11)) + ); + run!( + i32, + prepare_copy_i32_3d, + copy_i32_3d_async_owned, + 128, + 4, + 4, + -7, + |groups| LaunchConfig3D::auto_shared((groups, 2, 2), (8, 4, 4)), + |groups| LaunchConfig3D::auto_shared((groups, 2, 2), (16, 2, 4)) + ); let plan = module.prepare_transpose_scalar(LaunchConfig1D::auto_shared(1, 33))?; let target = DeviceBuffer::from_host(&stream, &[0_i16; 33])?; let target = module From 47861b64afa07476dcb17c1df18172f85f8bcaa2 Mon Sep 17 00:00:00 2001 From: Trent Nelson Date: Tue, 20 Oct 2026 12:38:37 -0700 Subject: [PATCH 32/33] Clarify block-shape ownership offsets and launch contracts Signed-off-by: Trent Nelson --- docs/rust/load_store.rst | 6 +++++- rust/cuda-coop-codegen/src/lib.rs | 1 + rust/cuda-coop-codegen/tests/native.rs | 10 ++++++++-- 3 files changed, 14 insertions(+), 3 deletions(-) diff --git a/docs/rust/load_store.rst b/docs/rust/load_store.rst index 35e407a7c9b1..74cbe55cafe8 100644 --- a/docs/rust/load_store.rst +++ b/docs/rust/load_store.rst @@ -69,8 +69,12 @@ For block shape ``[X,Y,Z]``, thread rank is x-fastest: ``offset + t * ITEMS_PER_THREAD + i``. The exact shape is part of the kernel launch contract; another shape with the same thread count is rejected. Use ``LaunchConfig2D`` or ``LaunchConfig3D`` for the corresponding launch domain. +Declare the matching kernel contract, for example +``#[launch_contract(domain = 3, block = (8, 4, 4), dynamic_shared = inferred)]``. The block origin in a multidimensional grid uses x-fastest block rank as well; -include that rank times the block tile extent in the element offset. +for a collective carrier, include that rank times its ownership grain +``block_size * CAPACITY`` in the element offset (see `Collective ownership and +completion`_). Read-only slice loads have no ownership grain. A masked Load without ``oob_default`` resets every out-of-prefix slot to unwritten, even if initialized before the Load. ``get`` returns ``None`` for an diff --git a/rust/cuda-coop-codegen/src/lib.rs b/rust/cuda-coop-codegen/src/lib.rs index dfc24b818e3d..3383bb5016a6 100644 --- a/rust/cuda-coop-codegen/src/lib.rs +++ b/rust/cuda-coop-codegen/src/lib.rs @@ -390,6 +390,7 @@ pub struct Layout { /// 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. +/// A lane is the x-fastest linear thread rank, `cub::RowMajorTid(X, Y, Z)`. /// /// 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. diff --git a/rust/cuda-coop-codegen/tests/native.rs b/rust/cuda-coop-codegen/tests/native.rs index a2abff03a481..1c782d59df39 100644 --- a/rust/cuda-coop-codegen/tests/native.rs +++ b/rust/cuda-coop-codegen/tests/native.rs @@ -11,7 +11,13 @@ 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()); } - for block in [[0, 4, 4], [8, 0, 4], [8, 4, 0], [8, 4, 33], [u32::MAX; 3]] { + for block in [ + [0, 4, 4], + [8, 0, 4], + [8, 4, 0], + [8, 4, 33], + [u32::MAX, u32::MAX, 1], + ] { assert!(Request::new_with_shape(Operation::LoadFull, Algorithm::Direct, block, 4).is_err()); } assert!( @@ -117,7 +123,7 @@ fn native_typed_validity_and_determinism() -> Result<(), Box Date: Tue, 20 Oct 2026 12:35:33 -0700 Subject: [PATCH 33/33] Describe the public Transpose recipes in the Rust guide Signed-off-by: Trent Nelson --- docs/rust/index.rst | 13 +++++++------ 1 file changed, 7 insertions(+), 6 deletions(-) diff --git a/docs/rust/index.rst b/docs/rust/index.rst index d6e7802cc427..4f235b07aa45 100644 --- a/docs/rust/index.rst +++ b/docs/rust/index.rst @@ -6,7 +6,7 @@ Rust development The Rust workspace contains per-thread payloads, CUDA group queries, and the native host compiler foundation for ``cuda.coop``. The ``cuda-coop`` crate supplies checked -``ThreadData`` values, physical/mapped hierarchy queries, and Direct block +``ThreadData`` values, physical/mapped hierarchy queries, and Direct/Transpose block Load/Store; see :doc:`payload`, :doc:`groups`, and :doc:`load_store`. Its private ``cuda-coop-codegen`` crate specializes CUB block operations and returns linkable device LTO-IR with exact temporary-storage requirements. It supports @@ -56,18 +56,19 @@ order, retaining initialized values in every slot, and uses the exact CUB storag layout. Compiler integrations must enforce collective participation, memory ownership, valid bounds, and synchronization before scratch reuse. -The separate validity-aware Direct recipes support all ten numeric payload +The separate validity-aware Direct and Transpose recipes support all ten numeric payload types: ``i8``, ``u8``, ``i16``, ``u16``, ``i32``, ``u32``, ``i64``, ``u64``, ``f32`` and ``f64``. ``Request::new_typed`` selects the exact type. These recipes -take opaque value storage and separate per-slot ``u8`` flags. Full and +take opaque value storage and a packed validity bitmap. Full and default-filled Load initialize every slot. Masked Load without a default copies only the loaded prefix and resets all other flags to zero. Full and partial -Store require flags equal to exactly one in their write domain and preserve -all payload bytes and flags. A failing Store traps; it is not transactional +Store require initialized slots in their write domain. Direct preserves all +payload bytes and validity bits; Transpose permutes both into striped order, +including slots outside its write prefix. A failing Store traps; it is not transactional across lanes. Partial Store forms destination pointers only within its valid prefix, including when the allocation is shorter than a complete tile. -Validity-aware Transpose remains under development. The public Direct API uses +The public Direct and Transpose API uses the same ``ThreadData`` representation without exposing an initialized slice or a public native layout.