From d1d7e5f87b3e4c89ed5a222f4f787f017c6d701b Mon Sep 17 00:00:00 2001 From: Trent Nelson Date: Mon, 19 Oct 2026 21:39:36 -0700 Subject: [PATCH 01/48] 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 00000000000..d5bde6e27fa --- /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 e1edae43a94..8906c9170e1 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 00000000000..23776f9573f --- /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 00000000000..b83d22266ac --- /dev/null +++ b/rust/.gitignore @@ -0,0 +1 @@ +/target/ diff --git a/rust/AGENTS.md b/rust/AGENTS.md new file mode 100644 index 00000000000..6493480b108 --- /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 00000000000..fdc6db4041d --- /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 00000000000..c90fafb1430 --- /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 00000000000..7476b2c6039 --- /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 00000000000..3a2cfe26364 --- /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 00000000000..7b910513322 --- /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 00000000000..298510237ea --- /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 00000000000..52daf38a400 --- /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 00000000000..4f8e16d38eb --- /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 00000000000..06d6cb6baf7 --- /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/48] 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 d5bde6e27fa..ad93cbf4b65 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/48] 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 fdc6db4041d..b1c2103ebcb 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 7476b2c6039..7a92fcaaded 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 298510237ea..872adf16c76 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 00000000000..09d57e186b6 --- /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 00000000000..d2b9d7f747a --- /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 00000000000..fc123a55f9d --- /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 00000000000..0d5afeb1bb0 --- /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 00000000000..76b2e828a01 --- /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 00000000000..5f555bfd695 --- /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 00000000000..e2393ccec57 --- /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 00000000000..ccf454ce08a --- /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 00000000000..a997768dc07 --- /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 00000000000..025d76c5fb6 --- /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 00000000000..6b357575759 --- /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/48] 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 fc123a55f9d..fb76581bf75 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 ccf454ce08a..eba29e10753 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/48] 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 23776f9573f..72be476dcb9 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 00000000000..ae2d0f55d26 --- /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 fdc6db4041d..650a0f45e4c 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 c90fafb1430..c6fd608c371 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 00000000000..1ca43ea3639 --- /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 00000000000..a7625da6c1b --- /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 00000000000..63e8f298760 --- /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 00000000000..ec97f008886 --- /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 00000000000..4d955cd2cbb --- /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 00000000000..8b24d6a7c2c --- /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/48] 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 00000000000..1d3f23f503c --- /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 fb76581bf75..07f3b0f1883 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 0d5afeb1bb0..450d33d5678 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 5f555bfd695..f0f464a2c76 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 00000000000..a0e57f30b2b --- /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 00000000000..c2df4f7d738 --- /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 e2393ccec57..eb584174339 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 eba29e10753..716937b82d5 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 a997768dc07..abb1aba763a 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 00000000000..f425c21babd --- /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 025d76c5fb6..019ca58eefc 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 6b357575759..a6644ccc769 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/48] 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 00000000000..8e411a01c50 --- /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 72be476dcb9..6250a0f5cf3 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 ae2d0f55d26..be5aebab8b9 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 650a0f45e4c..c0d1de3a583 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 1ca43ea3639..03e46ac0c20 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 00000000000..d3ea2e900b2 --- /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 ec97f008886..3d575be24c8 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 00000000000..74bffa8af4a --- /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 00000000000..2026e4b3bed --- /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 00000000000..7e29090be63 --- /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 00000000000..12ce62674ea --- /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/48] 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 1d3f23f503c..166ac4f7162 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 07f3b0f1883..3d95e7af549 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 eb584174339..ff096fc7ba9 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 716937b82d5..7b4b06eb74d 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 f425c21babd..9899bed7be6 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 a6644ccc769..52bfb966a21 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/48] 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 8e411a01c50..4c5050e0291 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 d3ea2e900b2..f1a1cdf3f5b 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 12ce62674ea..b7188d1def7 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/48] 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 23776f9573f..06b7c5d546d 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 7b910513322..dd55da52337 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 52daf38a400..fe0ca34e765 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 09d57e186b6..001a4b55d09 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 76b2e828a01..6c1abfb8573 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 5f555bfd695..5a6d130add0 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/48] 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 6250a0f5cf3..98d37567acb 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 f1a1cdf3f5b..97004caff67 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/48] 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 06b7c5d546d..2846d81e941 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 001a4b55d09..291e3f90004 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/48] 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 00000000000..6e63d79b9ac --- /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 d2b9d7f747a..3278c198b44 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 3d95e7af549..be7a15222ac 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 450d33d5678..2dce4d7120f 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 00000000000..e46c515f0d6 --- /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 f0f464a2c76..2f51b0c273b 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 00000000000..fe068a69bb7 --- /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 00000000000..2b39d7138a8 --- /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 52bfb966a21..c2899825874 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/48] 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 166ac4f7162..d67393be7fe 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 3d95e7af549..c47b65d48f9 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 00000000000..9646c1c74a3 --- /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 7b4b06eb74d..fb2f5eff8a1 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 9899bed7be6..9dcf388dc50 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 019ca58eefc..8f0145b72ba 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/48] 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 6e63d79b9ac..4a982dc6222 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 cb3e919abe5..8ae3934937a 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 2b39d7138a8..394c0da136c 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/48] 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 00000000000..f67f47730ab --- /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 cb3e919abe5..caeb786e742 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 2dce4d7120f..a20c32465f2 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 63ae2fe618d..4cd1eb1a8fd 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 00000000000..6b5478c358d --- /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 ff096fc7ba9..1d496140d00 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 fb2f5eff8a1..3f2b53b5f28 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 00000000000..9b7feab4bb0 --- /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 abb1aba763a..bce66cd3101 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 8f0145b72ba..77f9296ed30 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 00000000000..4c55b0de622 --- /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 c2899825874..2e4ede05486 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/48] 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 d67393be7fe..7689ec4308e 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/48] 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 25f80abf7c1..e2950e6e8e8 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 00000000000..d404f17e039 --- /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 872adf16c76..a16180ba536 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 291e3f90004..511f7a33bfe 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 00000000000..97b9399c70e --- /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 4f8e16d38eb..e61d2c2609c 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 cb3e919abe5..22afed6ab49 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 2dce4d7120f..d05b9378354 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 6c1abfb8573..6f0a40201f8 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 e46c515f0d6..d23d219467e 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 63ae2fe618d..df61f93fc98 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 00000000000..7891d885e42 --- /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 fe068a69bb7..0f96fdca727 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 00000000000..d28290bbe97 --- /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 c2899825874..b5176c2bad9 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 4d955cd2cbb..95f5170d63e 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 00000000000..1bbc9b447fc --- /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/48] 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 f67f47730ab..dde9b2417f8 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 a5a99baa5e8..e595bcc77fd 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 9b7feab4bb0..c1ab62f9323 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 4c55b0de622..b51cad4c48e 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/48] 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 4d955cd2cbb..1fef6629f29 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/48] 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 1fef6629f29..1992ec8939e 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 8b24d6a7c2c..88e8c24b4a9 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/48] 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 7689ec4308e..bb64773b308 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/48] 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 e2950e6e8e8..82d6bd1a140 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 00000000000..8515808b624 --- /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 e6a9256c8ed..e8b7b6d4af6 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 750b9fa4a00..41b532965f6 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 58d8590a91f..1d2f10a04c2 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 b4a21fbf539..cb14c141806 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 a0543a71b15..3118fb4bbaf 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 03e46ac0c20..0a903f2e1ed 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 00000000000..858fae4c024 --- /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 00000000000..74bffa8af4a --- /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 00000000000..7e39af8b87c --- /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 00000000000..0cf131efc56 --- /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 00000000000..1c6b5909485 --- /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 00000000000..986ac5ec32b --- /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 00000000000..3f534fcc3cc --- /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 3d575be24c8..a73e14c8592 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 00000000000..b73135d3c4f --- /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 00000000000..9ae7f1e54ab --- /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 11880d852c5..76203c84773 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 48da4cf37b3..a155164ad45 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 00000000000..d29d8c0b397 --- /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 00000000000..74bffa8af4a --- /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 00000000000..133f595cdf9 --- /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 00000000000..6dbdea6cee6 --- /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 00000000000..920e04b4c7d --- /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/48] 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 4a982dc6222..3a6b64af880 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 bb64773b308..c4fcf2e3ac2 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/48] 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 8515808b624..6de4481c36a 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 9ae7f1e54ab..4c0afdbc63e 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 d29d8c0b397..b881828fd73 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/48] 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 41b532965f6..a3f4e62bf98 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 1d2f10a04c2..c12d5b7e34c 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 e61d2c2609c..08bde41edc5 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/48] 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 e8b7b6d4af6..03bddcbec46 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 0a903f2e1ed..8bd4ee382a2 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 7e39af8b87c..f7068217c0a 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 0cf131efc56..442eaf871c8 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 133f595cdf9..5bcb0ea6a68 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 6dbdea6cee6..235f9d76fdc 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/48] 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 6de4481c36a..b0264afe105 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 858fae4c024..64adc8bb40e 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 986ac5ec32b..b587972831d 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 3f534fcc3cc..0f41cc7fdfa 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 a73e14c8592..c7536f88cf7 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 b73135d3c4f..094c83b7356 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 4c0afdbc63e..88e8a239e43 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 b881828fd73..e83c6dd7491 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 920e04b4c7d..142cfbe7bde 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 00000000000..8323f2223ed --- /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/48] 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 03bddcbec46..b092c7b6289 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 8bd4ee382a2..0f3c80862ca 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 f7068217c0a..161a9f9544e 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 442eaf871c8..258c8f4142b 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 5bcb0ea6a68..4e95c613b32 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 235f9d76fdc..2b2d125d2aa 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/48] 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 b0264afe105..be7ce6d1614 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 a3f4e62bf98..83851bb6cd2 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 094c83b7356..2b9b770f288 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 e83c6dd7491..e10ab4d21d1 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/48] 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 82d6bd1a140..d6e7802cc42 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 be7ce6d1614..35e407a7c9b 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 83851bb6cd2..dfc24b818e3 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 c12d5b7e34c..08972f32f85 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 c4fcf2e3ac2..7f422c751fc 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 08bde41edc5..a2abff03a48 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 142cfbe7bde..f93056d411d 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 8323f2223ed..6fc2cb9fbad 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 08fdb4b10cbe0255ed2f7b7dfbe3c54a326d83c3 Mon Sep 17 00:00:00 2001 From: Trent Nelson Date: Tue, 20 Oct 2026 12:35:33 -0700 Subject: [PATCH 32/48] 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 82d6bd1a140..0b5628041c4 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 @@ -53,18 +53,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. From 47861b64afa07476dcb17c1df18172f85f8bcaa2 Mon Sep 17 00:00:00 2001 From: Trent Nelson Date: Tue, 20 Oct 2026 12:38:37 -0700 Subject: [PATCH 33/48] 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 35e407a7c9b..74cbe55cafe 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 dfc24b818e3..3383bb5016a 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 a2abff03a48..1c782d59df3 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 34/48] 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 d6e7802cc42..4f235b07aa4 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. From 7e1f4d46344900cccc69e9a27ecef82aa4b4dd6c Mon Sep 17 00:00:00 2001 From: Trent Nelson Date: Tue, 20 Oct 2026 13:00:55 -0700 Subject: [PATCH 35/48] Add Rust Striped Load and Store Signed-off-by: Trent Nelson --- docs/rust/load_store.rst | 34 +- rust/cuda-coop-codegen/src/lib.rs | 36 +- rust/cuda-coop-codegen/src/oxide.rs | 14 +- rust/cuda-coop-codegen/tests/native.rs | 16 +- rust/cuda-coop/build.rs | 11 +- .../cuda-coop/examples/load_store/src/main.rs | 27 ++ rust/cuda-coop/src/lib.rs | 5 +- rust/cuda-coop/src/load_store.rs | 113 +++-- rust/cuda-coop/src/load_store/boundary.rs | 103 +++-- rust/cuda-coop/tests/check_load_store.py | 22 +- rust/cuda-coop/tests/load_store/src/main.rs | 7 + .../cuda-coop/tests/load_store/src/striped.rs | 413 ++++++++++++++++++ 12 files changed, 692 insertions(+), 109 deletions(-) create mode 100644 rust/cuda-coop/tests/load_store/src/striped.rs diff --git a/docs/rust/load_store.rst b/docs/rust/load_store.rst index 74cbe55cafe..237fb28e8ae 100644 --- a/docs/rust/load_store.rst +++ b/docs/rust/load_store.rst @@ -6,9 +6,10 @@ 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. 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 +first. Complete one-, two- and three-dimensional blocks support Direct, Striped +and Transpose algorithms. Direct and Transpose Load produce blocked payloads, +where each thread owns consecutive elements. Striped Load gives each thread +elements separated by the block size. 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 @@ -20,7 +21,7 @@ A complete program ------------------ The example is an independent Cargo workspace with a separate generic helper -crate. It copies three block tiles with Direct and Transpose at item counts 1 +crate. It copies three block tiles with Direct, Striped 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 @@ -40,7 +41,7 @@ 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 and includes the planned shared storage for Transpose. +Direct and Striped Load/Store and includes the planned shared storage for Transpose. .. literalinclude:: ../../rust/cuda-coop/examples/load_store/src/main.rs :language: rust @@ -48,7 +49,7 @@ Direct Load/Store and includes the planned shared storage for Transpose. 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: +Direct and Striped code may pass a shared payload borrow: .. literalinclude:: ../../rust/cuda-coop/examples/load_store/helper/src/lib.rs :language: rust @@ -66,7 +67,8 @@ 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 +``offset + t * ITEMS_PER_THREAD + i``; a Striped slot addresses +``offset + t + X*Y*Z * 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 @@ -83,12 +85,13 @@ 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. Direct Store preserves every payload value and validity flag, including +Direct and Striped Store accept ``&ThreadData``, a shared reborrow of ``&mut ThreadData``, +or a scalar numeric value. A scalar contributes one item per thread. Both preserve every payload value +and validity bit, including initialized slots outside its write domain. -Selecting Transpose -------------------- +Selecting an algorithm +---------------------- Select the algorithm through the trailing options. Selection is a Rust type, so each call compiles one algorithm; ordinary defaults continue to select Direct. @@ -100,6 +103,11 @@ so each call compiles one algorithm; ordinary defaults continue to select Direct coop::store(block, output, &mut values, coop::StoreOptions::default().valid_items(count).algorithm(coop::Transpose)); +Striped uses CUB's direct striped accesses without shared scratch. Select +``.algorithm(coop::Striped)`` on both operations to load and store a striped +payload. Its Store accepts a shared payload borrow and preserves the complete +payload, including a zero-count Store and initialized slots outside its prefix. + 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 @@ -107,7 +115,7 @@ 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. +uses a one-item temporary and accepts every algorithm. The compiler allocates the provider's complete shared storage, including the validity plane when needed, and inserts completion before reusing managed @@ -165,7 +173,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``, ``StoreValue``, ``Algorithm``, ``Direct`` and ``Transpose`` specify their full generic signatures and +``LoadSource``, ``StoreValue``, ``Algorithm``, ``Direct``, ``Striped`` 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-codegen/src/lib.rs b/rust/cuda-coop-codegen/src/lib.rs index 3383bb5016a..3d8c213bca7 100644 --- a/rust/cuda-coop-codegen/src/lib.rs +++ b/rust/cuda-coop-codegen/src/lib.rs @@ -11,11 +11,12 @@ //! # Supported specializations //! //! [`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 +//! items per thread, and [`Algorithm::Direct`], [`Algorithm::Striped`] or +//! [`Algorithm::Transpose`]. +//! The raw i32 recipes support all three algorithms. Their Store payloads must //! 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 +//! Masked Load resets validity outside its prefix; Direct/Striped Store preserves both //! arrays and checks initialization only in its write domain. Transpose Store //! replaces each lane's blocked payload and validity with striped items across //! the block, including when its write count is zero. @@ -86,18 +87,25 @@ impl From for Error { /// Result returned by native provider operations. pub type Result = std::result::Result; -/// CUB block layout algorithm; both variants produce/consume blocked payloads. +/// CUB block layout algorithm and per-thread item arrangement. #[derive(Clone, Copy, Debug, Eq, PartialEq)] pub enum Algorithm { /// Consecutive items per thread, without shared scratch or a reuse barrier. Direct, + /// Thread `t` accesses item `i` at `t + block_threads * i`, without scratch. + Striped, /// Coalesced access with shared-memory transposition to/from blocked layout. Transpose, } impl Algorithm { + fn scratch_free(self) -> bool { + matches!(self, Self::Direct | Self::Striped) + } + fn token(self) -> &'static str { match self { Self::Direct => "DIRECT", + Self::Striped => "STRIPED", Self::Transpose => "TRANSPOSE", } } @@ -258,7 +266,7 @@ impl Request { /// Validate an exact typed block specialization. /// - /// The five validity-aware recipes support every [`Dtype`] with either algorithm. + /// The five validity-aware recipes support every [`Dtype`] with every algorithm. /// Legacy raw recipes retain their i32 ABI and algorithm support. /// /// # Errors @@ -396,7 +404,7 @@ pub struct Layout { /// 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, +/// the provider does not add a trailing reuse barrier. Direct and Striped ignore scratch, /// so a null pointer is permitted for that argument only. All providers return /// unit; loads fill payloads in place and stores write memory. /// @@ -406,7 +414,7 @@ pub struct Layout { /// 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 +/// borrows both arrays; Direct/Striped Store receives const pointers and borrows them /// read-only. Transpose Store exclusively borrows both arrays and returns their /// permutation: destination `(lane, item)` receives source /// `((block_threads * item + lane) / items_per_thread, @@ -429,7 +437,7 @@ pub struct Specialization { /// Actual CUB TempStorage sizeof/alignof from the same NVRTC program. pub cub_layout: Layout, /// Complete shared region, including any validity plane and padding. - /// Direct uses zero bytes despite its empty C++ type's sizeof=1. + /// Direct and Striped use zero bytes despite their empty C++ types' sizeof=1. pub scratch: Layout, /// Whether this specialization accesses no shared scratch or scratch-reuse barrier. pub scratch_free: bool, @@ -586,7 +594,7 @@ impl Provider { size: 0, alignment: 1, }, - scratch_free: request.algorithm == Algorithm::Direct, + scratch_free: request.algorithm.scratch_free(), mutates_payload: !request.operation.is_load() && request.algorithm == Algorithm::Transpose, }); @@ -656,7 +664,7 @@ static_assert(alignof(scratch_storage) == alignof(primitive::TempStorage));"# controls.push_str(&format!(", {dtype} default_value")); operands.push_str(", default_value"); } - let constructor = if request.algorithm == Algorithm::Direct { + let constructor = if request.algorithm.scratch_free() { "primitive operation;" } else { "primitive operation{*reinterpret_cast(scratch)};" @@ -679,7 +687,11 @@ static_assert(alignof(scratch_storage) == alignof(primitive::TempStorage));"# op, Operation::LoadMaskedValidity | Operation::StorePartialValidity ) { - format!("tid * {items} + i < valid") + if request.algorithm == Algorithm::Striped { + format!("tid + {threads} * i < valid") + } else { + format!("tid * {items} + i < valid") + } } else { "true".to_owned() }; @@ -704,7 +716,7 @@ static_assert(alignof(scratch_storage) == alignof(primitive::TempStorage));"# ) } else { // The general Store permutes all source slots, including initialized - // slots outside its write domain. Direct only needs its write domain. + // slots outside its write domain. Direct/Striped only need their write domain. let copy_domain = if transpose { "true" } else { &domain }; let prepare = if mask_plane { format!( diff --git a/rust/cuda-coop-codegen/src/oxide.rs b/rust/cuda-coop-codegen/src/oxide.rs index 08972f32f85..511b5b80829 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/Transpose validity recipes +//! The protocol supports the raw i32 and typed Direct/Striped/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. @@ -134,6 +134,7 @@ impl SpecializationRequest { }; let algorithm = match self.algorithm.as_str() { "direct" => Algorithm::Direct, + "striped" => Algorithm::Striped, "transpose" => Algorithm::Transpose, _ => { return Err(Error::message(format!( @@ -193,7 +194,7 @@ 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; + operation.validity() && !operation.is_load() && request.algorithm().scratch_free(); let mut arguments = vec![ if operation.is_load() { Argument::ConstPointer(dtype) @@ -289,6 +290,7 @@ pub fn registration( operation: operation.token(), algorithm: match algorithm { Algorithm::Direct => "direct", + Algorithm::Striped => "striped", Algorithm::Transpose => "transpose", }, }) @@ -484,7 +486,7 @@ mod tests { ] { let scalar = Scalar::from(dtype); assert_eq!(scalar.dtype().unwrap(), dtype); - for algorithm in ["direct", "transpose"] { + for algorithm in ["direct", "striped", "transpose"] { let request = SpecializationRequest { operation: "load_filled_validity_bits_v1".into(), algorithm: algorithm.into(), @@ -510,16 +512,16 @@ mod tests { Operation::StoreFullValidity, Operation::StorePartialValidity, ] { - for algorithm in [Algorithm::Direct, Algorithm::Transpose] { + for algorithm in [Algorithm::Direct, Algorithm::Striped, 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 { + if algorithm.scratch_free() { Argument::ConstPointer(Scalar::Signed(16)) } else { Argument::MutPointer(Scalar::Signed(16)) }, - if algorithm == Algorithm::Direct { + if algorithm.scratch_free() { Argument::ConstPointer(Scalar::Unsigned(8)) } else { Argument::MutPointer(Scalar::Unsigned(8)) diff --git a/rust/cuda-coop-codegen/tests/native.rs b/rust/cuda-coop-codegen/tests/native.rs index 1c782d59df3..656d79751b4 100644 --- a/rust/cuda-coop-codegen/tests/native.rs +++ b/rust/cuda-coop-codegen/tests/native.rs @@ -81,7 +81,11 @@ fn native_typed_validity_and_determinism() -> Result<(), Box Result<(), Box Result<(), Box Result<(), Box Result<(), Box Result<(), Box( + 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::Striped, + ); + } } fn main() -> Result<(), Box> { @@ -90,6 +115,8 @@ fn main() -> Result<(), Box> { } run!(1, prepare_copy, copy_async_owned); run!(4, prepare_copy, copy_async_owned); + run!(1, prepare_striped, striped_async_owned); + run!(4, prepare_striped, striped_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 c7536f88cf7..fea98f683a0 100644 --- a/rust/cuda-coop/src/lib.rs +++ b/rust/cuda-coop/src/lib.rs @@ -30,7 +30,7 @@ //! [`this_block`] and the other dimension-free group constructors describe //! the executing thread's current launch. Queries grant no collective //! participation or synchronization guarantee. [`load`] and [`store`] support -//! Direct and Transpose on complete blocks with compiler-checked participation +//! Direct, Striped 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. @@ -45,7 +45,8 @@ mod thread_data; pub use alignment::{Alignment, SupportedAlignment}; pub use cuda_device::{CollectiveTiles, DisjointSlice}; pub use load_store::{ - Algorithm, Direct, LoadOptions, LoadSource, StoreOptions, StoreValue, Transpose, load, store, + Algorithm, Direct, LoadOptions, LoadSource, StoreOptions, StoreValue, Striped, Transpose, load, + store, }; pub use thread_data::{Numeric, ThreadData, ThreadDataLike}; diff --git a/rust/cuda-coop/src/load_store.rs b/rust/cuda-coop/src/load_store.rs index 5960d27edbe..c1bb82f526f 100644 --- a/rust/cuda-coop/src/load_store.rs +++ b/rust/cuda-coop/src/load_store.rs @@ -11,6 +11,13 @@ mod boundary; #[derive(Clone, Copy, Debug, Default)] pub struct Direct; +/// Direct striped access without shared scratch. +/// +/// Thread rank `t` and item `i` address `t + group.count() * i`. Load produces +/// this striped arrangement; Store consumes it and preserves the payload. +#[derive(Clone, Copy, Debug, Default)] +pub struct Striped; + /// Coalesced access through CUB's shared-memory block exchange. /// /// Load produces blocked payloads. Store permutes payload values and validity @@ -36,12 +43,16 @@ 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 Striped {} impl Algorithm for Transpose {} impl sealed::Algorithm for Direct { - const TRANSPOSE: bool = false; + const KIND: u8 = sealed::DIRECT; +} +impl sealed::Algorithm for Striped { + const KIND: u8 = sealed::STRIPED; } impl sealed::Algorithm for Transpose { - const TRANSPOSE: bool = true; + const KIND: u8 = sealed::TRANSPOSE; } /// Controls for a [`load`], using [`Direct`] by default. @@ -161,19 +172,21 @@ impl StoreOptions { } } -/// Load a block tile in blocked order into an existing per-thread payload. +/// Load a block tile 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. +/// `group.count() * ITEMS` elements. Direct and Transpose give each thread +/// consecutive items; Striped gives it items separated by the block size. /// [`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. [`Direct`] uses no shared scratch; [`Transpose`] +/// The operation returns unit. [`Direct`] and [`Striped`] use no shared scratch; [`Transpose`] /// uses compiler-planned shared storage and returns the same blocked layout. /// -/// 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. +/// Complete one-, two- and three-dimensional blocks support all three algorithms. +/// Thread rank is x-fastest: `t = x + X * (y + Y * z)` for shape `[X, Y, Z]`. +/// Slot `i` addresses `t * ITEMS + i` for Direct/Transpose Load and +/// `t + group.count() * i` for Striped Load, relative to the selected offset. /// 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 @@ -210,13 +223,15 @@ pub fn load, const ITEMS: usize, const ALIGN: usize source.load(group, output, options); } -/// Store a block tile from a blocked per-thread payload. +/// Store a block tile from a per-thread payload. /// /// `T` is inferred from the destination and value. [`StoreValue`] accepts a -/// shared or mutable borrow of [`ThreadData`] for [`Direct`], a mutable borrow -/// for [`Transpose`], or a numeric scalar (one item per thread) for either. +/// shared or mutable borrow of [`ThreadData`] for [`Direct`] and [`Striped`], a +/// mutable borrow for [`Transpose`], or a numeric scalar (one item per thread) +/// for any algorithm. /// `CAPACITY` describes the carrier's ownership grain; it need not equal the -/// payload extent. Direct preserves values and flags and uses no shared scratch. +/// payload extent. Direct consumes blocked items; Striped consumes striped items. +/// Both preserve every value and initialized bit and use 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`. @@ -320,9 +335,10 @@ impl LoadSource for &mut DisjointSlice<'_, T, Col /// A sealed per-thread value accepted by [`store`] for type `T` and algorithm `A`. /// -/// 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 +/// Numeric scalars and mutable [`ThreadData`] borrows support all algorithms. +/// Shared payload borrows support [`Direct`] and [`Striped`]. A scalar contributes +/// one item per thread. Direct and Striped preserve every value and bit; +/// [`Transpose`] mutates their /// positions as described by [`store`]. pub trait StoreValue: sealed::StoreValue { #[doc(hidden)] @@ -334,26 +350,33 @@ pub trait StoreValue: sealed::StoreValue 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); - } +macro_rules! shared_store_value { + ($algorithm:ty) => { + 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<$algorithm>, + ) { + boundary::store(destination, self, options); + } + } + }; } +shared_store_value!(Direct); +shared_store_value!(Striped); impl sealed::StoreValue for &mut ThreadData where @@ -373,10 +396,17 @@ where destination: &mut DisjointSlice<'_, T, CollectiveTiles>, options: StoreOptions, ) { - if A::TRANSPOSE { - boundary::store_transpose(destination, self, options.algorithm(Transpose)); - } else { - boundary::store(destination, self, options.algorithm(Direct)); + match A::KIND { + sealed::DIRECT => { + boundary::store(destination, self, options.algorithm(Direct)); + } + sealed::STRIPED => { + boundary::store(destination, self, options.algorithm(Striped)); + } + sealed::TRANSPOSE => { + boundary::store_transpose(destination, self, options.algorithm(Transpose)); + } + _ => panic!(), } } } @@ -398,7 +428,12 @@ impl StoreValue for T { mod sealed { pub trait LoadSource {} pub trait StoreValue {} + // Scalar tags let the compiler prune unused native operation arms before + // planning scratch. Only the sealed selector implementations assign them. + pub const DIRECT: u8 = 0; + pub const STRIPED: u8 = 1; + pub const TRANSPOSE: u8 = 2; pub trait Algorithm { - const TRANSPOSE: bool; + const KIND: u8; } } diff --git a/rust/cuda-coop/src/load_store/boundary.rs b/rust/cuda-coop/src/load_store/boundary.rs index 88e8a239e43..fc59bdfce71 100644 --- a/rust/cuda-coop/src/load_store/boundary.rs +++ b/rust/cuda-coop/src/load_store/boundary.rs @@ -1,7 +1,10 @@ // Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception -use super::{Algorithm, LoadOptions, StoreOptions, Transpose}; +use super::{ + Algorithm, LoadOptions, StoreOptions, Transpose, + sealed::{DIRECT, STRIPED, TRANSPOSE}, +}; use crate::{Alignment, Numeric, SupportedAlignment, ThreadData}; use cuda_device::{CollectiveTiles, DisjointSlice, library, thread}; @@ -18,6 +21,12 @@ registration!(load_filled, "COOP_VALIDITY_LOAD_FILLED"); registration!(store_full, "COOP_VALIDITY_STORE_FULL"); registration!(store_partial, "COOP_VALIDITY_STORE_PARTIAL"); +registration!(striped_load_full, "COOP_STRIPED_VALIDITY_LOAD_FULL"); +registration!(striped_load_masked, "COOP_STRIPED_VALIDITY_LOAD_MASKED"); +registration!(striped_load_filled, "COOP_STRIPED_VALIDITY_LOAD_FILLED"); +registration!(striped_store_full, "COOP_STRIPED_VALIDITY_STORE_FULL"); +registration!(striped_store_partial, "COOP_STRIPED_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"); @@ -97,19 +106,24 @@ pub(super) unsafe fn load< // 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, - (), - $controls, - payload, - ) - } else { - library::checked_block_resource_call_v1::( + ($direct:ident, $striped:ident, $transpose:ident, $controls:expr) => { + match A::KIND { + DIRECT => library::checked_block_resource_call_v1::( $direct, memory, $controls, payload, - ) + ), + STRIPED => library::checked_block_resource_call_v1::( + $striped, memory, $controls, payload, + ), + TRANSPOSE => { + library::checked_block_storage_call_v1::( + $transpose, + memory, + (), + $controls, + payload, + ) + } + _ => panic!(), } }; } @@ -119,12 +133,23 @@ pub(super) unsafe fn load< // 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,)), + (None, None) => call!( + load_full, + striped_load_full, + transpose_load_full, + (pointer,) + ), (Some(count), None) => { - call!(load_masked, transpose_load_masked, (pointer, count as i32)) + call!( + load_masked, + striped_load_masked, + transpose_load_masked, + (pointer, count as i32) + ) } (Some(count), Some(default)) => call!( load_filled, + striped_load_filled, transpose_load_filled, (pointer, count as i32, default) ), @@ -134,10 +159,16 @@ pub(super) unsafe fn load< } #[inline(always)] -pub(super) fn store( +pub(super) fn store< + T: Numeric, + const ITEMS: usize, + const ALIGN: usize, + const CAPACITY: usize, + A: Algorithm, +>( destination: &mut DisjointSlice<'_, T, CollectiveTiles>, value: &ThreadData, - options: StoreOptions, + options: StoreOptions, ) where Alignment: SupportedAlignment, { @@ -155,18 +186,36 @@ pub(super) fn store { + match A::KIND { + DIRECT => { + library::checked_block_resource_call_v1::( + $direct, + destination, + $controls, + payload, + ) + } + STRIPED => { + library::checked_block_resource_call_v1::( + $striped, + destination, + $controls, + payload, + ) + } + // Only the shared-preserving Store implementations call this boundary. + _ => panic!(), + } + }; + } match options.valid_items { - None => library::checked_block_resource_call_v1::( - store_full, - destination, - (pointer,), - payload, - ), - Some(count) => library::checked_block_resource_call_v1::( + None => call!(store_full, striped_store_full, (pointer,)), + Some(count) => call!( store_partial, - destination, - (pointer, count as i32), - payload, + striped_store_partial, + (pointer, count as i32) ), } } diff --git a/rust/cuda-coop/tests/check_load_store.py b/rust/cuda-coop/tests/check_load_store.py index e10ab4d21d1..8f08bb17555 100644 --- a/rust/cuda-coop/tests/check_load_store.py +++ b/rust/cuda-coop/tests/check_load_store.py @@ -115,7 +115,7 @@ cases.update( shared_transpose=( load + transpose_store.replace("&mut values", "&values"), - "expected `Direct`, found `Transpose`", + "StoreValue", ), transpose_managed_reuse=( transpose_load + transpose_store + transpose_store, @@ -152,6 +152,26 @@ ), ) +striped_load = load.replace( + ".oob_default(0)", ".oob_default(0).algorithm(coop::Striped)" +) +striped_store = store.replace( + ".valid_items(valid)", ".valid_items(valid).algorithm(coop::Striped)" +) +cases.update( + striped_shared=(striped_load + striped_store, None), + striped_varying_count=( + striped_load.replace( + ".valid_items(valid)", ".valid_items(thread::threadIdx_x() as usize)" + ), + "not proved block-uniform", + ), + striped_missing_completion=( + striped_load + striped_store + reload, + "requires a full-block barrier", + ), +) + 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 f93056d411d..9714f1a474c 100644 --- a/rust/cuda-coop/tests/load_store/src/main.rs +++ b/rust/cuda-coop/tests/load_store/src/main.rs @@ -2,6 +2,7 @@ // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception use coop::{Numeric, ThreadData, this_block}; +mod striped; mod transpose; use cuda_coop as coop; @@ -365,6 +366,12 @@ mod kernels { fn main() -> Result<(), Box> { if let Some(argument) = std::env::args().nth(1) { + if argument == "striped" { + return striped::run(None); + } + if argument.starts_with("striped-") { + return striped::run(Some(&argument)); + } if argument == "transpose" { return transpose::run(None); } diff --git a/rust/cuda-coop/tests/load_store/src/striped.rs b/rust/cuda-coop/tests/load_store/src/striped.rs new file mode 100644 index 00000000000..fa9b424a5a7 --- /dev/null +++ b/rust/cuda-coop/tests/load_store/src/striped.rs @@ -0,0 +1,413 @@ +// Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +//! Striped payload mapping and preservation, observed independently of roundtrip I/O. +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::*; + + #[kernel] + #[launch_contract(domain = 1, block = (33, 1, 1), dynamic_shared = inferred)] + pub fn copy_33( + input: &[i32], + mut output: DisjointSlice>, + ) { + let block = coop::this_block(); + let lane = thread::threadIdx_x() as usize; + let offset = thread::blockIdx_x() as usize * 33 * ITEMS_PER_THREAD; + let mut values = coop::ThreadData::<_, ITEMS_PER_THREAD>::new(); + coop::load( + block, + input, + &mut values, + coop::LoadOptions::default() + .offset(offset) + .algorithm(coop::Striped), + ); + let before = values; + coop::store( + block, + &mut output, + &values, + coop::StoreOptions::default() + .algorithm(coop::Striped) + .offset(offset), + ); + let mut observations_ok = true; + let mut item = 0; + while item < ITEMS_PER_THREAD { + let expected = input[offset + lane + 33 * item]; + observations_ok &= before.get(item) == Some(expected); + observations_ok &= values.get(item) == before.get(item); + item += 1; + } + assert!(observations_ok); + } + + fn exercise( + input: &[T], + output: &mut DisjointSlice>, + valid: usize, + shift: usize, + default: T, + mode: u32, + ) { + let block = coop::this_block(); + 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 = coop::ThreadData::<_, ITEMS>::new(); + match mode { + 0 => coop::load( + block, + input, + &mut values, + coop::LoadOptions::default() + .algorithm(coop::Striped) + .offset(offset), + ), + 1 | 5 | 6 => { + // A second masked Load must clear validity left by a full Load. + coop::load( + block, + input, + &mut values, + coop::LoadOptions::default() + .algorithm(coop::Striped) + .offset(offset), + ); + coop::load( + block, + input, + &mut values, + coop::LoadOptions::default() + .algorithm(coop::Striped) + .offset(offset) + .valid_items(valid), + ); + } + 2 | 4 => coop::load( + block, + input, + &mut values, + coop::LoadOptions::default() + .algorithm(coop::Striped) + .offset(offset) + .valid_items(valid) + .oob_default(default), + ), + 7 => {} + _ => { + let mut item = 0; + while item < ITEMS { + let selected = lane + threads * item < valid; + if (selected || (lane + item) & 1 == 0) + && !(mode == 8 && lane + threads * item == valid - 1) + { + values.write(item, default); + } + item += 1; + } + } + } + let before = values; + if mode == 0 || mode == 2 { + coop::store( + block, + &mut *output, + &values, + coop::StoreOptions::default() + .algorithm(coop::Striped) + .offset(offset), + ); + } else { + let written = if mode == 5 { + valid / 2 + } else if mode == 6 { + valid + 1 + } else { + valid + }; + coop::store( + block, + &mut *output, + &values, + coop::StoreOptions::default() + .algorithm(coop::Striped) + .offset(offset) + .valid_items(written), + ); + } + let mut observations_ok = true; + let mut item = 0; + while item < ITEMS { + let position = lane + threads * item; + let expected = match mode { + 3 | 8 => { + if position < valid || (lane + item) & 1 == 0 { + Some(default) + } else { + None + } + } + 7 => None, + _ => { + if position < valid || mode == 0 { + Some(input[offset + position]) + } else if mode == 2 || mode == 4 { + Some(default) + } else { + None + } + } + }; + observations_ok &= before.get(item) == expected; + observations_ok &= values.get(item) == before.get(item); + item += 1; + } + if mode == 3 || mode == 4 { + thread::sync_threads(); + // A repeated zero Store preserves arbitrary initialized/invalid slots. + coop::store( + block, + &mut *output, + &values, + coop::StoreOptions::default() + .algorithm(coop::Striped) + .offset(offset) + .valid_items(0), + ); + let mut item = 0; + while item < ITEMS { + observations_ok &= values.get(item) == before.get(item); + item += 1; + } + } + assert!(observations_ok); + } + + #[kernel] + #[launch_contract(domain = 1, block = (33, 1, 1), dynamic_shared = inferred)] + pub fn partial_33( + input: &[T], + mut output: DisjointSlice>, + valid: usize, + shift: usize, + default: T, + mode: u32, + ) { + exercise::<_, ITEMS>(input, &mut output, valid, shift, default, mode); + } + + #[kernel] + #[launch_contract(domain = 3, block = (8, 4, 4), dynamic_shared = inferred)] + pub fn partial_3d( + input: &[T], + mut output: DisjointSlice>, + valid: usize, + shift: usize, + default: T, + mode: u32, + ) { + exercise::<_, ITEMS>(input, &mut output, valid, shift, default, mode); + } + + #[kernel] + #[launch_contract(domain = 1, block = (33, 1, 1), dynamic_shared = inferred)] + pub fn scalar(mut output: DisjointSlice>, valid: usize) { + let offset = thread::blockIdx_x() as usize * 33; + coop::store( + coop::this_block(), + &mut output, + 17_i32, + coop::StoreOptions::default() + .algorithm(coop::Striped) + .offset(offset) + .valid_items(valid), + ); + } +} + +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, LaunchConfig3D}; + init_device_contexts(0, 1)?; + let context = with_cuda_context(0, |context| context.clone())?; + let stream = context.default_stream(); + // SAFETY: embedded kernels and typed declarations compile together. + let module = unsafe { kernels::load_async(0)? }; + if let Some(trap) = trap { + let mode = if trap == "striped-larger" { 6 } else { 8 }; + let plan = module.prepare_partial_33::(LaunchConfig1D::auto_shared(1, 33))?; + let source = DeviceBuffer::from_host(&stream, &vec![7_i32; 132])?; + let target = DeviceBuffer::from_host(&stream, &vec![0_i32; 132])?; + module + .partial_33_async_owned::(&plan, source, target, 34, 0, -7, mode) + .sync()?; + return Err("invalid Striped Store unexpectedly completed".into()); + } + let input: Vec = (0..2 * 33 * 4).map(|i| i * 7 + 3).collect(); + let source = DeviceBuffer::from_host(&stream, &input)?; + let target = DeviceBuffer::from_host(&stream, &vec![-1_i32; input.len()])?; + let plan = module.prepare_copy_33::<4>(LaunchConfig1D::auto_shared(2, 33))?; + assert_eq!(plan.__raw_config().shared_mem_bytes, 0); + let (_, target) = module + .copy_33_async_owned::<4, _, _>(&plan, source, target) + .sync()?; + let mut output = vec![0_i32; input.len()]; + target.copy_to_host(&stream, &mut output)?; + stream.synchronize()?; + assert_eq!(output, input); + let mut cases = 1; + macro_rules! run { + ($t:ty, $items:literal, $threads:literal, $prepare:ident, $launch:ident, $config:expr, $extra:expr) => {{ + let tile = $threads * $items; + let plan = module.$prepare::<$t, $items>($config)?; + assert_eq!(plan.__raw_config().shared_mem_bytes, 0); + let prefix = ($threads + 1).min(tile); + let mut selected = vec![ + (0, tile, 0), + (1, prefix, 0), + (2, 1, 0), + (3, prefix, 0), + (5, prefix, 0), + ]; + if $extra { + for valid in [0, 1, $threads - 1, $threads, prefix, tile] { + for mode in [1, 2, 3, 4, 5] { + selected.push((mode, valid, 0)); + } + } + for mode in [1, 3, 4, 5] { + selected.push((mode, tile - 7, 7)); + } + selected.push((7, 0, 0)); + } + selected.sort_unstable(); + selected.dedup(); + for (mode, valid, shift) in selected { + let written = if mode == 0 || mode == 2 { + tile + } else if mode == 5 { + valid / 2 + } else { + valid + }; + // Keep the final allocation shorter than the full tile when + // counted Store writes a prefix. Striped per-lane iterator + // formation must not create an out-of-allocation pointer. + let output_len = $threads * 4 + shift + written; + let input: Vec<$t> = (0..2 * $threads * 4 + shift) + .map(|i| (i % 101) as $t) + .collect(); + let default = 111 as $t; + let sentinel = 123 as $t; + let source = DeviceBuffer::from_host(&stream, &input)?; + let target = DeviceBuffer::from_host(&stream, &vec![sentinel; output_len])?; + let (_, target) = module + .$launch::<$t, $items, _, _>(&plan, source, target, valid, shift, default, mode) + .sync()?; + let mut output = vec![sentinel; output_len]; + target.copy_to_host(&stream, &mut output)?; + stream.synchronize()?; + for (index, actual) in output.into_iter().enumerate() { + let local = index % ($threads * 4); + let expected = if local >= shift && local - shift < written { + if mode == 3 || (mode == 2 && local - shift >= valid) { + default + } else { + input[index] + } + } else { + sentinel + }; + assert_eq!( + actual, + expected, + "type={} items={} mode={mode} valid={valid} shift={shift} index={index}", + stringify!($t), + $items + ); + } + cases += 1; + } + }}; + } + macro_rules! width { + ($t:ty) => { + run!( + $t, + 1, + 33, + prepare_partial_33, + partial_33_async_owned, + LaunchConfig1D::auto_shared(2, 33), + false + ); + run!( + $t, + 4, + 33, + prepare_partial_33, + partial_33_async_owned, + LaunchConfig1D::auto_shared(2, 33), + false + ); + }; + } + width!(i8); + width!(u8); + width!(i16); + width!(u16); + width!(i32); + width!(u32); + width!(i64); + width!(u64); + width!(f32); + width!(f64); + run!( + i32, + 4, + 33, + prepare_partial_33, + partial_33_async_owned, + LaunchConfig1D::auto_shared(2, 33), + true + ); + run!( + i32, + 4, + 128, + prepare_partial_3d, + partial_3d_async_owned, + LaunchConfig3D::auto_shared((1, 2, 1), (8, 4, 4)), + true + ); + for valid in [0, 1, 32, 33] { + let plan = module.prepare_scalar(LaunchConfig1D::auto_shared(2, 33))?; + assert_eq!(plan.__raw_config().shared_mem_bytes, 0); + let target = DeviceBuffer::from_host(&stream, &vec![-1_i32; 33 + valid])?; + let target = module.scalar_async_owned(&plan, target, valid).sync()?; + let mut output = vec![0_i32; 33 + valid]; + target.copy_to_host(&stream, &mut output)?; + stream.synchronize()?; + for (i, value) in output.into_iter().enumerate() { + assert_eq!(value, if i % 33 < valid { 17 } else { -1 }); + } + cases += 1; + } + println!("PASS: {cases} Striped cases with observed payload mapping and validity"); + Ok(()) +} From 193515598d12439c75b768e494a52670d040c077 Mon Sep 17 00:00:00 2001 From: Trent Nelson Date: Tue, 20 Oct 2026 13:44:23 -0700 Subject: [PATCH 36/48] Qualify Striped tests in the combined device image Signed-off-by: Trent Nelson --- docs/rust/load_store.rst | 10 ++++++++-- rust/cuda-coop/tests/load_store/src/striped.rs | 10 ++++++---- 2 files changed, 14 insertions(+), 6 deletions(-) diff --git a/docs/rust/load_store.rst b/docs/rust/load_store.rst index 237fb28e8ae..b495bd270d3 100644 --- a/docs/rust/load_store.rst +++ b/docs/rust/load_store.rst @@ -49,7 +49,6 @@ Direct and Striped Load/Store and includes the planned shared storage for Transp The reusable helper preserves the source element type. Its algorithm-generic form borrows the payload mutably because Transpose Store permutes it. -Direct and Striped code may pass a shared payload borrow: .. literalinclude:: ../../rust/cuda-coop/examples/load_store/helper/src/lib.rs :language: rust @@ -88,7 +87,7 @@ those slots to be initialized. Direct and Striped Store accept ``&ThreadData``, a shared reborrow of ``&mut ThreadData``, or a scalar numeric value. A scalar contributes one item per thread. Both preserve every payload value and validity bit, including -initialized slots outside its write domain. +initialized slots outside their write domain. Selecting an algorithm ---------------------- @@ -108,6 +107,13 @@ Striped uses CUB's direct striped accesses without shared scratch. Select payload. Its Store accepts a shared payload borrow and preserves the complete payload, including a zero-count Store and initialized slots outside its prefix. +.. code-block:: rust + + coop::load(block, input, &mut values, + coop::LoadOptions::default().algorithm(coop::Striped)); + coop::store(block, output, &values, + coop::StoreOptions::default().algorithm(coop::Striped)); + 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 diff --git a/rust/cuda-coop/tests/load_store/src/striped.rs b/rust/cuda-coop/tests/load_store/src/striped.rs index fa9b424a5a7..003de60cc87 100644 --- a/rust/cuda-coop/tests/load_store/src/striped.rs +++ b/rust/cuda-coop/tests/load_store/src/striped.rs @@ -225,7 +225,7 @@ mod kernels { #[kernel] #[launch_contract(domain = 1, block = (33, 1, 1), dynamic_shared = inferred)] - pub fn scalar(mut output: DisjointSlice>, valid: usize) { + pub fn striped_scalar(mut output: DisjointSlice>, valid: usize) { let offset = thread::blockIdx_x() as usize * 33; coop::store( coop::this_block(), @@ -256,7 +256,7 @@ pub fn run(trap: Option<&str>) -> Result<(), Box> { module .partial_33_async_owned::(&plan, source, target, 34, 0, -7, mode) .sync()?; - return Err("invalid Striped Store unexpectedly completed".into()); + panic!("invalid Striped Store unexpectedly completed"); } let input: Vec = (0..2 * 33 * 4).map(|i| i * 7 + 3).collect(); let source = DeviceBuffer::from_host(&stream, &input)?; @@ -396,10 +396,12 @@ pub fn run(trap: Option<&str>) -> Result<(), Box> { true ); for valid in [0, 1, 32, 33] { - let plan = module.prepare_scalar(LaunchConfig1D::auto_shared(2, 33))?; + let plan = module.prepare_striped_scalar(LaunchConfig1D::auto_shared(2, 33))?; assert_eq!(plan.__raw_config().shared_mem_bytes, 0); let target = DeviceBuffer::from_host(&stream, &vec![-1_i32; 33 + valid])?; - let target = module.scalar_async_owned(&plan, target, valid).sync()?; + let target = module + .striped_scalar_async_owned(&plan, target, valid) + .sync()?; let mut output = vec![0_i32; 33 + valid]; target.copy_to_host(&stream, &mut output)?; stream.synchronize()?; From aaf03fcae0c1ef6ac40d45f0ccff8d5c538aa87e Mon Sep 17 00:00:00 2001 From: Trent Nelson Date: Tue, 20 Oct 2026 13:46:09 -0700 Subject: [PATCH 37/48] Include Striped in the Rust development overview Signed-off-by: Trent Nelson --- docs/rust/index.rst | 17 +++++++++-------- 1 file changed, 9 insertions(+), 8 deletions(-) diff --git a/docs/rust/index.rst b/docs/rust/index.rst index 4f235b07aa4..02bb644f791 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/Transpose block +``ThreadData`` values, physical/mapped hierarchy queries, and Direct, Striped and 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 @@ -38,7 +38,7 @@ these forms: - Behavior - Payload requirement * - Full Load - - Fills every payload slot in blocked order + - Fills every payload slot in the algorithm's order - Writable storage * - Default-filled Load - Loads a valid prefix and fills the suffix with a supplied value @@ -50,25 +50,26 @@ these forms: - 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 +Direct, Striped and Transpose algorithms are supported. Direct and Striped access +no shared scratch. Striped loads and stores striped payload order. +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. -The separate validity-aware Direct and Transpose recipes support all ten numeric payload +The separate validity-aware 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 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 initialized slots in their write domain. Direct preserves all +Store require initialized slots in their write domain. Direct and Striped preserve 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. -The public Direct and Transpose API uses +The public Load/Store API uses the same ``ThreadData`` representation without exposing an initialized slice or a public native layout. @@ -106,7 +107,7 @@ Native regression checks are opt-in: 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 +is ``compute_120``. These checks compile the supported algorithm 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. From 8bc907b93cb402933e7f64bcc93f3c52682b47f4 Mon Sep 17 00:00:00 2001 From: Trent Nelson Date: Tue, 20 Oct 2026 13:23:52 -0700 Subject: [PATCH 38/48] Add vectorized Rust block Load and Store Signed-off-by: Trent Nelson --- docs/rust/index.rst | 49 +-- docs/rust/load_store.rst | 23 +- rust/cuda-coop-codegen/src/lib.rs | 49 ++- rust/cuda-coop-codegen/src/oxide.rs | 13 +- rust/cuda-coop-codegen/tests/native.rs | 26 +- rust/cuda-coop/build.rs | 1 + .../examples/load_store/helper/src/lib.rs | 59 ++- .../cuda-coop/examples/load_store/src/main.rs | 27 ++ rust/cuda-coop/src/lib.rs | 6 +- rust/cuda-coop/src/load_store.rs | 36 +- rust/cuda-coop/src/load_store/boundary.rs | 39 +- rust/cuda-coop/tests/check_load_store.py | 21 + rust/cuda-coop/tests/load_store/src/main.rs | 7 + .../tests/load_store/src/vectorize.rs | 378 ++++++++++++++++++ 14 files changed, 649 insertions(+), 85 deletions(-) create mode 100644 rust/cuda-coop/tests/load_store/src/vectorize.rs diff --git a/docs/rust/index.rst b/docs/rust/index.rst index 02bb644f791..b05a65b9092 100644 --- a/docs/rust/index.rst +++ b/docs/rust/index.rst @@ -6,8 +6,8 @@ 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, Striped and Transpose block -Load/Store; see :doc:`payload`, :doc:`groups`, and :doc:`load_store`. Its private ``cuda-coop-codegen`` crate +``ThreadData`` values, physical/mapped hierarchy queries, and block +Load/Store with Direct, Striped, Vectorize and Transpose; 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. Additional collective algorithms and public @@ -38,7 +38,7 @@ these forms: - Behavior - Payload requirement * - Full Load - - Fills every payload slot in the algorithm's order + - Fills every payload slot in the selected arrangement - Writable storage * - Default-filled Load - Loads a valid prefix and fills the suffix with a supplied value @@ -50,28 +50,28 @@ these forms: - Writes a valid prefix; leaves the destination suffix unchanged - Every slot initialized -Direct, Striped and Transpose algorithms are supported. Direct and Striped access -no shared scratch. Striped loads and stores striped payload order. -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. - -The separate validity-aware 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 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 initialized slots in their write domain. Direct and Striped preserve 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 +Direct, Striped and Vectorize access no shared scratch. Direct and Vectorize +use blocked payloads; Striped uses striped payloads. Full Vectorize operations +use wide accesses when the type, item count and adjusted pointer permit them; +misaligned pointers and counted overloads use Direct accesses. Transpose Store +redistributes blocked payloads across lanes into striped order and uses the +queried CUB storage layout. Compiler integrations enforce collective +participation, memory ownership, valid bounds and completion before scratch reuse. + +The validity-aware recipes support all four algorithms and all ten numeric +payload types: ``i8``, ``u8``, ``i16``, ``u16``, ``i32``, ``u32``, ``i64``, +``u64``, ``f32`` and ``f64``. ``Request::new_typed`` selects the exact type. +Separate packed validity bits track initialized slots. Full and default-filled +Load initialize every slot; masked Load resets slots outside its selected +prefix to unwritten. Store checks initialization throughout its write domain. +Direct, Striped and Vectorize preserve payloads and bits; Transpose permutes +both, including slots outside the prefix. Its complete scratch layout includes +the validity plane when needed. A failing Store traps and is not transactional +across lanes. Counted Store forms destination pointers only within its valid prefix, including when the allocation is shorter than a complete tile. -The public Load/Store API uses -the same ``ThreadData`` representation without exposing an initialized slice or -a public native layout. +The public API uses this representation without exposing an initialized slice +or a public native layout. Compile a provider ------------------ @@ -108,7 +108,8 @@ Native regression checks are opt-in: ``CUDA_COOP_TEST_TARGET`` can select another supported NVRTC target; its default is ``compute_120``. These checks compile the supported algorithm forms at -item counts 1 and 4, compare repeated LTO output, and check exact layout decoding +item counts 1 and 4 plus odd-count Vectorize cases, compare repeated LTO output, +and check exact layout decoding and specialization identity. They do not execute GPU kernels. Compiler integration and API reference diff --git a/docs/rust/load_store.rst b/docs/rust/load_store.rst index b495bd270d3..830f3aad617 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. Complete one-, two- and three-dimensional blocks support Direct, Striped -and Transpose algorithms. Direct and Transpose Load produce blocked payloads, +first. Complete one-, two- and three-dimensional blocks support Direct, Striped, +Vectorize and Transpose algorithms. Direct, Vectorize and Transpose Load produce blocked payloads, where each thread owns consecutive elements. Striped Load gives each thread elements separated by the block size. Other algorithms, warp groups, and public temporary storage remain under development. @@ -21,7 +21,7 @@ A complete program ------------------ The example is an independent Cargo workspace with a separate generic helper -crate. It copies three block tiles with Direct, Striped and Transpose at item counts 1 +crate. It copies three block tiles with Direct, Striped, Vectorize 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 @@ -41,7 +41,7 @@ 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 and Striped Load/Store and includes the planned shared storage for Transpose. +Direct, Striped and Vectorize Load/Store and includes the planned shared storage for Transpose. .. literalinclude:: ../../rust/cuda-coop/examples/load_store/src/main.rs :language: rust @@ -84,8 +84,8 @@ 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 and Striped Store accept ``&ThreadData``, a shared reborrow of ``&mut ThreadData``, -or a scalar numeric value. A scalar contributes one item per thread. Both preserve every payload value +Direct, Striped and Vectorize Store accept ``&ThreadData``, a shared reborrow of ``&mut ThreadData``, +or a scalar numeric value. A scalar contributes one item per thread. All three preserve every payload value and validity bit, including initialized slots outside their write domain. @@ -114,6 +114,15 @@ payload, including a zero-count Store and initialized slots outside its prefix. coop::store(block, output, &values, coop::StoreOptions::default().algorithm(coop::Striped)); +Vectorize preserves blocked payload order and uses no shared scratch. Full +Load and Store can use wide accesses when the numeric type, item count and +actual offset-adjusted pointer permit them. CUB falls back to Direct access +when the pointer is not sufficiently aligned; ordinary scalar alignment is +enough to call the operation. A counted overload uses Direct access even when +its count equals the tile size. Filled Load is also counted. The example helper +uses full calls for complete tiles and counted calls for its final short tile. +Its Store accepts a shared payload borrow and preserves values and validity. + 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 @@ -179,7 +188,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``, ``StoreValue``, ``Algorithm``, ``Direct``, ``Striped`` and ``Transpose`` specify their full generic signatures and +``LoadSource``, ``StoreValue``, ``Algorithm``, ``Direct``, ``Striped``, ``Vectorize`` 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-codegen/src/lib.rs b/rust/cuda-coop-codegen/src/lib.rs index 3d8c213bca7..d3fce9029b9 100644 --- a/rust/cuda-coop-codegen/src/lib.rs +++ b/rust/cuda-coop-codegen/src/lib.rs @@ -11,15 +11,17 @@ //! # Supported specializations //! //! [`Request`] describes an exact complete block shape, positive fixed -//! items per thread, and [`Algorithm::Direct`], [`Algorithm::Striped`] or -//! [`Algorithm::Transpose`]. -//! The raw i32 recipes support all three algorithms. Their Store payloads must +//! items per thread, and [`Algorithm::Direct`], [`Algorithm::Striped`], +//! [`Algorithm::Vectorize`] or [`Algorithm::Transpose`]. +//! The raw i32 recipes support all four algorithms. Their Store payloads must //! 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/Striped Store preserves both +//! Masked Load resets validity outside its prefix; Direct/Striped/Vectorize Store preserves both //! arrays and checks initialization only in its write domain. Transpose Store //! replaces each lane's blocked payload and validity with striped items across -//! the block, including when its write count is zero. +//! the block, including when its write count is zero. Full Vectorize operations +//! use numeric local arrays after full initialization or complete Store preflight; +//! counted Vectorize overloads retain the opaque Direct path. //! //! [`Provider::from_checkout`] snapshots the selected CCCL and CUDA headers; //! [`Provider::compile`] uses those exact bytes and loaded NVRTC to emit an @@ -94,18 +96,22 @@ pub enum Algorithm { Direct, /// Thread `t` accesses item `i` at `t + block_threads * i`, without scratch. Striped, + /// Blocked access with aligned vector operations for full tiles. + /// Misaligned pointers and counted overloads use Direct accesses. + Vectorize, /// Coalesced access with shared-memory transposition to/from blocked layout. Transpose, } impl Algorithm { fn scratch_free(self) -> bool { - matches!(self, Self::Direct | Self::Striped) + matches!(self, Self::Direct | Self::Striped | Self::Vectorize) } fn token(self) -> &'static str { match self { Self::Direct => "DIRECT", Self::Striped => "STRIPED", + Self::Vectorize => "VECTORIZE", Self::Transpose => "TRANSPOSE", } } @@ -404,7 +410,7 @@ pub struct Layout { /// 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 and Striped ignore scratch, +/// the provider does not add a trailing reuse barrier. Direct, Striped and Vectorize ignore scratch, /// so a null pointer is permitted for that argument only. All providers return /// unit; loads fill payloads in place and stores write memory. /// @@ -414,7 +420,7 @@ pub struct Layout { /// 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/Striped Store receives const pointers and borrows them +/// borrows both arrays; Direct/Striped/Vectorize Store receives const pointers and borrows them /// read-only. Transpose Store exclusively borrows both arrays and returns their /// permutation: destination `(lane, item)` receives source /// `((block_threads * item + lane) / items_per_thread, @@ -437,7 +443,7 @@ pub struct Specialization { /// Actual CUB TempStorage sizeof/alignof from the same NVRTC program. pub cub_layout: Layout, /// Complete shared region, including any validity plane and padding. - /// Direct and Striped use zero bytes despite their empty C++ types' sizeof=1. + /// Direct, Striped and Vectorize use zero bytes despite their empty C++ types' sizeof=1. pub scratch: Layout, /// Whether this specialization accesses no shared scratch or scratch-reuse barrier. pub scratch_free: bool, @@ -625,6 +631,13 @@ impl Provider { fn source(request: Request, symbol: &str) -> String { let op = request.operation; let transpose = request.algorithm == Algorithm::Transpose; + // CUB's full vectorized overload requires the same numeric T for the + // global pointer and local array. Counted overloads keep opaque cells. + let numeric_full = request.algorithm == Algorithm::Vectorize + && matches!( + op, + Operation::LoadFullValidity | Operation::StoreFullValidity + ); let (header, primitive, algorithm, qualifier, method) = if op.is_load() { ("load", "BlockLoad", "BLOCK_LOAD", "const ", "Load") } else { @@ -683,6 +696,11 @@ static_assert(alignof(scratch_storage) == alignof(primitive::TempStorage));"# "" }; let (cell, flags, body) = if op.validity() { + let cell_bytes = if numeric_full { + "&values[i]" + } else { + "values[i].bytes" + }; let domain = if matches!( op, Operation::LoadMaskedValidity | Operation::StorePartialValidity @@ -701,7 +719,7 @@ static_assert(alignof(scratch_storage) == alignof(primitive::TempStorage));"# 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)); + if ({domain}) cuda::std::memcpy(bytes + sizeof(value_type) * i, {cell_bytes}, sizeof(value_type)); }} #pragma unroll for (int byte = 0; byte < {bitmap_bytes}; ++byte) {{ @@ -716,7 +734,8 @@ static_assert(alignof(scratch_storage) == alignof(primitive::TempStorage));"# ) } else { // The general Store permutes all source slots, including initialized - // slots outside its write domain. Direct/Striped only need their write domain. + // slots outside its write domain. Nonmutating algorithms only + // need their write domain, checked before any numeric copy below. let copy_domain = if transpose { "true" } else { &domain }; let prepare = if mask_plane { format!( @@ -797,13 +816,17 @@ static_assert(alignof(scratch_storage) == alignof(primitive::TempStorage));"# cell values[{items}]; #pragma unroll for (int i = 0; i < {items}; ++i) {{ - if ({copy_domain}) cuda::std::memcpy(values[i].bytes, bytes + sizeof(value_type) * i, sizeof(value_type)); + if ({copy_domain}) cuda::std::memcpy({cell_bytes}, bytes + sizeof(value_type) * i, sizeof(value_type)); }} operation.Store({memory}, values{operands});{restore}"# ) }; ( - "cuda_coop_detail::opaque_cell", + if numeric_full { + "value_type" + } else { + "cuda_coop_detail::opaque_cell" + }, if op.is_load() || transpose { ", unsigned char* flags" } else { diff --git a/rust/cuda-coop-codegen/src/oxide.rs b/rust/cuda-coop-codegen/src/oxide.rs index 511b5b80829..e2937eb7105 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/Striped/Transpose validity recipes +//! The protocol supports the raw i32 and typed Direct/Striped/Vectorize/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. @@ -135,6 +135,7 @@ impl SpecializationRequest { let algorithm = match self.algorithm.as_str() { "direct" => Algorithm::Direct, "striped" => Algorithm::Striped, + "vectorize" => Algorithm::Vectorize, "transpose" => Algorithm::Transpose, _ => { return Err(Error::message(format!( @@ -291,6 +292,7 @@ pub fn registration( algorithm: match algorithm { Algorithm::Direct => "direct", Algorithm::Striped => "striped", + Algorithm::Vectorize => "vectorize", Algorithm::Transpose => "transpose", }, }) @@ -486,7 +488,7 @@ mod tests { ] { let scalar = Scalar::from(dtype); assert_eq!(scalar.dtype().unwrap(), dtype); - for algorithm in ["direct", "striped", "transpose"] { + for algorithm in ["direct", "striped", "vectorize", "transpose"] { let request = SpecializationRequest { operation: "load_filled_validity_bits_v1".into(), algorithm: algorithm.into(), @@ -512,7 +514,12 @@ mod tests { Operation::StoreFullValidity, Operation::StorePartialValidity, ] { - for algorithm in [Algorithm::Direct, Algorithm::Striped, Algorithm::Transpose] { + for algorithm in [ + Algorithm::Direct, + Algorithm::Striped, + Algorithm::Vectorize, + Algorithm::Transpose, + ] { let request = Request::new_typed(Dtype::I16, operation, algorithm, 128, 4).unwrap(); let mut expected = vec![ Argument::MutPointer(Scalar::Signed(16)), diff --git a/rust/cuda-coop-codegen/tests/native.rs b/rust/cuda-coop-codegen/tests/native.rs index 656d79751b4..547917363d6 100644 --- a/rust/cuda-coop-codegen/tests/native.rs +++ b/rust/cuda-coop-codegen/tests/native.rs @@ -84,6 +84,7 @@ fn native_typed_validity_and_determinism() -> Result<(), Box Result<(), Box Result<(), Box Result<(), Box Result<(), Box::new(); - coop::load( - block, - input, - &mut values, - coop::LoadOptions::default() - .valid_items(valid) - .oob_default(default) - .offset(offset) - .algorithm(algorithm), - ); - coop::store( - block, - output, - &mut values, - coop::StoreOptions::default() - .valid_items(valid) - .offset(offset) - .algorithm(algorithm), - ); + // Full calls let Vectorize use wide accesses for complete tiles. Its counted + // calls retain CUB's guarded Direct path for the final partial tile. + if valid == block.count() as usize * ITEMS_PER_THREAD { + coop::load( + block, + input, + &mut values, + coop::LoadOptions::default() + .offset(offset) + .algorithm(algorithm), + ); + coop::store( + block, + output, + &mut values, + coop::StoreOptions::default() + .offset(offset) + .algorithm(algorithm), + ); + } else { + coop::load( + block, + input, + &mut values, + coop::LoadOptions::default() + .valid_items(valid) + .oob_default(default) + .offset(offset) + .algorithm(algorithm), + ); + coop::store( + block, + output, + &mut values, + coop::StoreOptions::default() + .valid_items(valid) + .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 13821b3db0c..8570bbd0415 100644 --- a/rust/cuda-coop/examples/load_store/src/main.rs +++ b/rust/cuda-coop/examples/load_store/src/main.rs @@ -83,6 +83,31 @@ mod kernels { coop::Striped, ); } + #[kernel] + #[launch_contract(domain = 1, block = (128, 1, 1), dynamic_shared = inferred)] + pub fn vectorize( + 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::Vectorize, + ); + } } fn main() -> Result<(), Box> { @@ -117,6 +142,8 @@ fn main() -> Result<(), Box> { run!(4, prepare_copy, copy_async_owned); run!(1, prepare_striped, striped_async_owned); run!(4, prepare_striped, striped_async_owned); + run!(1, prepare_vectorize, vectorize_async_owned); + run!(4, prepare_vectorize, vectorize_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 fea98f683a0..34933c9ecad 100644 --- a/rust/cuda-coop/src/lib.rs +++ b/rust/cuda-coop/src/lib.rs @@ -30,7 +30,7 @@ //! [`this_block`] and the other dimension-free group constructors describe //! the executing thread's current launch. Queries grant no collective //! participation or synchronization guarantee. [`load`] and [`store`] support -//! Direct, Striped and Transpose on complete blocks with compiler-checked participation +//! Direct, Striped, Vectorize 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. @@ -45,8 +45,8 @@ mod thread_data; pub use alignment::{Alignment, SupportedAlignment}; pub use cuda_device::{CollectiveTiles, DisjointSlice}; pub use load_store::{ - Algorithm, Direct, LoadOptions, LoadSource, StoreOptions, StoreValue, Striped, Transpose, load, - store, + Algorithm, Direct, LoadOptions, LoadSource, StoreOptions, StoreValue, Striped, Transpose, + Vectorize, load, store, }; pub use thread_data::{Numeric, ThreadData, ThreadDataLike}; diff --git a/rust/cuda-coop/src/load_store.rs b/rust/cuda-coop/src/load_store.rs index c1bb82f526f..235ea9c3f98 100644 --- a/rust/cuda-coop/src/load_store.rs +++ b/rust/cuda-coop/src/load_store.rs @@ -18,6 +18,15 @@ pub struct Direct; #[derive(Clone, Copy, Debug, Default)] pub struct Striped; +/// Blocked access using aligned vector operations for full tiles. +/// +/// CUB checks the actual pointer after the element offset and falls back to +/// Direct access when it is not sufficiently aligned. Counts, including a +/// full-tile count, select Direct accesses. No extra alignment precondition or +/// shared scratch is required. Store preserves every payload value and bit. +#[derive(Clone, Copy, Debug, Default)] +pub struct Vectorize; + /// Coalesced access through CUB's shared-memory block exchange. /// /// Load produces blocked payloads. Store permutes payload values and validity @@ -44,6 +53,7 @@ pub struct Transpose; pub trait Algorithm: sealed::Algorithm + Copy {} impl Algorithm for Direct {} impl Algorithm for Striped {} +impl Algorithm for Vectorize {} impl Algorithm for Transpose {} impl sealed::Algorithm for Direct { const KIND: u8 = sealed::DIRECT; @@ -51,6 +61,9 @@ impl sealed::Algorithm for Direct { impl sealed::Algorithm for Striped { const KIND: u8 = sealed::STRIPED; } +impl sealed::Algorithm for Vectorize { + const KIND: u8 = sealed::VECTORIZE; +} impl sealed::Algorithm for Transpose { const KIND: u8 = sealed::TRANSPOSE; } @@ -176,16 +189,16 @@ impl StoreOptions { /// /// `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. Direct and Transpose give each thread +/// `group.count() * ITEMS` elements. Direct, Vectorize and Transpose give each thread /// consecutive items; Striped gives it items separated by the block size. /// [`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. [`Direct`] and [`Striped`] use no shared scratch; [`Transpose`] +/// The operation returns unit. [`Direct`], [`Striped`] and [`Vectorize`] use no shared scratch; [`Transpose`] /// uses compiler-planned shared storage and returns the same blocked layout. /// -/// Complete one-, two- and three-dimensional blocks support all three algorithms. +/// Complete one-, two- and three-dimensional blocks support all four algorithms. /// Thread rank is x-fastest: `t = x + X * (y + Y * z)` for shape `[X, Y, Z]`. -/// Slot `i` addresses `t * ITEMS + i` for Direct/Transpose Load and +/// Slot `i` addresses `t * ITEMS + i` for Direct/Vectorize/Transpose Load and /// `t + group.count() * i` for Striped Load, relative to the selected offset. /// Compile with the supported Oxide backend and an exact block launch contract. /// Every block thread must reach matching calls with uniform source, controls @@ -226,12 +239,12 @@ pub fn load, const ITEMS: usize, const ALIGN: usize /// Store a block tile from a per-thread payload. /// /// `T` is inferred from the destination and value. [`StoreValue`] accepts a -/// shared or mutable borrow of [`ThreadData`] for [`Direct`] and [`Striped`], a +/// shared or mutable borrow of [`ThreadData`] for [`Direct`], [`Striped`] and [`Vectorize`], a /// mutable borrow for [`Transpose`], or a numeric scalar (one item per thread) /// for any algorithm. /// `CAPACITY` describes the carrier's ownership grain; it need not equal the -/// payload extent. Direct consumes blocked items; Striped consumes striped items. -/// Both preserve every value and initialized bit and use no shared scratch. +/// payload extent. Direct and Vectorize consume blocked items; Striped consumes striped items. +/// These algorithms preserve every value and initialized bit and use 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`. @@ -336,8 +349,8 @@ impl LoadSource for &mut DisjointSlice<'_, T, Col /// A sealed per-thread value accepted by [`store`] for type `T` and algorithm `A`. /// /// Numeric scalars and mutable [`ThreadData`] borrows support all algorithms. -/// Shared payload borrows support [`Direct`] and [`Striped`]. A scalar contributes -/// one item per thread. Direct and Striped preserve every value and bit; +/// Shared payload borrows support [`Direct`], [`Striped`] and [`Vectorize`]. A scalar contributes +/// one item per thread. These algorithms preserve every value and bit; /// [`Transpose`] mutates their /// positions as described by [`store`]. pub trait StoreValue: sealed::StoreValue { @@ -377,6 +390,7 @@ macro_rules! shared_store_value { } shared_store_value!(Direct); shared_store_value!(Striped); +shared_store_value!(Vectorize); impl sealed::StoreValue for &mut ThreadData where @@ -403,6 +417,9 @@ where sealed::STRIPED => { boundary::store(destination, self, options.algorithm(Striped)); } + sealed::VECTORIZE => { + boundary::store(destination, self, options.algorithm(Vectorize)); + } sealed::TRANSPOSE => { boundary::store_transpose(destination, self, options.algorithm(Transpose)); } @@ -433,6 +450,7 @@ mod sealed { pub const DIRECT: u8 = 0; pub const STRIPED: u8 = 1; pub const TRANSPOSE: u8 = 2; + pub const VECTORIZE: u8 = 3; pub trait Algorithm { const KIND: u8; } diff --git a/rust/cuda-coop/src/load_store/boundary.rs b/rust/cuda-coop/src/load_store/boundary.rs index fc59bdfce71..fdf32328c12 100644 --- a/rust/cuda-coop/src/load_store/boundary.rs +++ b/rust/cuda-coop/src/load_store/boundary.rs @@ -3,7 +3,7 @@ use super::{ Algorithm, LoadOptions, StoreOptions, Transpose, - sealed::{DIRECT, STRIPED, TRANSPOSE}, + sealed::{DIRECT, STRIPED, TRANSPOSE, VECTORIZE}, }; use crate::{Alignment, Numeric, SupportedAlignment, ThreadData}; use cuda_device::{CollectiveTiles, DisjointSlice, library, thread}; @@ -27,6 +27,15 @@ registration!(striped_load_filled, "COOP_STRIPED_VALIDITY_LOAD_FILLED"); registration!(striped_store_full, "COOP_STRIPED_VALIDITY_STORE_FULL"); registration!(striped_store_partial, "COOP_STRIPED_VALIDITY_STORE_PARTIAL"); +registration!(vectorize_load_full, "COOP_VECTORIZE_VALIDITY_LOAD_FULL"); +registration!(vectorize_load_masked, "COOP_VECTORIZE_VALIDITY_LOAD_MASKED"); +registration!(vectorize_load_filled, "COOP_VECTORIZE_VALIDITY_LOAD_FILLED"); +registration!(vectorize_store_full, "COOP_VECTORIZE_VALIDITY_STORE_FULL"); +registration!( + vectorize_store_partial, + "COOP_VECTORIZE_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"); @@ -106,7 +115,7 @@ pub(super) unsafe fn load< // 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, $striped:ident, $transpose:ident, $controls:expr) => { + ($direct:ident, $striped:ident, $vectorize:ident, $transpose:ident, $controls:expr) => { match A::KIND { DIRECT => library::checked_block_resource_call_v1::( $direct, memory, $controls, payload, @@ -114,6 +123,11 @@ pub(super) unsafe fn load< STRIPED => library::checked_block_resource_call_v1::( $striped, memory, $controls, payload, ), + VECTORIZE => { + library::checked_block_resource_call_v1::( + $vectorize, memory, $controls, payload, + ) + } TRANSPOSE => { library::checked_block_storage_call_v1::( $transpose, @@ -136,6 +150,7 @@ pub(super) unsafe fn load< (None, None) => call!( load_full, striped_load_full, + vectorize_load_full, transpose_load_full, (pointer,) ), @@ -143,6 +158,7 @@ pub(super) unsafe fn load< call!( load_masked, striped_load_masked, + vectorize_load_masked, transpose_load_masked, (pointer, count as i32) ) @@ -150,6 +166,7 @@ pub(super) unsafe fn load< (Some(count), Some(default)) => call!( load_filled, striped_load_filled, + vectorize_load_filled, transpose_load_filled, (pointer, count as i32, default) ), @@ -187,7 +204,7 @@ pub(super) fn store< unsafe { let pointer = pointer.add(options.offset); macro_rules! call { - ($direct:ident, $striped:ident, $controls:expr) => { + ($direct:ident, $striped:ident, $vectorize:ident, $controls:expr) => { match A::KIND { DIRECT => { library::checked_block_resource_call_v1::( @@ -205,16 +222,30 @@ pub(super) fn store< payload, ) } + VECTORIZE => { + library::checked_block_resource_call_v1::( + $vectorize, + destination, + $controls, + payload, + ) + } // Only the shared-preserving Store implementations call this boundary. _ => panic!(), } }; } match options.valid_items { - None => call!(store_full, striped_store_full, (pointer,)), + None => call!( + store_full, + striped_store_full, + vectorize_store_full, + (pointer,) + ), Some(count) => call!( store_partial, striped_store_partial, + vectorize_store_partial, (pointer, count as i32) ), } diff --git a/rust/cuda-coop/tests/check_load_store.py b/rust/cuda-coop/tests/check_load_store.py index 8f08bb17555..349c5d65b98 100644 --- a/rust/cuda-coop/tests/check_load_store.py +++ b/rust/cuda-coop/tests/check_load_store.py @@ -172,6 +172,27 @@ ), ) + +vectorize_load = load.replace( + ".oob_default(0)", ".oob_default(0).algorithm(coop::Vectorize)" +) +vectorize_store = store.replace( + ".valid_items(valid)", ".valid_items(valid).algorithm(coop::Vectorize)" +) +cases.update( + vectorize_shared=(vectorize_load + vectorize_store, None), + vectorize_varying_offset=( + vectorize_load.replace( + ".offset(offset)", ".offset(offset+thread::threadIdx_x() as usize)" + ), + "not proved block-uniform", + ), + vectorize_missing_completion=( + vectorize_load + vectorize_store + reload, + "requires a full-block barrier", + ), +) + 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 9714f1a474c..4e3e20e33dd 100644 --- a/rust/cuda-coop/tests/load_store/src/main.rs +++ b/rust/cuda-coop/tests/load_store/src/main.rs @@ -4,6 +4,7 @@ use coop::{Numeric, ThreadData, this_block}; mod striped; mod transpose; +mod vectorize; use cuda_coop as coop; use cuda_device::{CollectiveTiles, DisjointSlice, kernel, launch_contract, thread}; @@ -366,6 +367,12 @@ mod kernels { fn main() -> Result<(), Box> { if let Some(argument) = std::env::args().nth(1) { + if argument == "vectorize" { + return vectorize::run(None); + } + if argument.starts_with("vectorize-") { + return vectorize::run(Some(&argument)); + } if argument == "striped" { return striped::run(None); } diff --git a/rust/cuda-coop/tests/load_store/src/vectorize.rs b/rust/cuda-coop/tests/load_store/src/vectorize.rs new file mode 100644 index 00000000000..8e99067afbd --- /dev/null +++ b/rust/cuda-coop/tests/load_store/src/vectorize.rs @@ -0,0 +1,378 @@ +// Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +//! Vectorize full/guarded access and post-Store validity, including misaligned offsets. +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 exercise( + input: &[T], + output: &mut DisjointSlice>, + valid: usize, + shift: usize, + default: T, + mode: u32, + ) { + let block = coop::this_block(); + 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 * 8 + shift; + let mut values = coop::ThreadData::<_, ITEMS>::new(); + match mode { + 0 => coop::load( + block, + input, + &mut values, + coop::LoadOptions::default() + .algorithm(coop::Vectorize) + .offset(offset), + ), + 1 | 5 | 6 | 9 => { + // A second masked Load must clear validity left by a full Load. + coop::load( + block, + input, + &mut values, + coop::LoadOptions::default() + .algorithm(coop::Vectorize) + .offset(offset), + ); + coop::load( + block, + input, + &mut values, + coop::LoadOptions::default() + .algorithm(coop::Vectorize) + .offset(offset) + .valid_items(valid), + ); + } + 2 | 4 => coop::load( + block, + input, + &mut values, + coop::LoadOptions::default() + .algorithm(coop::Vectorize) + .offset(offset) + .valid_items(valid) + .oob_default(default), + ), + 7 => {} + _ => { + let mut item = 0; + while item < ITEMS { + let selected = lane * ITEMS + item < valid; + if (selected || (lane + item) & 1 == 0) + && !(mode == 8 && lane * ITEMS + item == valid - 1) + { + values.write(item, default); + } + item += 1; + } + } + } + let before = values; + if mode == 0 || mode == 2 || mode == 9 { + coop::store( + block, + &mut *output, + &values, + coop::StoreOptions::default() + .algorithm(coop::Vectorize) + .offset(offset), + ); + } else { + let written = if mode == 5 { + valid / 2 + } else if mode == 6 { + valid + 1 + } else { + valid + }; + coop::store( + block, + &mut *output, + &values, + coop::StoreOptions::default() + .algorithm(coop::Vectorize) + .offset(offset) + .valid_items(written), + ); + } + let mut observations_ok = true; + let mut item = 0; + while item < ITEMS { + let position = lane * ITEMS + item; + let expected = match mode { + 3 | 8 => { + if position < valid || (lane + item) & 1 == 0 { + Some(default) + } else { + None + } + } + 7 => None, + _ => { + if position < valid || mode == 0 { + Some(input[offset + position]) + } else if mode == 2 || mode == 4 { + Some(default) + } else { + None + } + } + }; + observations_ok &= before.get(item) == expected; + observations_ok &= values.get(item) == before.get(item); + item += 1; + } + if mode == 3 || mode == 4 { + thread::sync_threads(); + // A repeated zero Store preserves arbitrary initialized/invalid slots. + coop::store( + block, + &mut *output, + &values, + coop::StoreOptions::default() + .algorithm(coop::Vectorize) + .offset(offset) + .valid_items(0), + ); + let mut item = 0; + while item < ITEMS { + observations_ok &= values.get(item) == before.get(item); + item += 1; + } + } + assert!(observations_ok); + } + + #[kernel] + #[launch_contract(domain = 1, block = (33, 1, 1), dynamic_shared = inferred)] + pub fn partial_33( + input: &[T], + mut output: DisjointSlice>, + valid: usize, + shift: usize, + default: T, + mode: u32, + ) { + exercise::<_, ITEMS>(input, &mut output, valid, shift, default, mode); + } + + #[kernel] + #[launch_contract(domain = 3, block = (8, 4, 4), dynamic_shared = inferred)] + pub fn partial_3d( + input: &[T], + mut output: DisjointSlice>, + valid: usize, + shift: usize, + default: T, + mode: u32, + ) { + exercise::<_, ITEMS>(input, &mut output, valid, shift, default, mode); + } + + #[kernel] + #[launch_contract(domain = 1, block = (33, 1, 1), dynamic_shared = inferred)] + pub fn scalar(mut output: DisjointSlice>, valid: usize) { + let offset = thread::blockIdx_x() as usize * 33; + coop::store( + coop::this_block(), + &mut output, + 17_i32, + coop::StoreOptions::default() + .algorithm(coop::Vectorize) + .offset(offset) + .valid_items(valid), + ); + } +} + +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, LaunchConfig3D}; + init_device_contexts(0, 1)?; + let context = with_cuda_context(0, |context| context.clone())?; + let stream = context.default_stream(); + // SAFETY: embedded kernels and typed declarations compile together. + let module = unsafe { kernels::load_async(0)? }; + if let Some(trap) = trap { + let mode = match trap { + "vectorize-larger" => 6, + "vectorize-full-hole" => 9, + _ => 8, + }; + let plan = module.prepare_partial_33::(LaunchConfig1D::auto_shared(1, 33))?; + let source = DeviceBuffer::from_host(&stream, &vec![7_i32; 132])?; + let target = DeviceBuffer::from_host(&stream, &vec![0_i32; 132])?; + module + .partial_33_async_owned::(&plan, source, target, 34, 0, -7, mode) + .sync()?; + return Err("invalid Vectorize Store unexpectedly completed".into()); + } + let mut cases = 0; + macro_rules! run { + ($t:ty, $items:literal, $threads:literal, $prepare:ident, $launch:ident, $config:expr, $extra:expr) => {{ + let tile = $threads * $items; + let plan = module.$prepare::<$t, $items>($config)?; + assert_eq!(plan.__raw_config().shared_mem_bytes, 0); + let prefix = ($threads + 1).min(tile); + let mut selected = vec![ + (0, tile, 0), + (0, tile, 1), + (0, tile, 4), + (1, prefix, 0), + (2, 1, 0), + (3, prefix, 0), + (5, prefix, 0), + ]; + if $extra { + for valid in [0, 1, $threads - 1, $threads, prefix, tile] { + for mode in [1, 2, 3, 4, 5] { + selected.push((mode, valid, 0)); + } + } + for mode in [1, 3, 4, 5] { + selected.push((mode, tile - 7, 7)); + } + selected.push((7, 0, 0)); + } + selected.sort_unstable(); + selected.dedup(); + for (mode, valid, shift) in selected { + let written = if mode == 0 || mode == 2 || mode == 9 { + tile + } else if mode == 5 { + valid / 2 + } else { + valid + }; + // Keep the final allocation shorter than the full tile when + // counted Store writes a prefix. Its guarded iterator must not + // form an out-of-allocation pointer. + let output_len = $threads * 8 + shift + written; + let input: Vec<$t> = (0..2 * $threads * 8 + shift) + .map(|i| (i % 101) as $t) + .collect(); + let default = 111 as $t; + let sentinel = 123 as $t; + let source = DeviceBuffer::from_host(&stream, &input)?; + let target = DeviceBuffer::from_host(&stream, &vec![sentinel; output_len])?; + let (_, target) = module + .$launch::<$t, $items, _, _>(&plan, source, target, valid, shift, default, mode) + .sync()?; + let mut output = vec![sentinel; output_len]; + target.copy_to_host(&stream, &mut output)?; + stream.synchronize()?; + for (index, actual) in output.into_iter().enumerate() { + let local = index % ($threads * 8); + let expected = if local >= shift && local - shift < written { + if mode == 3 || (mode == 2 && local - shift >= valid) { + default + } else { + input[index] + } + } else { + sentinel + }; + assert_eq!( + actual, + expected, + "type={} items={} mode={mode} valid={valid} shift={shift} index={index}", + stringify!($t), + $items + ); + } + cases += 1; + } + }}; + } + macro_rules! width { + ($t:ty) => { + run!( + $t, + 1, + 33, + prepare_partial_33, + partial_33_async_owned, + LaunchConfig1D::auto_shared(2, 33), + false + ); + run!( + $t, + 4, + 33, + prepare_partial_33, + partial_33_async_owned, + LaunchConfig1D::auto_shared(2, 33), + false + ); + }; + } + width!(i8); + width!(u8); + width!(i16); + width!(u16); + width!(i32); + width!(u32); + width!(i64); + width!(u64); + width!(f32); + width!(f64); + run!( + i32, + 4, + 33, + prepare_partial_33, + partial_33_async_owned, + LaunchConfig1D::auto_shared(2, 33), + true + ); + run!( + i32, + 4, + 128, + prepare_partial_3d, + partial_3d_async_owned, + LaunchConfig3D::auto_shared((1, 2, 1), (8, 4, 4)), + true + ); + run!( + i32, + 3, + 33, + prepare_partial_33, + partial_33_async_owned, + LaunchConfig1D::auto_shared(2, 33), + false + ); + for valid in [0, 1, 32, 33] { + let plan = module.prepare_scalar(LaunchConfig1D::auto_shared(2, 33))?; + assert_eq!(plan.__raw_config().shared_mem_bytes, 0); + let target = DeviceBuffer::from_host(&stream, &vec![-1_i32; 33 + valid])?; + let target = module.scalar_async_owned(&plan, target, valid).sync()?; + let mut output = vec![0_i32; 33 + valid]; + target.copy_to_host(&stream, &mut output)?; + stream.synchronize()?; + for (i, value) in output.into_iter().enumerate() { + assert_eq!(value, if i % 33 < valid { 17 } else { -1 }); + } + cases += 1; + } + println!("PASS: {cases} Vectorize cases with observed payload mapping and validity"); + Ok(()) +} From 101c0c4844d4347617320f4e7cc75f664a57904c Mon Sep 17 00:00:00 2001 From: Trent Nelson Date: Tue, 20 Oct 2026 13:27:42 -0700 Subject: [PATCH 39/48] Use a distinct Vectorize fixture kernel name Signed-off-by: Trent Nelson --- rust/cuda-coop/tests/load_store/src/vectorize.rs | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/rust/cuda-coop/tests/load_store/src/vectorize.rs b/rust/cuda-coop/tests/load_store/src/vectorize.rs index 8e99067afbd..ddd39a1b920 100644 --- a/rust/cuda-coop/tests/load_store/src/vectorize.rs +++ b/rust/cuda-coop/tests/load_store/src/vectorize.rs @@ -187,7 +187,7 @@ mod kernels { #[kernel] #[launch_contract(domain = 1, block = (33, 1, 1), dynamic_shared = inferred)] - pub fn scalar(mut output: DisjointSlice>, valid: usize) { + pub fn vectorize_scalar(mut output: DisjointSlice>, valid: usize) { let offset = thread::blockIdx_x() as usize * 33; coop::store( coop::this_block(), @@ -361,10 +361,12 @@ pub fn run(trap: Option<&str>) -> Result<(), Box> { false ); for valid in [0, 1, 32, 33] { - let plan = module.prepare_scalar(LaunchConfig1D::auto_shared(2, 33))?; + let plan = module.prepare_vectorize_scalar(LaunchConfig1D::auto_shared(2, 33))?; assert_eq!(plan.__raw_config().shared_mem_bytes, 0); let target = DeviceBuffer::from_host(&stream, &vec![-1_i32; 33 + valid])?; - let target = module.scalar_async_owned(&plan, target, valid).sync()?; + let target = module + .vectorize_scalar_async_owned(&plan, target, valid) + .sync()?; let mut output = vec![0_i32; 33 + valid]; target.copy_to_host(&stream, &mut output)?; stream.synchronize()?; From acf9ab4e4db383b5d669229e5f5068a2d896cc5c Mon Sep 17 00:00:00 2001 From: Trent Nelson Date: Tue, 20 Oct 2026 13:45:14 -0700 Subject: [PATCH 40/48] Distinguish unexpected Vectorize Store completion Signed-off-by: Trent Nelson --- rust/cuda-coop/tests/load_store/src/vectorize.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/rust/cuda-coop/tests/load_store/src/vectorize.rs b/rust/cuda-coop/tests/load_store/src/vectorize.rs index ddd39a1b920..c4708797870 100644 --- a/rust/cuda-coop/tests/load_store/src/vectorize.rs +++ b/rust/cuda-coop/tests/load_store/src/vectorize.rs @@ -222,7 +222,7 @@ pub fn run(trap: Option<&str>) -> Result<(), Box> { module .partial_33_async_owned::(&plan, source, target, 34, 0, -7, mode) .sync()?; - return Err("invalid Vectorize Store unexpectedly completed".into()); + panic!("invalid Vectorize Store unexpectedly completed"); } let mut cases = 0; macro_rules! run { From 2e2d9fc6fe3b0385fa399b8bd6909d52c384d42c Mon Sep 17 00:00:00 2001 From: Trent Nelson Date: Tue, 20 Oct 2026 17:36:08 -0700 Subject: [PATCH 41/48] Add Rust block warp-transpose Load and Store Support both CUB warp-transpose algorithms for complete blocks. Preserve payload validity through Store permutation and query exact scratch layouts for regular and timesliced exchanges. Signed-off-by: Trent Nelson --- docs/rust/index.rst | 19 +- docs/rust/load_store.rst | 30 +- rust/cuda-coop-codegen/src/block_io.cuh | 5 + rust/cuda-coop-codegen/src/lib.rs | 65 +++- rust/cuda-coop-codegen/src/oxide.rs | 17 +- rust/cuda-coop-codegen/tests/native.rs | 85 +++- rust/cuda-coop/build.rs | 5 + .../cuda-coop/examples/load_store/src/main.rs | 62 +++ rust/cuda-coop/src/lib.rs | 7 +- rust/cuda-coop/src/load_store.rs | 68 +++- rust/cuda-coop/src/load_store/boundary.rs | 130 ++++++- rust/cuda-coop/tests/check_load_store.py | 30 +- .../tests/load_store/src/block_warp.rs | 366 ++++++++++++++++++ rust/cuda-coop/tests/load_store/src/main.rs | 7 + 14 files changed, 822 insertions(+), 74 deletions(-) create mode 100644 rust/cuda-coop/tests/load_store/src/block_warp.rs diff --git a/docs/rust/index.rst b/docs/rust/index.rst index b05a65b9092..39bb62e3c36 100644 --- a/docs/rust/index.rst +++ b/docs/rust/index.rst @@ -7,7 +7,9 @@ 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 block -Load/Store with Direct, Striped, Vectorize and Transpose; see :doc:`payload`, :doc:`groups`, and :doc:`load_store`. Its private ``cuda-coop-codegen`` crate +Load/Store with Direct, Striped, Vectorize, Transpose, WarpTranspose and +WarpTransposeTimesliced; 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. Additional collective algorithms and public @@ -55,18 +57,23 @@ use blocked payloads; Striped uses striped payloads. Full Vectorize operations use wide accesses when the type, item count and adjusted pointer permit them; misaligned pointers and counted overloads use Direct accesses. Transpose Store redistributes blocked payloads across lanes into striped order and uses the -queried CUB storage layout. Compiler integrations enforce collective +queried CUB storage layout. WarpTranspose and WarpTransposeTimesliced exchange +within physical warps and require a complete block whose thread count is +divisible by 32. The timesliced form reuses one warp's CUB exchange buffer. +Compiler integrations enforce collective participation, memory ownership, valid bounds and completion before scratch reuse. -The validity-aware recipes support all four algorithms and all ten numeric +The validity-aware recipes support all six algorithms and all ten numeric payload types: ``i8``, ``u8``, ``i16``, ``u16``, ``i32``, ``u32``, ``i64``, ``u64``, ``f32`` and ``f64``. ``Request::new_typed`` selects the exact type. Separate packed validity bits track initialized slots. Full and default-filled Load initialize every slot; masked Load resets slots outside its selected prefix to unwritten. Store checks initialization throughout its write domain. -Direct, Striped and Vectorize preserve payloads and bits; Transpose permutes -both, including slots outside the prefix. Its complete scratch layout includes -the validity plane when needed. A failing Store traps and is not transactional +Direct, Striped and Vectorize preserve payloads and bits; the three transposing +algorithms permute both, including slots outside the prefix. Their complete +scratch layouts include a validity plane when needed. For counted warp-transpose +Store with more than one item per thread, that plane covers every warp even +when the CUB exchange buffer is timesliced. A failing Store traps and is not transactional across lanes. Counted Store forms destination pointers only within its valid prefix, including when the allocation is shorter than a complete tile. diff --git a/docs/rust/load_store.rst b/docs/rust/load_store.rst index 830f3aad617..434eb2b4227 100644 --- a/docs/rust/load_store.rst +++ b/docs/rust/load_store.rst @@ -7,10 +7,11 @@ 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. Complete one-, two- and three-dimensional blocks support Direct, Striped, -Vectorize and Transpose algorithms. Direct, Vectorize and Transpose Load produce blocked payloads, -where each thread owns consecutive elements. Striped Load gives each thread -elements separated by the block size. Other algorithms, warp -groups, and public temporary storage remain under development. +Vectorize, Transpose, WarpTranspose and WarpTransposeTimesliced. The two warp +transpose algorithms require a total block thread count divisible by 32. +Striped Load gives each thread elements separated by the block size; the other +algorithms produce blocked payloads, where each thread owns consecutive elements. +Warp-group invocation 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. @@ -21,7 +22,7 @@ A complete program ------------------ The example is an independent Cargo workspace with a separate generic helper -crate. It copies three block tiles with Direct, Striped, Vectorize and Transpose at item counts 1 +crate. It copies three block tiles with each of the six algorithms 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 @@ -41,14 +42,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, Striped and Vectorize Load/Store and includes the planned shared storage for Transpose. +Direct, Striped and Vectorize Load/Store and includes the planned shared storage +for the three transposing algorithms. .. 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. Its algorithm-generic -form borrows the payload mutably because Transpose Store permutes it. +form borrows the payload mutably because the transposing Stores permute it. .. literalinclude:: ../../rust/cuda-coop/examples/load_store/helper/src/lib.rs :language: rust @@ -132,6 +134,17 @@ 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 every algorithm. +WarpTranspose and WarpTransposeTimesliced use the same blocked Load layout. +Their Store requires a mutable payload and permutes every slot within its +physical warp: ``(t, i)`` receives old blocked slot +``q = (t / 32) * (32 * I) + 32 * i + t % 32``. Both values and validity follow +that mapping, including zero-count and repeated Stores. The timesliced variant +reuses one warp's CUB exchange storage across successive warps; its counted +Store retains a separate validity plane for all warps. Both algorithms still +require the entire block to participate, and their internal warp barriers do +not complete a conflicting full-block memory access. For shape ``[X,Y,Z]``, +``X*Y*Z`` must be divisible by 32; X alone need not be. + 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`` @@ -188,7 +201,8 @@ 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``, ``StoreValue``, ``Algorithm``, ``Direct``, ``Striped``, ``Vectorize`` and ``Transpose`` specify their full generic signatures and +``LoadSource``, ``StoreValue``, ``Algorithm`` and the six algorithm selectors +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-codegen/src/block_io.cuh b/rust/cuda-coop-codegen/src/block_io.cuh index d404f17e039..44aaca2295d 100644 --- a/rust/cuda-coop-codegen/src/block_io.cuh +++ b/rust/cuda-coop-codegen/src/block_io.cuh @@ -14,6 +14,11 @@ struct alignas(T) opaque_cell unsigned char bytes[sizeof(T)]; __device__ opaque_cell() = default; + // Warp-striped CUB Load constructs a cell from an initialized input value. + __device__ opaque_cell(T value) + { + cuda::std::memcpy(bytes, &value, sizeof(T)); + } __device__ opaque_cell(const opaque_cell& other) { cuda::std::memcpy(bytes, other.bytes, sizeof(T)); diff --git a/rust/cuda-coop-codegen/src/lib.rs b/rust/cuda-coop-codegen/src/lib.rs index d3fce9029b9..a737a706078 100644 --- a/rust/cuda-coop-codegen/src/lib.rs +++ b/rust/cuda-coop-codegen/src/lib.rs @@ -12,14 +12,16 @@ //! //! [`Request`] describes an exact complete block shape, positive fixed //! items per thread, and [`Algorithm::Direct`], [`Algorithm::Striped`], -//! [`Algorithm::Vectorize`] or [`Algorithm::Transpose`]. -//! The raw i32 recipes support all four algorithms. Their Store payloads must +//! [`Algorithm::Vectorize`], [`Algorithm::Transpose`], [`Algorithm::WarpTranspose`] +//! or [`Algorithm::WarpTransposeTimesliced`]. The two warp-transpose algorithms +//! require a total block thread count divisible by 32. +//! The raw i32 recipes support all six algorithms. Their Store payloads must //! 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/Striped/Vectorize Store preserves both -//! arrays and checks initialization only in its write domain. Transpose Store +//! arrays and checks initialization only in its write domain. Transposing Store //! replaces each lane's blocked payload and validity with striped items across -//! the block, including when its write count is zero. Full Vectorize operations +//! the block or physical warp, including when its write count is zero. Full Vectorize operations //! use numeric local arrays after full initialization or complete Store preflight; //! counted Vectorize overloads retain the opaque Direct path. //! @@ -101,6 +103,12 @@ pub enum Algorithm { Vectorize, /// Coalesced access with shared-memory transposition to/from blocked layout. Transpose, + /// Coalesced access with a separate shared exchange for each physical warp. + /// Requires a block thread count divisible by 32. + WarpTranspose, + /// Warp transposition that reuses one physical warp's CUB exchange storage. + /// Requires a block thread count divisible by 32. + WarpTransposeTimesliced, } impl Algorithm { fn scratch_free(self) -> bool { @@ -113,6 +121,8 @@ impl Algorithm { Self::Striped => "STRIPED", Self::Vectorize => "VECTORIZE", Self::Transpose => "TRANSPOSE", + Self::WarpTranspose => "WARP_TRANSPOSE", + Self::WarpTransposeTimesliced => "WARP_TRANSPOSE_TIMESLICED", } } } @@ -176,10 +186,10 @@ pub enum Operation { /// Counted Load with an exact typed default; marks every slot initialized. LoadFilledValidity, /// Full Store after checking every validity bit is set. - /// Transpose redistributes the payload; all resulting slots remain initialized. + /// Transposing algorithms redistribute the payload; all resulting slots remain initialized. StoreFullValidity, /// Counted Store after checking its write domain. - /// Transpose redistributes all values and validity bits, even outside that domain. + /// Transposing algorithms redistribute all values and validity bits, even outside that domain. StorePartialValidity, } impl Operation { @@ -323,7 +333,8 @@ impl Request { /// # Errors /// /// Rejects the shapes documented by [`Self::new_with_shape`] and a non-i32 - /// raw recipe. + /// raw recipe. Warp-transpose algorithms require a total thread count + /// divisible by 32; the X extent need not be divisible by 32. pub fn new_typed_with_shape( dtype: Dtype, operation: Operation, @@ -348,6 +359,15 @@ impl Request { "require 1..=1024 block threads, positive items per thread and tile extent <= i32::MAX", )); } + if matches!( + algorithm, + Algorithm::WarpTranspose | Algorithm::WarpTransposeTimesliced + ) && block_threads.unwrap() % 32 != 0 + { + return Err(Error::message( + "block warp transpose requires a block thread count divisible by 32", + )); + } Ok(Self { operation, algorithm, @@ -408,7 +428,7 @@ pub struct Layout { /// /// 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 +/// Scratch for transposing algorithms 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, Striped and Vectorize ignore scratch, /// so a null pointer is permitted for that argument only. All providers return @@ -421,15 +441,16 @@ pub struct Layout { /// 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/Striped/Vectorize Store receives const pointers and borrows them -/// 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)`. +/// read-only. Transposing Stores exclusively borrow both arrays. Destination +/// `(lane, item)` receives source `(q / items_per_thread, q % items_per_thread)`: +/// Transpose uses `q = block_threads * item + lane`; both warp variants use +/// `q = (lane / 32) * (32 * items_per_thread) + 32 * item + lane % 32`. +/// Lane rank is x-fastest within the complete block. /// 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. Its counted Transpose form transports +/// write domain before numeric conversion. The counted transposing forms transport /// 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. @@ -448,7 +469,7 @@ 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 for Transpose stores, including counted stores with zero writes. + /// True for transposing stores, including counted stores with zero writes. /// This describes Store rearrangement, not whether Load writes its payload. pub mutates_payload: bool, } @@ -601,8 +622,7 @@ impl Provider { alignment: 1, }, scratch_free: request.algorithm.scratch_free(), - mutates_payload: !request.operation.is_load() - && request.algorithm == Algorithm::Transpose, + mutates_payload: !request.operation.is_load() && !request.algorithm.scratch_free(), }); } let output = self @@ -630,7 +650,11 @@ impl Provider { fn source(request: Request, symbol: &str) -> String { let op = request.operation; - let transpose = request.algorithm == Algorithm::Transpose; + let transpose = !request.algorithm.scratch_free(); + let warp_transpose = matches!( + request.algorithm, + Algorithm::WarpTranspose | Algorithm::WarpTransposeTimesliced + ); // CUB's full vectorized overload requires the same numeric T for the // global pointer and local array. Counted overloads keep opaque cells. let numeric_full = request.algorithm == Algorithm::Vectorize @@ -649,6 +673,11 @@ fn source(request: Request, symbol: &str) -> String { let threads = request.block_threads(); let bitmap_bytes = items.div_ceil(8); let tile = threads * items; + let source_index = if warp_transpose { + format!("(tid / 32) * (32 * {items}) + 32 * i + (tid % 32)") + } else { + format!("{threads} * i + tid") + }; let mask_plane = op == Operation::StorePartialValidity && transpose && items > 1; let storage = if mask_plane { let plane_bytes = u64::from(threads.div_ceil(32)) * u64::from(items) * 4; @@ -780,7 +809,7 @@ static_assert(alignof(scratch_storage) == alignof(primitive::TempStorage));"# for (int bit = 0; bit < 8; ++bit) {{ const int i = byte * 8 + bit; if (i < {items}) {{ - const int q = {threads} * i + tid; + const int q = {source_index}; const int source_thread = q / {items}; const int source_item = q % {items}; const auto* const address = reinterpret_cast( diff --git a/rust/cuda-coop-codegen/src/oxide.rs b/rust/cuda-coop-codegen/src/oxide.rs index e2937eb7105..8ca40f07aa6 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/Striped/Vectorize/Transpose validity recipes +//! The protocol supports raw i32 and typed validity recipes for all six block algorithms //! documented by [`Operation`]. //! Device calls retain all participation, validity and scratch obligations of //! [`crate::Specialization`]. No allocation or synchronization policy is supplied. @@ -137,6 +137,8 @@ impl SpecializationRequest { "striped" => Algorithm::Striped, "vectorize" => Algorithm::Vectorize, "transpose" => Algorithm::Transpose, + "warp_transpose" => Algorithm::WarpTranspose, + "warp_transpose_timesliced" => Algorithm::WarpTransposeTimesliced, _ => { return Err(Error::message(format!( "unsupported algorithm {:?}", @@ -294,6 +296,8 @@ pub fn registration( Algorithm::Striped => "striped", Algorithm::Vectorize => "vectorize", Algorithm::Transpose => "transpose", + Algorithm::WarpTranspose => "warp_transpose", + Algorithm::WarpTransposeTimesliced => "warp_transpose_timesliced", }, }) .map_err(|error| Error::message(format!("cannot encode provider registration: {error}"))) @@ -488,7 +492,14 @@ mod tests { ] { let scalar = Scalar::from(dtype); assert_eq!(scalar.dtype().unwrap(), dtype); - for algorithm in ["direct", "striped", "vectorize", "transpose"] { + for algorithm in [ + "direct", + "striped", + "vectorize", + "transpose", + "warp_transpose", + "warp_transpose_timesliced", + ] { let request = SpecializationRequest { operation: "load_filled_validity_bits_v1".into(), algorithm: algorithm.into(), @@ -519,6 +530,8 @@ mod tests { Algorithm::Striped, Algorithm::Vectorize, Algorithm::Transpose, + Algorithm::WarpTranspose, + Algorithm::WarpTransposeTimesliced, ] { let request = Request::new_typed(Dtype::I16, operation, algorithm, 128, 4).unwrap(); let mut expected = vec![ diff --git a/rust/cuda-coop-codegen/tests/native.rs b/rust/cuda-coop-codegen/tests/native.rs index 547917363d6..b43d9f63f70 100644 --- a/rust/cuda-coop-codegen/tests/native.rs +++ b/rust/cuda-coop-codegen/tests/native.rs @@ -11,6 +11,16 @@ 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 algorithm in [Algorithm::WarpTranspose, Algorithm::WarpTransposeTimesliced] { + assert!( + Request::new(Operation::LoadFull, algorithm, 33, 4) + .unwrap_err() + .to_string() + .contains("divisible by 32") + ); + assert!(Request::new_with_shape(Operation::LoadFull, algorithm, [11, 3, 1], 4).is_err()); + assert!(Request::new_with_shape(Operation::LoadFull, algorithm, [8, 4, 2], 4).is_ok()); + } for block in [ [0, 4, 4], [8, 0, 4], @@ -86,6 +96,8 @@ fn native_typed_validity_and_determinism() -> Result<(), Box Result<(), Box Result<(), Box Result<(), Box (256, 272), + (Algorithm::WarpTranspose, 4) => (1024, 1040), + (Algorithm::WarpTransposeTimesliced, 1) => (128, 144), + (Algorithm::WarpTransposeTimesliced, 4) => (512, 528), + _ => unreachable!(), + }; + let is_store = matches!( + request.operation(), + Operation::StoreFullValidity | Operation::StorePartialValidity + ); + assert_eq!( + entry.cub_layout, + Layout { + size: if is_store { store } else { load }, + alignment: 16 + } + ); + } if request.operation() == Operation::StorePartialValidity && request.items_per_thread() > 1 { @@ -230,6 +283,8 @@ fn native_specializations_and_determinism() -> Result<(), Box Result<(), Box Result<(), Box Result<(), Box= 128 * request.items_per_thread() as usize * 4); + let exchange_threads = if request.algorithm() == Algorithm::WarpTransposeTimesliced { + 32 + } else { + 128 + }; + assert!( + entry.scratch.size >= exchange_threads * request.items_per_thread() as usize * 4 + ); } } // A specialization's symbol is independent of unrelated batch members. diff --git a/rust/cuda-coop/build.rs b/rust/cuda-coop/build.rs index 65d476ab5da..1052620aed5 100644 --- a/rust/cuda-coop/build.rs +++ b/rust/cuda-coop/build.rs @@ -24,6 +24,11 @@ fn main() { (Algorithm::Striped, "STRIPED"), (Algorithm::Vectorize, "VECTORIZE"), (Algorithm::Transpose, "TRANSPOSE"), + (Algorithm::WarpTranspose, "WARP_TRANSPOSE"), + ( + Algorithm::WarpTransposeTimesliced, + "WARP_TRANSPOSE_TIMESLICED", + ), ] { let value = oxide::registration(&executable, operation, algorithm) .expect("cuda.coop algorithm registration"); diff --git a/rust/cuda-coop/examples/load_store/src/main.rs b/rust/cuda-coop/examples/load_store/src/main.rs index 8570bbd0415..6e827fe382f 100644 --- a/rust/cuda-coop/examples/load_store/src/main.rs +++ b/rust/cuda-coop/examples/load_store/src/main.rs @@ -108,6 +108,56 @@ mod kernels { coop::Vectorize, ); } + #[kernel] + #[launch_contract(domain = 1, block = (128, 1, 1), dynamic_shared = inferred)] + pub fn warp_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::WarpTranspose, + ); + } + #[kernel] + #[launch_contract(domain = 1, block = (128, 1, 1), dynamic_shared = inferred)] + pub fn warp_transpose_timesliced( + 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::WarpTransposeTimesliced, + ); + } } fn main() -> Result<(), Box> { @@ -146,5 +196,17 @@ fn main() -> Result<(), Box> { run!(4, prepare_vectorize, vectorize_async_owned); run!(1, prepare_transpose, transpose_async_owned); run!(4, prepare_transpose, transpose_async_owned); + run!(1, prepare_warp_transpose, warp_transpose_async_owned); + run!(4, prepare_warp_transpose, warp_transpose_async_owned); + run!( + 1, + prepare_warp_transpose_timesliced, + warp_transpose_timesliced_async_owned + ); + run!( + 4, + prepare_warp_transpose_timesliced, + warp_transpose_timesliced_async_owned + ); Ok(()) } diff --git a/rust/cuda-coop/src/lib.rs b/rust/cuda-coop/src/lib.rs index 34933c9ecad..fff2d90160d 100644 --- a/rust/cuda-coop/src/lib.rs +++ b/rust/cuda-coop/src/lib.rs @@ -30,8 +30,9 @@ //! [`this_block`] and the other dimension-free group constructors describe //! the executing thread's current launch. Queries grant no collective //! participation or synchronization guarantee. [`load`] and [`store`] support -//! Direct, Striped, Vectorize and Transpose on complete blocks with compiler-checked participation -//! and collective memory ownership. Transpose uses compiler-planned shared +//! Direct, Striped, Vectorize and three transpose algorithms on complete blocks +//! with compiler-checked participation and collective memory ownership. The +//! transpose algorithms use compiler-planned shared //! scratch. Host payload tests need no GPU; //! executing hardware queries requires an Oxide-compiled CUDA kernel. @@ -46,7 +47,7 @@ pub use alignment::{Alignment, SupportedAlignment}; pub use cuda_device::{CollectiveTiles, DisjointSlice}; pub use load_store::{ Algorithm, Direct, LoadOptions, LoadSource, StoreOptions, StoreValue, Striped, Transpose, - Vectorize, load, store, + Vectorize, WarpTranspose, WarpTransposeTimesliced, load, store, }; pub use thread_data::{Numeric, ThreadData, ThreadDataLike}; diff --git a/rust/cuda-coop/src/load_store.rs b/rust/cuda-coop/src/load_store.rs index 235ea9c3f98..cc799eedf55 100644 --- a/rust/cuda-coop/src/load_store.rs +++ b/rust/cuda-coop/src/load_store.rs @@ -49,12 +49,32 @@ pub struct Vectorize; #[derive(Clone, Copy, Debug, Default)] pub struct Transpose; +/// Coalesced blocked access through a shared-memory exchange within each warp. +/// +/// The complete block participates; its total thread count must be divisible +/// by 32. Load produces blocked payloads. Store mutates every value and validity +/// bit into warp-striped order, including when its write count is zero. See +/// [`store`] for the exact permutation and synchronization requirements. +#[derive(Clone, Copy, Debug, Default)] +pub struct WarpTranspose; + +/// [`WarpTranspose`] using one warp's CUB exchange storage in successive slices. +/// +/// Every thread in the complete block participates, including the block barriers +/// between slices. The block thread count must be divisible by 32. Load, Store +/// mutation and validity follow [`WarpTranspose`]. Counted Store additionally +/// reserves its validity plane for every warp; that plane is not timesliced. +#[derive(Clone, Copy, Debug, Default)] +pub struct WarpTransposeTimesliced; + /// A sealed, compile-time algorithm selected through Load/Store options. pub trait Algorithm: sealed::Algorithm + Copy {} impl Algorithm for Direct {} impl Algorithm for Striped {} impl Algorithm for Vectorize {} impl Algorithm for Transpose {} +impl Algorithm for WarpTranspose {} +impl Algorithm for WarpTransposeTimesliced {} impl sealed::Algorithm for Direct { const KIND: u8 = sealed::DIRECT; } @@ -67,6 +87,12 @@ impl sealed::Algorithm for Vectorize { impl sealed::Algorithm for Transpose { const KIND: u8 = sealed::TRANSPOSE; } +impl sealed::Algorithm for WarpTranspose { + const KIND: u8 = sealed::WARP_TRANSPOSE; +} +impl sealed::Algorithm for WarpTransposeTimesliced { + const KIND: u8 = sealed::WARP_TRANSPOSE_TIMESLICED; +} /// Controls for a [`load`], using [`Direct`] by default. /// @@ -189,17 +215,20 @@ impl StoreOptions { /// /// `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. Direct, Vectorize and Transpose give each thread -/// consecutive items; Striped gives it items separated by the block size. +/// `group.count() * ITEMS` elements. Striped gives each thread items separated +/// by the block size; the other algorithms give it 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. [`Direct`], [`Striped`] and [`Vectorize`] use no shared scratch; [`Transpose`] -/// uses compiler-planned shared storage and returns the same blocked layout. +/// The operation returns unit. [`Direct`], [`Striped`] and [`Vectorize`] use no +/// shared scratch. The three transposing algorithms use compiler-planned shared +/// storage and return blocked payloads. /// -/// Complete one-, two- and three-dimensional blocks support all four algorithms. +/// Complete one-, two- and three-dimensional blocks support all six algorithms. +/// [`WarpTranspose`] and [`WarpTransposeTimesliced`] require a total thread count +/// divisible by 32; they still require participation by every block thread. /// Thread rank is x-fastest: `t = x + X * (y + Y * z)` for shape `[X, Y, Z]`. -/// Slot `i` addresses `t * ITEMS + i` for Direct/Vectorize/Transpose Load and -/// `t + group.count() * i` for Striped Load, relative to the selected offset. +/// Slot `i` addresses `t + group.count() * i` for Striped Load and `t * ITEMS + i` +/// for the other algorithms, relative to the selected offset. /// 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 @@ -240,7 +269,8 @@ pub fn load, const ITEMS: usize, const ALIGN: usize /// /// `T` is inferred from the destination and value. [`StoreValue`] accepts a /// shared or mutable borrow of [`ThreadData`] for [`Direct`], [`Striped`] and [`Vectorize`], a -/// mutable borrow for [`Transpose`], or a numeric scalar (one item per thread) +/// mutable borrow for [`Transpose`], [`WarpTranspose`] or [`WarpTransposeTimesliced`], +/// or a numeric scalar (one item per thread) /// for any algorithm. /// `CAPACITY` describes the carrier's ownership grain; it need not equal the /// payload extent. Direct and Vectorize consume blocked items; Striped consumes striped items. @@ -248,7 +278,10 @@ pub fn load, const ITEMS: usize, const ALIGN: usize /// 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 +/// Both warp-transpose algorithms instead use warp-striped index +/// `(t / 32) * (32 * ITEMS) + 32 * i + t % 32`, where `ITEMS` is the payload extent. +/// They also mutate every slot, including a zero-count Store. All three preserve +/// the initialized bit of the moved 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 @@ -351,8 +384,7 @@ impl LoadSource for &mut DisjointSlice<'_, T, Col /// Numeric scalars and mutable [`ThreadData`] borrows support all algorithms. /// Shared payload borrows support [`Direct`], [`Striped`] and [`Vectorize`]. A scalar contributes /// one item per thread. These algorithms preserve every value and bit; -/// [`Transpose`] mutates their -/// positions as described by [`store`]. +/// the transposing algorithms mutate their positions as described by [`store`]. pub trait StoreValue: sealed::StoreValue { #[doc(hidden)] fn store( @@ -421,7 +453,17 @@ where boundary::store(destination, self, options.algorithm(Vectorize)); } sealed::TRANSPOSE => { - boundary::store_transpose(destination, self, options.algorithm(Transpose)); + boundary::store_permuting(destination, self, options.algorithm(Transpose)); + } + sealed::WARP_TRANSPOSE => { + boundary::store_permuting(destination, self, options.algorithm(WarpTranspose)); + } + sealed::WARP_TRANSPOSE_TIMESLICED => { + boundary::store_permuting( + destination, + self, + options.algorithm(WarpTransposeTimesliced), + ); } _ => panic!(), } @@ -451,6 +493,8 @@ mod sealed { pub const STRIPED: u8 = 1; pub const TRANSPOSE: u8 = 2; pub const VECTORIZE: u8 = 3; + pub const WARP_TRANSPOSE: u8 = 4; + pub const WARP_TRANSPOSE_TIMESLICED: u8 = 5; pub trait Algorithm { const KIND: u8; } diff --git a/rust/cuda-coop/src/load_store/boundary.rs b/rust/cuda-coop/src/load_store/boundary.rs index fdf32328c12..131845aa2a7 100644 --- a/rust/cuda-coop/src/load_store/boundary.rs +++ b/rust/cuda-coop/src/load_store/boundary.rs @@ -2,8 +2,8 @@ // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception use super::{ - Algorithm, LoadOptions, StoreOptions, Transpose, - sealed::{DIRECT, STRIPED, TRANSPOSE, VECTORIZE}, + Algorithm, LoadOptions, StoreOptions, + sealed::{DIRECT, STRIPED, TRANSPOSE, VECTORIZE, WARP_TRANSPOSE, WARP_TRANSPOSE_TIMESLICED}, }; use crate::{Alignment, Numeric, SupportedAlignment, ThreadData}; use cuda_device::{CollectiveTiles, DisjointSlice, library, thread}; @@ -45,6 +45,47 @@ registration!( "COOP_TRANSPOSE_VALIDITY_STORE_PARTIAL" ); +registration!( + warp_transpose_load_full, + "COOP_WARP_TRANSPOSE_VALIDITY_LOAD_FULL" +); +registration!( + warp_transpose_load_masked, + "COOP_WARP_TRANSPOSE_VALIDITY_LOAD_MASKED" +); +registration!( + warp_transpose_load_filled, + "COOP_WARP_TRANSPOSE_VALIDITY_LOAD_FILLED" +); +registration!( + warp_transpose_store_full, + "COOP_WARP_TRANSPOSE_VALIDITY_STORE_FULL" +); +registration!( + warp_transpose_store_partial, + "COOP_WARP_TRANSPOSE_VALIDITY_STORE_PARTIAL" +); +registration!( + warp_transpose_timesliced_load_full, + "COOP_WARP_TRANSPOSE_TIMESLICED_VALIDITY_LOAD_FULL" +); +registration!( + warp_transpose_timesliced_load_masked, + "COOP_WARP_TRANSPOSE_TIMESLICED_VALIDITY_LOAD_MASKED" +); +registration!( + warp_transpose_timesliced_load_filled, + "COOP_WARP_TRANSPOSE_TIMESLICED_VALIDITY_LOAD_FILLED" +); +registration!( + warp_transpose_timesliced_store_full, + "COOP_WARP_TRANSPOSE_TIMESLICED_VALIDITY_STORE_FULL" +); +registration!( + warp_transpose_timesliced_store_partial, + "COOP_WARP_TRANSPOSE_TIMESLICED_VALIDITY_STORE_PARTIAL" +); + #[inline(always)] pub(super) fn tile() -> usize { let threads = thread::blockDim_x() as usize @@ -115,7 +156,7 @@ pub(super) unsafe fn load< // 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, $striped:ident, $vectorize:ident, $transpose:ident, $controls:expr) => { + ($direct:ident, $striped:ident, $vectorize:ident, $transpose:ident, $warp:ident, $timesliced:ident, $controls:expr) => { match A::KIND { DIRECT => library::checked_block_resource_call_v1::( $direct, memory, $controls, payload, @@ -137,6 +178,24 @@ pub(super) unsafe fn load< payload, ) } + WARP_TRANSPOSE => { + library::checked_block_storage_call_v1::( + $warp, + memory, + (), + $controls, + payload, + ) + } + WARP_TRANSPOSE_TIMESLICED => { + library::checked_block_storage_call_v1::( + $timesliced, + memory, + (), + $controls, + payload, + ) + } _ => panic!(), } }; @@ -144,7 +203,7 @@ pub(super) unsafe fn load< // 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. + // compiler allocates the exact algorithm scratch and completes its previous use. unsafe { match (options.valid_items, options.oob_default) { (None, None) => call!( @@ -152,6 +211,8 @@ pub(super) unsafe fn load< striped_load_full, vectorize_load_full, transpose_load_full, + warp_transpose_load_full, + warp_transpose_timesliced_load_full, (pointer,) ), (Some(count), None) => { @@ -160,6 +221,8 @@ pub(super) unsafe fn load< striped_load_masked, vectorize_load_masked, transpose_load_masked, + warp_transpose_load_masked, + warp_transpose_timesliced_load_masked, (pointer, count as i32) ) } @@ -168,6 +231,8 @@ pub(super) unsafe fn load< striped_load_filled, vectorize_load_filled, transpose_load_filled, + warp_transpose_load_filled, + warp_transpose_timesliced_load_filled, (pointer, count as i32, default) ), (None, Some(_)) => panic!(), @@ -253,15 +318,16 @@ pub(super) fn store< } #[inline(always)] -pub(super) fn store_transpose< +pub(super) fn store_permuting< T: Numeric, const ITEMS: usize, const ALIGN: usize, const CAPACITY: usize, + A: Algorithm, >( destination: &mut DisjointSlice<'_, T, CollectiveTiles>, value: &mut ThreadData, - options: StoreOptions, + options: StoreOptions, ) where Alignment: SupportedAlignment, { @@ -279,20 +345,52 @@ pub(super) fn store_transpose< // that mutation. The compiler plans total scratch and completion before reuse. unsafe { let pointer = pointer.add(options.offset); + macro_rules! call { + ($transpose:ident, $warp:ident, $timesliced:ident, $controls:expr) => { + match A::KIND { + TRANSPOSE => { + library::checked_block_storage_call_v1::( + $transpose, + destination, + (), + $controls, + payload, + ) + } + WARP_TRANSPOSE => { + library::checked_block_storage_call_v1::( + $warp, + destination, + (), + $controls, + payload, + ) + } + WARP_TRANSPOSE_TIMESLICED => { + library::checked_block_storage_call_v1::( + $timesliced, + destination, + (), + $controls, + payload, + ) + } + _ => panic!(), + } + }; + } match options.valid_items { - None => library::checked_block_storage_call_v1::( + None => call!( transpose_store_full, - destination, - (), - (pointer,), - payload, + warp_transpose_store_full, + warp_transpose_timesliced_store_full, + (pointer,) ), - Some(count) => library::checked_block_storage_call_v1::( + Some(count) => call!( transpose_store_partial, - destination, - (), - (pointer, count as i32), - payload, + warp_transpose_store_partial, + warp_transpose_timesliced_store_partial, + (pointer, count as i32) ), } } diff --git a/rust/cuda-coop/tests/check_load_store.py b/rust/cuda-coop/tests/check_load_store.py index 349c5d65b98..949e9804162 100644 --- a/rust/cuda-coop/tests/check_load_store.py +++ b/rust/cuda-coop/tests/check_load_store.py @@ -193,6 +193,29 @@ ), ) +# Both algorithms participate as complete blocks and mutate the payload. +for name, selector in [ + ("warp_transpose", "WarpTranspose"), + ("warp_timesliced", "WarpTransposeTimesliced"), +]: + warp_load = transpose_load.replace("coop::Transpose", f"coop::{selector}") + warp_store = transpose_store.replace("coop::Transpose", f"coop::{selector}") + cases[name + "_shared"] = ( + load + warp_store.replace("&mut values", "&values"), + f"StoreValue", + ) + cases[name + "_managed_reuse"] = (warp_load + warp_store + warp_store, None) + cases[name + "_scalar"] = ( + warp_load + + f"coop::store(block,&mut output,-17_i32,coop::StoreOptions::default().offset(offset).algorithm(coop::{selector}));", + None, + ) + cases[name + "_partial_block"] = (warp_load, "divisible by 32") + cases[name + "_warp_only_participation"] = ( + "if thread::threadIdx_x()<32 {" + warp_load + "}", + "lane-varying", + ) + command = [ str(args.cargo_oxide.resolve()), "oxide", @@ -213,7 +236,12 @@ ) results = [] for name, (body, diagnostic) in cases.items(): - text = prefix + body + "\n}\nfn main() {}\n" + case_prefix = ( + prefix.replace("block=(128,1,1)", "block=(33,1,1)") + if name.endswith("_partial_block") + else prefix + ) + text = case_prefix + body + "\n}\nfn main() {}\n" source.write_text(text) (output / f"{name}.rs").write_text(text) run = subprocess.run( diff --git a/rust/cuda-coop/tests/load_store/src/block_warp.rs b/rust/cuda-coop/tests/load_store/src/block_warp.rs new file mode 100644 index 00000000000..42e058a840b --- /dev/null +++ b/rust/cuda-coop/tests/load_store/src/block_warp.rs @@ -0,0 +1,366 @@ +// Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +//! Full-block warp exchange observes every payload slot before and after 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< + T: coop::Numeric + PartialEq, + const ITEMS: usize, + const CAPACITY: usize, + A: coop::Algorithm, + >( + input: &[T], + output: &mut DisjointSlice>, + valid: usize, + default: T, + mode: u32, + algorithm: A, + ) { + let block = coop::this_block(); + 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( + block, + input, + &mut values, + coop::LoadOptions::default() + .offset(offset) + .algorithm(algorithm), + ); + } + if mode == 1 || mode == 5 || mode == 6 { + coop::load( + block, + input, + &mut values, + coop::LoadOptions::default() + .offset(offset) + .valid_items(valid) + .algorithm(algorithm), + ); + } else if mode == 2 || mode == 4 { + coop::load( + block, + input, + &mut values, + coop::LoadOptions::default() + .offset(offset) + .valid_items(valid) + .oob_default(default) + .algorithm(algorithm), + ); + } 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; + } + } + let mut correct = true; + let mut item = 0; + while item < ITEMS { + correct &= values.get(item) + == expected( + input, + offset, + lane * ITEMS + item, + ITEMS, + valid, + default, + mode, + ); + item += 1; + } + if mode == 0 || mode == 2 { + coop::store( + block, + &mut *output, + &mut values, + coop::StoreOptions::default() + .offset(offset) + .algorithm(algorithm), + ); + } 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(algorithm), + ); + } + item = 0; + while item < ITEMS { + let source = (lane / 32) * (32 * ITEMS) + 32 * item + lane % 32; + 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(algorithm), + ); + item = 0; + while item < ITEMS { + let first = (lane / 32) * (32 * ITEMS) + 32 * item + lane % 32; + let source = + ((first / ITEMS) / 32) * (32 * ITEMS) + 32 * (first % ITEMS) + (first / ITEMS) % 32; + 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(algorithm), + ); + coop::store( + block, + &mut *output, + &mut values, + coop::StoreOptions::default() + .offset(offset) + .valid_items(valid) + .algorithm(algorithm), + ); + item = 0; + while item < ITEMS { + correct &= values.get(item) + == expected( + input, + offset, + (lane / 32) * (32 * ITEMS) + 32 * item + lane % 32, + ITEMS, + valid, + default, + mode, + ); + item += 1; + } + } + assert!(correct); + } + #[kernel] + #[launch_contract(domain=3,block=(8,4,2),dynamic_shared=inferred)] + pub fn block_warp( + input: &[i32], + mut output: DisjointSlice>, + valid: usize, + mode: u32, + ) { + exercise::<_, ITEMS_PER_THREAD, 4, _>( + input, + &mut output, + valid, + -17, + mode, + coop::WarpTranspose, + ); + } + #[kernel] + #[launch_contract(domain=3,block=(8,4,2),dynamic_shared=inferred)] + pub fn block_warp_timesliced( + input: &[i32], + mut output: DisjointSlice>, + valid: usize, + mode: u32, + ) { + exercise::<_, ITEMS_PER_THREAD, 4, _>( + input, + &mut output, + valid, + -17, + mode, + coop::WarpTransposeTimesliced, + ); + } +} + +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, LaunchConfig3D}; + init_device_contexts(0, 1)?; + let context = with_cuda_context(0, |context| context.clone())?; + let stream = context.default_stream(); + // SAFETY: the typed declarations and device kernels are compiled together. + let module = unsafe { kernels::load_async(0)? }; + if let Some(trap) = trap { + let plan = module.prepare_block_warp_timesliced::<4>(LaunchConfig3D::auto_shared( + (1, 1, 1), + (8, 4, 2), + ))?; + let source = DeviceBuffer::from_host(&stream, &vec![11_i32; 256])?; + let target = DeviceBuffer::from_host(&stream, &vec![0_i32; 256])?; + let (valid, mode) = match trap { + "block-warp-count" => (257, 1), + "block-warp-unwritten" => (256, 9), + "block-warp-different-count" => (31, 5), + _ => panic!("unknown block warp trap"), + }; + let _ = module + .block_warp_timesliced_async_owned::<4, _, _>(&plan, source, target, valid, mode) + .sync()?; + panic!("expected block warp trap returned successfully"); + } + let mut cases = 0; + macro_rules! run { + ($prepare:ident,$launch:ident,$items:literal,$scratch:literal) => {{ + let plan = + module.$prepare::<$items>(LaunchConfig3D::auto_shared((2, 1, 1), (8, 4, 2)))?; + assert_eq!(plan.__raw_config().shared_mem_bytes, $scratch); + assert!( + module + .$prepare::<$items>(LaunchConfig3D::auto_shared((2, 1, 1), (64, 1, 1))) + .is_err() + ); + for mode in [0, 1, 2, 3, 4, 6] { + for valid in [0, 1, 32 * $items + 1, 64 * $items - 1, 64 * $items] { + if mode == 0 && valid != 64 * $items { + continue; + } + let input: Vec = (0..512).map(|i| i + 19).collect(); + let source = DeviceBuffer::from_host(&stream, &input)?; + let target = DeviceBuffer::from_host(&stream, &vec![0_i32; 512])?; + let (_, target) = module + .$launch::<$items, _, _>(&plan, source, target, valid, mode) + .sync()?; + let mut result = vec![0_i32; 512]; + target.copy_to_host(&stream, &mut result)?; + stream.synchronize()?; + for (index, &value) in result.iter().enumerate() { + let within = index % 256; + let written = if mode == 0 || mode == 2 { + 64 * $items + } else if mode == 6 { + valid / 2 + } else { + valid + }; + let expected = if within < written { + if mode == 2 && within >= valid { + -17 + } else { + input[index] + } + } else { + 0 + }; + assert_eq!( + value, + expected, + "{} items={} mode={} valid={} index={}", + stringify!($launch), + $items, + mode, + valid, + index + ); + } + cases += 1; + } + } + }}; + } + // CUB maximum plus the independent counted-Store bitmap plane, alignment16. + run!(prepare_block_warp, block_warp_async_owned, 1, 272); + run!(prepare_block_warp, block_warp_async_owned, 4, 1072); + run!( + prepare_block_warp_timesliced, + block_warp_timesliced_async_owned, + 1, + 144 + ); + run!( + prepare_block_warp_timesliced, + block_warp_timesliced_async_owned, + 4, + 560 + ); + println!("PASS {cases} public block warp Load/Store cases"); + Ok(()) +} diff --git a/rust/cuda-coop/tests/load_store/src/main.rs b/rust/cuda-coop/tests/load_store/src/main.rs index 4e3e20e33dd..ac6a04ca133 100644 --- a/rust/cuda-coop/tests/load_store/src/main.rs +++ b/rust/cuda-coop/tests/load_store/src/main.rs @@ -2,6 +2,7 @@ // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception use coop::{Numeric, ThreadData, this_block}; +mod block_warp; mod striped; mod transpose; mod vectorize; @@ -367,6 +368,12 @@ mod kernels { fn main() -> Result<(), Box> { if let Some(argument) = std::env::args().nth(1) { + if argument == "block-warp" { + return block_warp::run(None); + } + if argument.starts_with("block-warp-") { + return block_warp::run(Some(&argument)); + } if argument == "vectorize" { return vectorize::run(None); } From 4fa55e60c45a66efa6ddce6950224d5b2b569b94 Mon Sep 17 00:00:00 2001 From: Trent Nelson Date: Fri, 23 Oct 2026 08:33:21 -0700 Subject: [PATCH 42/48] Align pr12088 with the stable compiler Keep the owning implementation and tests unchanged while applying its scoped compiler, lock and documentation compatibility updates. Signed-off-by: Trent Nelson --- rust/rust-toolchain.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/rust/rust-toolchain.toml b/rust/rust-toolchain.toml index 06d6cb6baf7..fcb8eb44bbf 100644 --- a/rust/rust-toolchain.toml +++ b/rust/rust-toolchain.toml @@ -1,6 +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" +channel = "1.99.1" components = ["clippy", "rustfmt"] profile = "minimal" From e753fe04393fa80a73bc4836311fbac80541c54a Mon Sep 17 00:00:00 2001 From: Trent Nelson Date: Fri, 23 Oct 2026 09:30:10 -0700 Subject: [PATCH 43/48] Use matching compiler artifact metadata Resolve artifact metadata from the pinned compiler source. Keep the owning implementation and dependency versions unchanged. Signed-off-by: Trent Nelson --- rust/cuda-coop/tests/device_queries/Cargo.lock | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/rust/cuda-coop/tests/device_queries/Cargo.lock b/rust/cuda-coop/tests/device_queries/Cargo.lock index 1575c3564e2..15179597328 100644 --- a/rust/cuda-coop/tests/device_queries/Cargo.lock +++ b/rust/cuda-coop/tests/device_queries/Cargo.lock @@ -439,8 +439,7 @@ checksum = "9f7c3e4beb33f85d45ae3e3a1792185706c8e16d043238c593331cc7cd313b50" [[package]] name = "oxide-artifacts" version = "0.2.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "74d9a941c0c43d2dceaf25f5c9ca4c9f9cceb28dfa5e4e3cfd14e476f003e7c9" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" dependencies = [ "object", ] From 64f5d94d9e44edce85794921b71631311e4e9bc1 Mon Sep 17 00:00:00 2001 From: Trent Nelson Date: Fri, 23 Oct 2026 10:19:57 -0700 Subject: [PATCH 44/48] Retain the query provider build dependency Record the existing provider build dependency in the query lock. Keep implementation and third-party dependency records unchanged. Signed-off-by: Trent Nelson --- rust/cuda-coop/tests/device_queries/Cargo.lock | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/rust/cuda-coop/tests/device_queries/Cargo.lock b/rust/cuda-coop/tests/device_queries/Cargo.lock index 15179597328..069a9b1a123 100644 --- a/rust/cuda-coop/tests/device_queries/Cargo.lock +++ b/rust/cuda-coop/tests/device_queries/Cargo.lock @@ -144,9 +144,20 @@ dependencies = [ 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.5.0" From 32757242dba53a6c06dfec22dfad4c3acd78bfb3 Mon Sep 17 00:00:00 2001 From: Trent Nelson Date: Fri, 23 Oct 2026 11:12:08 -0700 Subject: [PATCH 45/48] Align Direct resource-conflict diagnostics Use the existing compiler diagnostic for three original cases. Keep their program bodies and all other expectations unchanged. Signed-off-by: Trent Nelson --- rust/cuda-coop/tests/check_load_store.py | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/rust/cuda-coop/tests/check_load_store.py b/rust/cuda-coop/tests/check_load_store.py index b881828fd73..e125688acb8 100644 --- a/rust/cuda-coop/tests/check_load_store.py +++ b/rust/cuda-coop/tests/check_load_store.py @@ -72,15 +72,18 @@ "lane-varying", ), "early_return": ("if thread::threadIdx_x()!=0 { return; }" + load, "lane-varying"), - "missing_completion": (load + store + reload, "requires a full-block barrier"), + "missing_completion": ( + load + store + reload, + "requires a barrier covering both accesses", + ), "write_after_read": ( reload + "coop::store(block,&mut output,1_i32,coop::StoreOptions::default().offset(offset));", - "requires a full-block barrier", + "requires a barrier covering both accesses", ), "loop_without_completion": ( load + "let mut iteration=0;while iteration<2 {" + store + "iteration+=1;}", - "requires a full-block barrier", + "requires a barrier covering both accesses", ), "loop_with_completion": ( load From 44f98765f2a3ea9f655a82edd1fc4c4fd6ae8be7 Mon Sep 17 00:00:00 2001 From: Trent Nelson Date: Fri, 23 Oct 2026 11:47:04 -0700 Subject: [PATCH 46/48] Align Transpose completion diagnostic Use the current compiler diagnostic for the original case. Keep its program body and all other expectations unchanged. Signed-off-by: Trent Nelson --- rust/cuda-coop/tests/check_load_store.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/rust/cuda-coop/tests/check_load_store.py b/rust/cuda-coop/tests/check_load_store.py index 131be4beefb..e46d2217cda 100644 --- a/rust/cuda-coop/tests/check_load_store.py +++ b/rust/cuda-coop/tests/check_load_store.py @@ -126,7 +126,7 @@ ), transpose_no_trailing_completion=( transpose_load + transpose_store + reload, - "requires a full-block barrier", + "requires a barrier covering both accesses", ), transpose_varying_count=( transpose_load.replace( From 3afd04fb04049c1ed1c85eddbc83d76bf03fa507 Mon Sep 17 00:00:00 2001 From: Trent Nelson Date: Fri, 23 Oct 2026 12:55:04 -0700 Subject: [PATCH 47/48] Align Striped completion diagnostic Use the current compiler diagnostic for the original case. Keep its program body and all other expectations unchanged. Signed-off-by: Trent Nelson --- rust/cuda-coop/tests/check_load_store.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/rust/cuda-coop/tests/check_load_store.py b/rust/cuda-coop/tests/check_load_store.py index 58c2512d3e9..a77f6a369b5 100644 --- a/rust/cuda-coop/tests/check_load_store.py +++ b/rust/cuda-coop/tests/check_load_store.py @@ -171,7 +171,7 @@ ), striped_missing_completion=( striped_load + striped_store + reload, - "requires a full-block barrier", + "requires a barrier covering both accesses", ), ) From e38a261c3b4e84d5ee2ac3086dcb41730b95a0ea Mon Sep 17 00:00:00 2001 From: Trent Nelson Date: Fri, 23 Oct 2026 14:07:55 -0700 Subject: [PATCH 48/48] Align Vectorize completion diagnostic Use the current compiler diagnostic for the original case. Keep its program body and all other expectations unchanged. Signed-off-by: Trent Nelson --- rust/cuda-coop/tests/check_load_store.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/rust/cuda-coop/tests/check_load_store.py b/rust/cuda-coop/tests/check_load_store.py index beb2119427d..90f5c0db8be 100644 --- a/rust/cuda-coop/tests/check_load_store.py +++ b/rust/cuda-coop/tests/check_load_store.py @@ -192,7 +192,7 @@ ), vectorize_missing_completion=( vectorize_load + vectorize_store + reload, - "requires a full-block barrier", + "requires a barrier covering both accesses", ), )