diff --git a/.github/workflows/rust-host.yml b/.github/workflows/rust-host.yml new file mode 100644 index 000000000000..ad93cbf4b65b --- /dev/null +++ b/.github/workflows/rust-host.yml @@ -0,0 +1,43 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +name: Rust host checks + +on: + push: + branches: + - main + - "pull-request/[0-9]+" + paths: + - "rust/**" + - ".github/workflows/rust-host.yml" + workflow_dispatch: + +permissions: + contents: read + +concurrency: + group: ${{ github.workflow }}-${{ github.ref }} + cancel-in-progress: true + +jobs: + host: + runs-on: ubuntu-latest + defaults: + run: + working-directory: rust + shell: bash --noprofile --norc -euo pipefail {0} + env: + RUSTDOCFLAGS: -D warnings + steps: + - uses: actions/checkout@v6 + with: + persist-credentials: false + - name: Check the pinned Rust workspace + run: | + rustup toolchain install --no-self-update + rustup show active-toolchain + cargo fmt --all -- --check + cargo test --workspace --locked + cargo clippy --workspace --all-targets --locked -- -D warnings + cargo doc --workspace --no-deps --locked diff --git a/docs/index.rst b/docs/index.rst index e1edae43a940..8906c9170e1c 100644 --- a/docs/index.rst +++ b/docs/index.rst @@ -7,6 +7,7 @@ CUDA Core Compute Libraries cccl/index python/index + rust/index contributors/index maintainers/index infrastructure/index @@ -18,6 +19,8 @@ make CUDA C++ and Python more delightful. - :doc:`Python Libraries ` +- :doc:`Rust Development ` + - :doc:`Contributor Docs ` - :doc:`Maintainer Docs ` diff --git a/docs/rust/groups.rst b/docs/rust/groups.rst new file mode 100644 index 000000000000..4c5050e02914 --- /dev/null +++ b/docs/rust/groups.rst @@ -0,0 +1,99 @@ +.. 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 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 +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. 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. + +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 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 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 +last mode indexes an unwritten payload slot. 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 new file mode 100644 index 000000000000..4f235b07aa45 --- /dev/null +++ b/docs/rust/index.rst @@ -0,0 +1,151 @@ +.. Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. +.. SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +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 +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. 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 +``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. + +Supported provider forms +------------------------ + +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:: + :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. + +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 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 +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 same ``ThreadData`` representation without exposing an initialized slice or +a public native layout. + +Compile a provider +------------------ + +Use the pinned toolchain in ``rust/rust-toolchain.toml``. Host tests, lint checks, +and rustdoc builds need no CUDA installation. Running native compilation requires +a Linux CUDA toolkit with NVRTC LTO support; CUDA 13.3 is the initial validation +toolkit. A GPU and CUDA context are not required. + +From the CCCL checkout: + +.. code-block:: shell + + cd rust + cargo test --workspace + cargo clippy --workspace --all-targets -- -D warnings + RUSTDOCFLAGS="-D warnings" cargo doc --workspace --no-deps + cargo run -p cuda-coop-codegen --example specialize -- \ + /usr/local/cuda-13.3 compute_120 /tmp/block_load.ltoir + +The example below is the executed Cargo source. It obtains its CCCL source root +from its build location, passes an explicit toolkit root, compiles a filled Load, +and writes the resulting LTO-IR. It does not link or launch a device kernel. + +.. literalinclude:: ../../rust/cuda-coop-codegen/examples/specialize.rs + :language: rust + :start-after: //! Compile a host-side provider without a CUDA context or a device consumer. + +Native regression checks are opt-in: + +.. code-block:: shell + + CUDA_HOME=/usr/local/cuda-13.3 cargo test --workspace -- --include-ignored + +``CUDA_COOP_TEST_TARGET`` can select another supported NVRTC target; its default +is ``compute_120``. These checks compile the supported Direct/Transpose forms at +item counts 1 and 4, compare repeated LTO output, and check exact layout decoding +and specialization identity. They do not execute GPU kernels. + +Compiler integration and API reference +-------------------------------------- + +``Provider::from_checkout`` snapshots headers from the explicitly selected CCCL +checkout and toolkit, then loads NVRTC from that toolkit's ``lib64`` directory. +Those native libraries are trusted executable compiler inputs. No Python wheel, +working-directory header lookup, or runtime query kernel is involved. Reuse the +provider and batch requests for a common target into one NVRTC compilation. + +The same NVRTC program produces both LTO-IR and each CUB ``TempStorage`` layout. +The result distinguishes scratch-free specializations from an empty C++ type's +nonzero ``sizeof``. Input identities include the request, header contents, +provider source, loaded NVRTC/builtins bytes, version, target, and compiler options. +Symbols include a specialization digest. Request planning must assign each +specialization to exactly one batch per final link: a batch is one indivisible +LTO-IR buffer, so partially overlapping batches would duplicate strong definitions. + +Rustdoc is the API reference, including the generated ABI's pointer, validity, +mutation, and collective requirements. Open ``rust/target/doc/cuda_coop_codegen/index.html`` +after the documentation command above, or run ``cargo doc --no-deps --open`` from +``rust/``. Device discovery, lowering, linking, launch enforcement, and Cargo +invalidation of external compiler inputs belong to the compiler integration. + +The Oxide adapter's ``oxide::serve_inputs`` answers +``--oxide-library-inputs-v1`` with the selected header roots and NVRTC library +paths, without loading CUDA. A provider build script dispatches this query +alongside ``--oxide-library-request-v1``. Device registration getters use +``cuda_device::library::registration!("REGISTRATION_ENV_NAME")`` so the frontend +tracks those files, directory contents and selected toolkit before emitting +Cargo dependency information. Host checks and rustdoc do not query CUDA inputs. +Cargo's timestamp freshness rules apply; replacements with unchanged or older +timestamps may require an explicit rebuild. + +.. toctree:: + :maxdepth: 1 + + payload + groups + load_store diff --git a/docs/rust/load_store.rst b/docs/rust/load_store.rst new file mode 100644 index 000000000000..74cbe55cafe8 --- /dev/null +++ b/docs/rust/load_store.rst @@ -0,0 +1,171 @@ +.. Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. +.. SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +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 +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 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. + +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 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. 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 + :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. + +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. +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; +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 +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. 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 +----------------------------------- + +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. + +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 +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. + +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 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 +---------------- + +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``, ``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/docs/rust/payload.rst b/docs/rust/payload.rst new file mode 100644 index 000000000000..be5aebab8b98 --- /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 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 +------------------------------ + +``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/.gitignore b/rust/.gitignore new file mode 100644 index 000000000000..b83d22266ac8 --- /dev/null +++ b/rust/.gitignore @@ -0,0 +1 @@ +/target/ diff --git a/rust/AGENTS.md b/rust/AGENTS.md new file mode 100644 index 000000000000..6493480b108a --- /dev/null +++ b/rust/AGENTS.md @@ -0,0 +1,63 @@ +# Rust `cuda.coop` guidance + +Read the repository's root guidance first. Apply this file to the Rust +workspace, its examples and tests, and `docs/rust/`. + +## Design tenets + +The design tenets of `cuda.coop` are **performant, consistent, delightful**. + +- **Performant:** target negligible overhead over CUB through compile-time + planning, inlining and LTO. Avoid hidden runtime work. Inspect the final + device artifact and measure representative execution before making claims + about register use, spills or overhead. +- **Consistent:** preserve common operation names, group-first arguments, + defaults, CUB behavior, mutation and result ownership. Document necessary + Rust syntax and safety adaptations. Additional capabilities extend shared + behavior without changing it. +- **Delightful:** use idiomatic Rust, useful defaults, inferred element types, + reusable generic kernels, ordinary Cargo builds and actionable errors. + +## API and examples + +- Use the single `cuda_coop` API, normally with `use cuda_coop as coop;`. + State the supported Oxide compiler/toolchain explicitly. Do not duplicate + primitive namespaces to separate built-in operators from static callables. +- Make the main per-thread item count a kernel const generic named + `ITEMS_PER_THREAD`, propagated through helpers and host specialization. + Keep operation-specific auxiliary and output extents independent. +- Infer payload element types from typed producers or assignments while + preserving widths. Investigate inference failures instead of adding a + redundant element type to an ordinary example. Document full signatures. +- Load fills its payload in place and returns unit. Preserve fresh results + where the common operation requires them, and reflect Store mutation in + borrowing and validity. Undefined, partial and root-only results must not + become unrestricted initialized Rust values. +- Preserve hierarchy units, supported groups, storage sharing and explicit + synchronization obligations. A group descriptor or Rust borrow alone is + not a proof of collective participation or shared-memory safety. +- Required collective completion must not depend solely on `Drop` running. + +## Ownership and validation + +Keep CUB specialization, exact scratch layout queries and operation semantics +in the native host provider. Keep compiler discovery, lowering, allocation, +device linking and launch enforcement in Oxide. Host compiler libraries and +filesystem dependencies must stay out of device code. CMake consumers do not +acquire a Rust requirement. + +Every public item needs useful rustdoc: arguments and generic parameters, +defaults, supported types/groups, participation, layout, mutation and result +validity, scratch/synchronization, errors and examples. Every public unsafe +item needs a concrete `# Safety` contract. Use rustdoc as the API reference +and the CCCL Rust guide for complete programs and explanations. + +Prefer end-user primitive tests, with item counts 1 and 4 where supported. +Use focused compiler/provider fixtures for independently reviewable compiler +capabilities. Keep displayed examples and their executed source identical. +Distinguish host tests, rustdoc/doctests, device compilation, GPU execution, +sanitizers and current-head CI. A sanitizer pass is not a Rust soundness proof. + +Run formatting, applicable lints, warning-free documentation and meaningful +targeted tests. Correct foundational defects at their owning layer and +requalify dependents instead of hiding a redesign in a later primitive family. diff --git a/rust/Cargo.lock b/rust/Cargo.lock new file mode 100644 index 000000000000..b092c7b62893 --- /dev/null +++ b/rust/Cargo.lock @@ -0,0 +1,284 @@ +# 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" +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-device" +version = "0.2.1" +source = "git+https://github.com/tpn/cuda-rust.git?rev=204e4e6f721bd111bf5e0e00caa69dc672a4d378#204e4e6f721bd111bf5e0e00caa69dc672a4d378" +dependencies = [ + "cuda-macros", +] + +[[package]] +name = "cuda-macros" +version = "0.2.1" +source = "git+https://github.com/tpn/cuda-rust.git?rev=204e4e6f721bd111bf5e0e00caa69dc672a4d378#204e4e6f721bd111bf5e0e00caa69dc672a4d378" +dependencies = [ + "oxide-tile-export", + "proc-macro2", + "quote", + "reserved-oxide-symbols", + "serde", + "serde_json", + "syn 2.0.119", +] + +[[package]] +name = "digest" +version = "0.10.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9ed9a281f7bc9b7576e61468ba615a66a5c8cfdff42420a70aa82701a3b1e292" +dependencies = [ + "block-buffer", + "crypto-common", +] + +[[package]] +name = "generic-array" +version = "0.14.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "85649ca51fd72272d7821adaf274ad91c288277713d9c18820d8499a7ff69e9a" +dependencies = [ + "typenum", + "version_check", +] + +[[package]] +name = "itoa" +version = "1.0.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8f42a60cbdf9a97f5d2305f08a87dc4e09308d1276d28c869c684d7777685682" + +[[package]] +name = "libc" +version = "0.2.191" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "65cd87c9a6d60536580b316ad985726104730a5151c99845f2490428ef91240a" + +[[package]] +name = "libloading" +version = "0.8.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d7c4b02199fee7c5d21a5ae7d8cfa79a6ef5bb2fc834d6e9058e89c825efdc55" +dependencies = [ + "cfg-if", + "windows-link", +] + +[[package]] +name = "memchr" +version = "2.8.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cf8baf1c55e62ffcace7a9f06f4bd9cd3f0c4beb022d3b367256b91b87513d98" + +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +source = "git+https://github.com/tpn/cuda-rust.git?rev=204e4e6f721bd111bf5e0e00caa69dc672a4d378#204e4e6f721bd111bf5e0e00caa69dc672a4d378" +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=204e4e6f721bd111bf5e0e00caa69dc672a4d378#204e4e6f721bd111bf5e0e00caa69dc672a4d378" + +[[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 = "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 = "windows-link" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f0805222e57f7521d6a62e36fa9163bc891acd422f971defe97d64e70d0a4fe5" + +[[package]] +name = "zmij" +version = "1.0.23" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "29666d0abbfad1e3dc4dcf6144730dd3a3ab225bbbdac83319345b1b44ccfc1b" diff --git a/rust/Cargo.toml b/rust/Cargo.toml new file mode 100644 index 000000000000..c6fd608c3716 --- /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", "cuda-coop-codegen"] +resolver = "3" + +[workspace.package] +edition = "2024" +license = "Apache-2.0 WITH LLVM-exception" +rust-version = "1.85" + +[workspace.lints.rust] +missing_docs = "deny" +unsafe_op_in_unsafe_fn = "deny" + +[workspace.lints.rustdoc] +broken_intra_doc_links = "deny" diff --git a/rust/cuda-coop-codegen/Cargo.toml b/rust/cuda-coop-codegen/Cargo.toml new file mode 100644 index 000000000000..7a92fcaaded4 --- /dev/null +++ b/rust/cuda-coop-codegen/Cargo.toml @@ -0,0 +1,19 @@ +# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. +# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +[package] +name = "cuda-coop-codegen" +version = "0.0.0" +edition.workspace = true +license.workspace = true +rust-version.workspace = true +publish = false + +[dependencies] +libloading = "0.8" +serde = { version = "1", features = ["derive"] } +serde_json = "1" +sha2 = "0.10" + +[lints] +workspace = true diff --git a/rust/cuda-coop-codegen/examples/specialize.rs b/rust/cuda-coop-codegen/examples/specialize.rs new file mode 100644 index 000000000000..3a2cfe26364b --- /dev/null +++ b/rust/cuda-coop-codegen/examples/specialize.rs @@ -0,0 +1,35 @@ +// Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +//! Compile a host-side provider without a CUDA context or a device consumer. + +use cuda_coop_codegen::{Algorithm, Operation, Provider, Request}; +use std::{env, path::Path, time::Instant}; + +fn main() -> Result<(), Box> { + let arguments: Vec<_> = env::args().collect(); + if arguments.len() != 4 { + return Err("usage: specialize CUDA_TOOLKIT COMPUTE_TARGET OUTPUT_LTOIR".into()); + } + // This example is built inside the selected source checkout. Consumers + // supply their own explicit source root instead of relying on their cwd. + let checkout = Path::new(env!("CARGO_MANIFEST_DIR")).join("../.."); + let started = Instant::now(); + let provider = Provider::from_checkout(checkout, &arguments[1])?; + let initialized = started.elapsed(); + let request = Request::new(Operation::LoadFilled, Algorithm::Transpose, 128, 4)?; + let started = Instant::now(); + let artifact = provider.compile(&[request], &arguments[2])?; + let compiled = started.elapsed(); + std::fs::write(&arguments[3], &artifact.ltoir)?; + println!("symbol: {}", artifact.specializations[0].symbol); + println!( + "CUB layout: {:?}; scratch: {:?}; scratch free: {}", + artifact.specializations[0].cub_layout, + artifact.specializations[0].scratch, + artifact.specializations[0].scratch_free + ); + println!("initialization: {initialized:?}; compilation: {compiled:?}"); + print!("{}", artifact.compiler_log); + Ok(()) +} diff --git a/rust/cuda-coop-codegen/src/block_io.cuh b/rust/cuda-coop-codegen/src/block_io.cuh new file mode 100644 index 000000000000..d404f17e0396 --- /dev/null +++ b/rust/cuda-coop-codegen/src/block_io.cuh @@ -0,0 +1,62 @@ +// Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +#include +#include + +namespace cuda_coop_detail +{ +// Copying this cell never evaluates a possibly uninitialized numeric value. +// Conversion to T is used only after the caller's complete Store preflight. +template +struct alignas(T) opaque_cell +{ + unsigned char bytes[sizeof(T)]; + + __device__ opaque_cell() = default; + __device__ opaque_cell(const opaque_cell& other) + { + cuda::std::memcpy(bytes, other.bytes, sizeof(T)); + } + __device__ opaque_cell& operator=(const opaque_cell& other) + { + cuda::std::memcpy(bytes, other.bytes, sizeof(T)); + return *this; + } + __device__ opaque_cell& operator=(T value) + { + cuda::std::memcpy(bytes, &value, sizeof(T)); + return *this; + } + __device__ operator T() const + { + T value; + cuda::std::memcpy(&value, bytes, sizeof(T)); + return value; + } +}; + +// CUB forms a per-lane iterator before testing its valid count. Keep that +// displacement as an integer; form a pointer only for an actually stored slot. +template +struct guarded_output +{ + using value_type = T; + using difference_type = int; + using pointer = T*; + using reference = T&; + using iterator_category = cuda::std::random_access_iterator_tag; + + T* base; + int offset; + + __device__ guarded_output operator+(int delta) const + { + return {base, offset + delta}; + } + __device__ T& operator[](int index) const + { + return base[offset + index]; + } +}; +} // namespace cuda_coop_detail diff --git a/rust/cuda-coop-codegen/src/headers.rs b/rust/cuda-coop-codegen/src/headers.rs new file mode 100644 index 000000000000..dd55da52337e --- /dev/null +++ b/rust/cuda-coop-codegen/src/headers.rs @@ -0,0 +1,95 @@ +// Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +//! Snapshot compiler inputs so identity describes the actual compilation bytes. + +use crate::{Error, Result}; +use sha2::{Digest, Sha256}; +use std::{ + collections::BTreeMap, + ffi::CString, + fs, + path::{Path, PathBuf}, +}; + +pub(crate) struct Headers { + pub(crate) entries: Vec<(CString, CString)>, + pub(crate) digest: [u8; 32], +} + +impl Headers { + pub(crate) fn roots(checkout: &Path, toolkit: &Path) -> [(PathBuf, &'static str); 4] { + [ + (checkout.join("cub/cub"), "cub/"), + (checkout.join("libcudacxx/include"), ""), + (checkout.join("thrust/thrust"), "thrust/"), + (toolkit.join("include"), ""), + ] + } + + pub(crate) fn read(checkout: &Path, toolkit: &Path) -> Result { + let mut files = BTreeMap::new(); + // First entry wins: the selected checkout always owns CCCL headers, + // including when the toolkit also bundles an older CCCL installation. + for (root, prefix) in Self::roots(checkout, toolkit) { + visit(&root, &root, prefix, &mut files)?; + } + let mut hash = Sha256::new(); + let mut entries = Vec::with_capacity(files.len()); + for (name, contents) in files { + hash_field(&mut hash, name.as_bytes()); + hash_field(&mut hash, &contents); + let content = CString::new(contents) + .map_err(|_| Error::message(format!("NUL byte in header: {name}")))?; + entries.push((CString::new(name)?, content)); + } + Ok(Self { + entries, + digest: hash.finalize().into(), + }) + } +} + +fn visit( + root: &Path, + directory: &Path, + prefix: &str, + files: &mut BTreeMap>, +) -> Result<()> { + for entry in fs::read_dir(directory).map_err(|e| Error::input(directory, e))? { + let entry = entry?; + let path = entry.path(); + if path.is_dir() { + // Header-directory symlinks would need cycle handling. The include + // root itself may be a symlink (as in normal CUDA installations). + if entry.file_type()?.is_symlink() { + return Err(Error::message(format!( + "header directory symlink is unsupported: {}", + path.display() + ))); + } + visit(root, &path, prefix, files)?; + } else if path.is_file() { + let relative = path.strip_prefix(root).expect("child of include root"); + let name = format!( + "{prefix}{}", + relative + .to_str() + .ok_or_else(|| Error::message(format!( + "non-UTF-8 header path: {}", + path.display() + )))? + .replace('\\', "/") + ); + if let std::collections::btree_map::Entry::Vacant(slot) = files.entry(name) { + slot.insert(fs::read(&path).map_err(|e| Error::input(&path, e))?); + } + } + } + Ok(()) +} + +pub(crate) fn hash_field(hash: &mut Sha256, bytes: &[u8]) { + hash.update((bytes.len() as u64).to_le_bytes()); + hash.update(bytes); +} diff --git a/rust/cuda-coop-codegen/src/lib.rs b/rust/cuda-coop-codegen/src/lib.rs new file mode 100644 index 000000000000..3383bb5016a6 --- /dev/null +++ b/rust/cuda-coop-codegen/src/lib.rs @@ -0,0 +1,837 @@ +// 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 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 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 +//! [`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 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 +//! [`Artifact`]. Each returned device ABI is unsafe to call: its documented +//! participation, pointer, capacity, validity and synchronization obligations +//! belong to the compiler integration. This crate allocates no device memory, +//! links no final kernel and supplies no launch policy. +//! +//! # Host example +//! +//! ```no_run +//! use cuda_coop_codegen::{Algorithm, Operation, Provider, Request}; +//! +//! # fn main() -> Result<(), Box> { +//! let provider = Provider::from_checkout("/path/to/cccl", "/usr/local/cuda")?; +//! let request = Request::new(Operation::LoadFull, Algorithm::Transpose, 128, 4)?; +//! let artifact = provider.compile(&[request], "compute_120")?; +//! std::fs::write("block_load.ltoir", &artifact.ltoir)?; +//! println!("{}: {:?}", artifact.specializations[0].symbol, artifact.specializations[0].scratch); +//! # Ok(()) } +//! ``` +//! +//! The example compiles on the host; executing it requires a compatible Linux +//! CUDA toolkit. The `specialize` example writes the artifact and prints its +//! identity/layout. `cargo doc --no-deps` generates this API reference. + +#![deny(unsafe_code)] + +mod headers; +#[allow(unsafe_code)] +mod nvrtc; +pub mod oxide; + +use headers::{Headers, hash_field}; +use sha2::{Digest, Sha256}; +use std::{fmt, path::Path}; + +/// A provider input, loader or compiler error, including NVRTC's compile log. +#[derive(Debug)] +pub struct Error(String); +impl Error { + pub(crate) fn message(message: impl Into) -> Self { + Self(message.into()) + } + pub(crate) fn input(path: &Path, error: std::io::Error) -> Self { + Self(format!("cannot read {}: {error}", path.display())) + } +} +impl fmt::Display for Error { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.write_str(&self.0) + } +} +impl std::error::Error for Error {} +impl From for Error { + fn from(error: std::io::Error) -> Self { + Self(error.to_string()) + } +} +impl From for Error { + fn from(error: std::ffi::NulError) -> Self { + Self(error.to_string()) + } +} +/// Result returned by native provider operations. +pub type Result = std::result::Result; + +/// CUB block layout algorithm; both variants produce/consume blocked payloads. +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub enum Algorithm { + /// Consecutive items per thread, without shared scratch or a reuse barrier. + Direct, + /// Coalesced access with shared-memory transposition to/from blocked layout. + Transpose, +} +impl Algorithm { + fn token(self) -> &'static str { + match self { + Self::Direct => "DIRECT", + Self::Transpose => "TRANSPOSE", + } + } +} + +/// 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)`. + 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, + /// Full Load with separate value storage and packed validity bits; marks every slot initialized. + LoadFullValidity, + /// Counted Load without a default; clears validity bits outside the loaded prefix. + LoadMaskedValidity, + /// 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. + StoreFullValidity, + /// Counted Store after checking its write domain. + /// Transpose redistributes all values and validity bits, even outside that domain. + StorePartialValidity, +} +impl Operation { + fn is_load(self) -> bool { + 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 { + Self::LoadFull => "load_full", + Self::LoadFilled => "load_filled", + Self::StoreFull => "store_full", + Self::StorePartial => "store_partial", + Self::LoadFullValidity => "load_full_validity_bits_v1", + Self::LoadMaskedValidity => "load_masked_validity_bits_v1", + Self::LoadFilledValidity => "load_filled_validity_bits_v1", + Self::StoreFullValidity => "store_full_validity_bits_v1", + Self::StorePartialValidity => "store_partial_validity_bits_v1", + } + } +} + +/// 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 +/// uses runtime `valid` and `default_value` only when selected by [`Operation`]. +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub struct Request { + operation: Operation, + algorithm: Algorithm, + dtype: Dtype, + block: [u32; 3], + items_per_thread: u32, +} +impl Request { + /// 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 + /// 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 { + 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 either algorithm. + /// Legacy raw recipes retain their i32 ABI and algorithm support. + /// + /// # Errors + /// + /// Rejects the invalid shapes documented by [`Self::new`], or a non-i32 raw + /// recipe. + pub fn new_typed( + dtype: Dtype, + operation: Operation, + 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", + )); + } + 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.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", + )); + } + Ok(Self { + operation, + algorithm, + dtype, + block, + 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 numeric element type. + pub fn dtype(self) -> Dtype { + self.dtype + } + /// 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[0] * self.block[1] * self.block[2] + } + /// Number of 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 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. +/// 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. +/// 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. +/// +/// 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. 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. 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. + 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, + /// 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 for Transpose stores, including counted stores with zero writes. + /// This describes Store rearrangement, not whether Load writes its payload. + 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. + /// 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 + /// + /// 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"), + include_bytes!("block_io.cuh"), + ] { + 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(include_str!("block_io.cuh")); + source.push_str("\ntemplate \nextern __device__ unsigned char layout_probe;\n"); + 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(); + 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")); + 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, layouts) in specializations + .iter_mut() + .zip(output.layouts.chunks_exact(2)) + { + specialization.cub_layout = layouts[0]; + if !specialization.scratch_free { + specialization.scratch = layouts[1]; + } + } + 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 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 { + ("store", "BlockStore", "BLOCK_STORE", "", "Store") + }; + 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 = threads * items; + 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; + 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() { + 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 memory = if matches!( + op, + Operation::StorePartial | Operation::StorePartialValidity + ) { + "cuda_coop_detail::guarded_output{memory, 0}" + } else { + "memory" + }; + let payload_const = if op.validity() && !op.is_load() && !transpose { + "const " + } else { + "" + }; + let (cell, flags, body) = if op.validity() { + let domain = if matches!( + op, + Operation::LoadMaskedValidity | Operation::StorePartialValidity + ) { + format!("tid * {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)); + }} +#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} && ({domain})) bits |= 1u << bit; + }} + flags[byte] = static_cast(bits); + }}"# + ) + } 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 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); + }} + }}"# + ) + } 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); + }}"# + ) + } 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 ({copy_domain}) cuda::std::memcpy(values[i].bytes, bytes + sizeof(value_type) * i, sizeof(value_type)); + }} + operation.Store({memory}, values{operands});{restore}"# + ) + }; + ( + "cuda_coop_detail::opaque_cell", + if op.is_load() || transpose { + ", unsigned char* flags" + } else { + ", const unsigned char* flags" + }, + format!( + "auto* const bytes = reinterpret_cast<{payload_const}unsigned char*>(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 +#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}) {{ + const int tid = cub::RowMajorTid({block_x}, {block_y}, {block_z}); + {check} + {constructor} + {body} +}} +}} +"#, + choice = request.algorithm.token() + ) +} diff --git a/rust/cuda-coop-codegen/src/nvrtc.rs b/rust/cuda-coop-codegen/src/nvrtc.rs new file mode 100644 index 000000000000..fe0ca34e7658 --- /dev/null +++ b/rust/cuda-coop-codegen/src/nvrtc.rs @@ -0,0 +1,307 @@ +// Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +//! Native ABI boundary. Library handles outlive every copied entry point. + +use crate::{ + Error, Layout, Result, + headers::{Headers, hash_field}, +}; +use libloading::Library; +use sha2::{Digest, Sha256}; +use std::{ + ffi::{CStr, CString, c_char, c_int, c_void}, + fs, + path::{Path, PathBuf}, + ptr, +}; + +type Handle = *mut c_void; +type Status = c_int; +type Create = unsafe extern "C" fn( + *mut Handle, + *const c_char, + *const c_char, + c_int, + *const *const c_char, + *const *const c_char, +) -> Status; +type Destroy = unsafe extern "C" fn(*mut Handle) -> Status; +type AddExpression = unsafe extern "C" fn(Handle, *const c_char) -> Status; +type Compile = unsafe extern "C" fn(Handle, c_int, *const *const c_char) -> Status; +type GetSize = unsafe extern "C" fn(Handle, *mut usize) -> Status; +type GetBytes = unsafe extern "C" fn(Handle, *mut c_char) -> Status; +type GetName = unsafe extern "C" fn(Handle, *const c_char, *mut *const c_char) -> Status; +type GetError = unsafe extern "C" fn(Status) -> *const c_char; + +pub(crate) struct Compiler { + _library: Library, + _builtins: Library, + pub(crate) identity: [u8; 32], + create: Create, + destroy: Destroy, + add_expression: AddExpression, + compile_program: Compile, + log_size: GetSize, + log: GetBytes, + lto_size: GetSize, + lto: GetBytes, + name: GetName, + error: GetError, +} +pub(crate) struct Output { + pub(crate) ltoir: Vec, + pub(crate) layouts: Vec, + pub(crate) log: String, +} + +impl Compiler { + pub(crate) fn input_paths(toolkit: &Path) -> [PathBuf; 2] { + ["libnvrtc-builtins.so", "libnvrtc.so"].map(|name| toolkit.join("lib64").join(name)) + } + + pub(crate) fn load(toolkit: &Path) -> Result { + if !cfg!(target_os = "linux") { + return Err(Error::message( + "the native provider currently supports Linux hosts", + )); + } + let mut hash = Sha256::new(); + let mut load = |input: &Path| -> Result { + let path = input.canonicalize().map_err(|e| Error::input(input, e))?; + let name = input + .file_name() + .and_then(|name| name.to_str()) + .expect("fixed library name"); + hash_field(&mut hash, name.as_bytes()); + hash_field( + &mut hash, + &fs::read(&path).map_err(|e| Error::input(&path, e))?, + ); + // SAFETY: the explicit toolkit is trusted native executable input; + // only documented NVRTC entry points are called with their C ABI. + unsafe { Library::new(&path) } + .map_err(|e| Error::message(format!("cannot load {}: {e}", path.display()))) + }; + let [builtins_path, library_path] = Self::input_paths(toolkit); + let builtins = load(&builtins_path)?; + let library = load(&library_path)?; + // SAFETY: symbols use signatures from nvrtc.h. Missing entry points + // return errors before a program is created; handles remain owned. + unsafe { + let version = library + .get:: Status>(b"nvrtcVersion\0") + .map_err(|e| Error::message(format!("missing NVRTC entry point: {e}")))?; + let (mut major, mut minor) = (0, 0); + if version(&mut major, &mut minor) != 0 { + return Err(Error::message("nvrtcVersion failed")); + } + hash_field(&mut hash, &major.to_le_bytes()); + hash_field(&mut hash, &minor.to_le_bytes()); + macro_rules! symbol { + ($name:literal, $ty:ty) => { + *library + .get::<$ty>(concat!($name, "\0").as_bytes()) + .map_err(|e| { + Error::message(format!("missing NVRTC entry point {}: {e}", $name)) + })? + }; + } + Ok(Self { + create: symbol!("nvrtcCreateProgram", Create), + destroy: symbol!("nvrtcDestroyProgram", Destroy), + add_expression: symbol!("nvrtcAddNameExpression", AddExpression), + compile_program: symbol!("nvrtcCompileProgram", Compile), + log_size: symbol!("nvrtcGetProgramLogSize", GetSize), + log: symbol!("nvrtcGetProgramLog", GetBytes), + lto_size: symbol!("nvrtcGetLTOIRSize", GetSize), + lto: symbol!("nvrtcGetLTOIR", GetBytes), + name: symbol!("nvrtcGetLoweredName", GetName), + error: symbol!("nvrtcGetErrorString", GetError), + identity: hash.finalize().into(), + _library: library, + _builtins: builtins, + }) + } + } + fn check(&self, status: Status, operation: &str) -> Result<()> { + if status == 0 { + return Ok(()); + } + // SAFETY: NVRTC returns a non-null static NUL-terminated error string. + let message = unsafe { CStr::from_ptr((self.error)(status)) }.to_string_lossy(); + Err(Error::message(format!("{operation}: {message} ({status})"))) + } + pub(crate) fn compile( + &self, + source: &str, + options: &[String], + expressions: &[String], + headers: &Headers, + ) -> Result { + let source = CString::new(source)?; + let options = options + .iter() + .map(|s| CString::new(s.as_str())) + .collect::, _>>()?; + let option_pointers: Vec<_> = options.iter().map(|s| s.as_ptr()).collect(); + let header_names: Vec<_> = headers + .entries + .iter() + .map(|(name, _)| name.as_ptr()) + .collect(); + let header_contents: Vec<_> = headers + .entries + .iter() + .map(|(_, contents)| contents.as_ptr()) + .collect(); + let count = c_int::try_from(headers.entries.len()) + .map_err(|_| Error::message("too many input headers"))?; + let mut program = Program { + raw: ptr::null_mut(), + compiler: self, + }; + // SAFETY: C strings and arrays remain live until after program destroy; + // NVRTC initializes the uniquely owned handle. + self.check( + unsafe { + (self.create)( + &mut program.raw, + source.as_ptr(), + c"cuda_coop.cu".as_ptr(), + count, + header_contents.as_ptr(), + header_names.as_ptr(), + ) + }, + "nvrtcCreateProgram", + )?; + let expressions = expressions + .iter() + .map(|s| CString::new(s.as_str())) + .collect::, _>>()?; + for expression in &expressions { + // SAFETY: program and terminated expression remain live through compilation. + self.check( + unsafe { (self.add_expression)(program.raw, expression.as_ptr()) }, + "nvrtcAddNameExpression", + )?; + } + // SAFETY: program and option strings are live for this synchronous call. + let status = unsafe { + (self.compile_program)( + program.raw, + option_pointers.len() as c_int, + option_pointers.as_ptr(), + ) + }; + let mut log = self.bytes(program.raw, self.log_size, self.log)?; + if log.last() == Some(&0) { + log.pop(); + } + let log = String::from_utf8_lossy(&log).into_owned(); + self.check(status, "nvrtcCompileProgram") + .map_err(|e| Error::message(format!("{e}\n{log}")))?; + let mut layouts = Vec::with_capacity(expressions.len()); + for expression in &expressions { + let mut lowered = ptr::null(); + // SAFETY: successful compilation owns the returned name until destroy. + self.check( + unsafe { (self.name)(program.raw, expression.as_ptr(), &mut lowered) }, + "nvrtcGetLoweredName", + )?; + if lowered.is_null() { + return Err(Error::message("NVRTC returned a null layout name")); + } + // SAFETY: successful query returned a program-owned C string. + layouts.push(decode_layout( + unsafe { CStr::from_ptr(lowered) } + .to_str() + .map_err(|_| Error::message("non-UTF-8 layout name"))?, + )?); + } + let ltoir = self.bytes(program.raw, self.lto_size, self.lto)?; + if ltoir.is_empty() { + return Err(Error::message("NVRTC emitted empty LTO-IR")); + } + Ok(Output { + ltoir, + layouts, + log, + }) + } + fn bytes(&self, program: Handle, size: GetSize, get: GetBytes) -> Result> { + let mut length = 0; + // SAFETY: program is live; NVRTC writes one size_t. + self.check(unsafe { size(program, &mut length) }, "NVRTC output size")?; + let mut bytes = vec![0; length]; + if length != 0 { + // SAFETY: exact reported byte capacity is allocated and writable. + self.check( + unsafe { get(program, bytes.as_mut_ptr().cast()) }, + "NVRTC output", + )?; + } + Ok(bytes) + } +} +struct Program<'a> { + raw: Handle, + compiler: &'a Compiler, +} +impl Drop for Program<'_> { + fn drop(&mut self) { + if !self.raw.is_null() { + // SAFETY: unique owner; the compiler library remains loaded. + unsafe { + (self.compiler.destroy)(&mut self.raw); + } + } + } +} +fn decode_layout(name: &str) -> Result { + let invalid = || Error::message(format!("unsupported NVRTC layout name: {name}")); + let values = name + .strip_prefix("_Z12layout_probeILy") + .and_then(|s| s.strip_suffix("EE")) + .ok_or_else(invalid)?; + let (size, alignment) = values.split_once("ELy").ok_or_else(invalid)?; + let size: usize = size.parse().map_err(|_| invalid())?; + let alignment: usize = alignment.parse().map_err(|_| invalid())?; + if size == 0 || !alignment.is_power_of_two() || size % alignment != 0 { + return Err(invalid()); + } + Ok(Layout { size, alignment }) +} +#[cfg(test)] +mod tests { + use super::*; + #[test] + fn layout_literals_and_rejection() { + assert_eq!( + decode_layout("_Z12layout_probeILy1ELy1EE").unwrap(), + Layout { + size: 1, + alignment: 1 + } + ); + assert_eq!( + decode_layout("_Z12layout_probeILy2064ELy16EE").unwrap(), + Layout { + size: 2064, + alignment: 16 + } + ); + for invalid in [ + "_Z12layout_probeILy0ELy1EE", + "_Z12layout_probeILy32ELy3EE", + "_Z12layout_probeILy33ELy16EE", + "_Z12layout_probeILy32ELy0EE", + "_Z12layout_probeILy18446744073709551616ELy1EE", + "other", + "_Z12layout_probeILy32ELy16EE.extra", + ] { + assert!(decode_layout(invalid).is_err(), "{invalid}"); + } + } +} diff --git a/rust/cuda-coop-codegen/src/oxide.rs b/rust/cuda-coop-codegen/src/oxide.rs new file mode 100644 index 000000000000..08972f32f85e --- /dev/null +++ b/rust/cuda-coop-codegen/src/oxide.rs @@ -0,0 +1,647 @@ +// Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +//! Version-one host transport for the Oxide library-operation marker. +//! +//! A dependency build script can serve as its already-built provider executable: +//! normal Cargo execution embeds [`registration`] in a device crate constant; +//! `--oxide-library-request-v1` execution passes stdin/stdout to [`serve`]; +//! `--oxide-library-inputs-v1` passes them to [`serve_inputs`] during macro expansion. +//! Use `cuda_device::library::registration!` in the getter to track these inputs. +//! The compiler reads the registration getter's constant MIR and invokes this +//! host process. This module never invokes Cargo or executes device Rust code. +//! The protocol supports the 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. + +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}; + +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(Clone, Copy, Debug, Deserialize, Serialize, PartialEq)] +#[serde( + tag = "kind", + content = "bits", + rename_all = "snake_case", + deny_unknown_fields +)] +enum Scalar { + Signed(u16), + Unsigned(u16), + 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", + 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 { + let operation = match self.operation.as_str() { + "load_full" => Operation::LoadFull, + "load_filled" => Operation::LoadFilled, + "store_full" => Operation::StoreFull, + "store_partial" => Operation::StorePartial, + "load_full_validity_bits_v1" => Operation::LoadFullValidity, + "load_masked_validity_bits_v1" => Operation::LoadMaskedValidity, + "load_filled_validity_bits_v1" => Operation::LoadFilledValidity, + "store_full_validity_bits_v1" => Operation::StoreFullValidity, + "store_partial_validity_bits_v1" => Operation::StorePartialValidity, + _ => { + 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_typed_with_shape( + self.dtype.dtype()?, + operation, + algorithm, + self.block, + 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(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 shared_payload { + Argument::ConstPointer(dtype) + } else { + Argument::MutPointer(dtype) + }, + ]; + if operation.validity() { + arguments.push(if shared_payload { + Argument::ConstPointer(Scalar::Unsigned(8)) + } else { + Argument::MutPointer(Scalar::Unsigned(8)) + }); + } + if operation.counted() { + arguments.push(Argument::Value(Scalar::Signed(32))); + } + if operation.filled() { + arguments.push(Argument::Value(dtype)); + } + 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 +/// obtained from `std::env::current_exe()`. Embed the returned string in the +/// device crate through `cargo::rustc-env`; a named getter must return that +/// constant directly. The marker uses the getter's function-item type. +/// +/// # Errors +/// +/// Rejects a nonabsolute executable or a path that cannot be encoded as JSON. +pub fn registration( + executable: &Path, + operation: Operation, + algorithm: Algorithm, +) -> Result { + if !executable.is_absolute() { + return Err(Error::message("provider executable path must be absolute")); + } + serde_json::to_string(&Registration { + version: VERSION, + provider: "cuda_coop", + executable, + operation: operation.token(), + algorithm: match algorithm { + Algorithm::Direct => "direct", + Algorithm::Transpose => "transpose", + }, + }) + .map_err(|error| Error::message(format!("cannot encode provider registration: {error}"))) +} + +#[derive(Deserialize)] +#[serde(deny_unknown_fields)] +struct InputRequest { + version: u32, + cuda_root: PathBuf, + target: String, +} + +#[derive(Serialize)] +struct InputResponse { + version: u32, + paths: Vec, + environment: Vec, +} + +/// Declare native inputs for Oxide's registration macro without loading CUDA. +/// +/// `checkout` selects the same CCCL tree as [`serve`]. The bounded version-one +/// JSON request contains `cuda_root` and `target`; the response lists header +/// directories and NVRTC libraries. The macro tracks files and directories +/// recursively before rustc emits dep-info, including new header files. +/// No GPU, native compiler invocation or library loading occurs here. +/// Cargo's timestamp freshness rules apply; replacements with unchanged or older +/// timestamps can remain undetected. Rebuild after restoring or replacing inputs +/// this way, including a symlink switched to an older installation. +/// +/// # Errors +/// +/// Rejects malformed, oversized or unsupported requests, nonabsolute toolkit +/// paths, invalid targets, and inaccessible checkout paths. +pub fn serve_inputs(checkout: &Path, input: impl Read, mut output: impl Write) -> Result<()> { + let mut bytes = Vec::new(); + input + .take((MAX_HEADER + 1) as u64) + .read_to_end(&mut bytes)?; + if bytes.len() > MAX_HEADER { + return Err(Error::message("provider input request exceeds 1 MiB")); + } + let request: InputRequest = serde_json::from_slice(&bytes) + .map_err(|error| Error::message(format!("invalid provider input request: {error}")))?; + if request.version != VERSION + || !request.cuda_root.is_absolute() + || !request.target.starts_with("compute_") + || request.target.len() <= 8 + || !request.target[8..] + .chars() + .all(|character| character.is_ascii_alphanumeric()) + { + return Err(Error::message( + "provider inputs require version 1, an absolute toolkit path and compute_XX target", + )); + } + let checkout = checkout + .canonicalize() + .map_err(|error| Error::input(checkout, error))?; + let paths = crate::headers::Headers::roots(&checkout, &request.cuda_root) + .into_iter() + .map(|(path, _)| path) + .chain(crate::nvrtc::Compiler::input_paths(&request.cuda_root)) + .collect(); + let response = serde_json::to_vec(&InputResponse { + version: VERSION, + paths, + environment: Vec::new(), + }) + .map_err(|error| Error::message(format!("cannot encode provider inputs: {error}")))?; + if response.len() > MAX_HEADER { + return Err(Error::message("provider input response exceeds 1 MiB")); + } + output.write_all(&response)?; + Ok(()) +} + +/// Compile one versioned stdin request and write its framed native response. +/// +/// `checkout` selects the CCCL checkout whose headers the provider captures. +/// The request supplies an explicit toolkit path and target. Stdout contains +/// only the frame; the caller reports errors on stderr and exits nonzero. +/// The 24-byte prefix is `OXDLIB01`, JSON header length, and raw LTO length +/// (both lengths are little-endian u64), followed by the two byte buffers. +/// Requests and response headers are limited to 1 MiB; LTO is limited to 256 MiB. +/// Runtime arguments are a scalar/raw-pointer tuple; scratch is inserted at +/// parameter two and uses a generic byte pointer after shared address conversion. +/// +/// # Errors +/// +/// Rejects malformed/version-mismatched requests, unsupported types/recipes, +/// invalid geometry, and duplicate specializations. Native compiler errors +/// retain NVRTC diagnostics. No response is written until compilation succeeds. +pub fn serve(checkout: &Path, input: impl Read, mut output: impl Write) -> Result<()> { + let mut bytes = Vec::new(); + input + .take((MAX_HEADER + 1) as u64) + .read_to_end(&mut bytes)?; + if bytes.len() > MAX_HEADER { + return Err(Error::message("provider request exceeds 1 MiB")); + } + let request: BatchRequest = serde_json::from_slice(&bytes) + .map_err(|error| Error::message(format!("invalid provider request: {error}")))?; + if request.version != VERSION { + return Err(Error::message(format!( + "unsupported provider protocol {}; expected {VERSION}", + request.version + ))); + } + let requests: Vec<_> = request + .requests + .iter() + .map(SpecializationRequest::native) + .collect::>()?; + let artifact = Provider::from_checkout(checkout, &request.cuda_root)? + .compile(&requests, &request.target)?; + write_response(&requests, artifact, &mut output) +} + +fn write_response(requests: &[Request], artifact: Artifact, output: &mut impl Write) -> Result<()> { + if artifact.ltoir.is_empty() || artifact.ltoir.len() > MAX_LTO { + return Err(Error::message( + "native LTO buffer must contain 1..=256 MiB bytes", + )); + } + let specializations = artifact + .specializations + .into_iter() + .zip(requests) + .enumerate() + .map(|(index, (entry, request))| { + Ok(SpecializationResponse { + request_index: u32::try_from(index) + .map_err(|_| Error::message("request index exceeds u32"))?, + identity: entry.identity, + symbol: entry.symbol, + arguments: arguments(*request), + 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 typed_validity_abi_preserves_width_and_separate_flags() { + // Old byte-per-slot registrations must not bind to the packed format. + assert!( + SpecializationRequest { + operation: "load_filled_validity_v1".into(), + algorithm: "direct".into(), + dtype: Scalar::Signed(32), + block: [128, 1, 1], + items_per_thread: 4, + } + .native() + .is_err() + ); + 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); + 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), + ] + ); + } + } + for operation in [ + Operation::StoreFullValidity, + Operation::StorePartialValidity, + ] { + 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); + } + } + 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")); + let toolkit = std::env::temp_dir().join("cuda-coop-input-query-no-toolkit"); + let request = serde_json::json!({"version":1,"cuda_root":toolkit,"target":"compute_120"}); + let mut bytes = Vec::new(); + serve_inputs(checkout, request.to_string().as_bytes(), &mut bytes).unwrap(); + let response: serde_json::Value = serde_json::from_slice(&bytes).unwrap(); + assert_eq!(response["version"], 1); + assert_eq!(response["paths"].as_array().unwrap().len(), 6); + assert!( + response["paths"] + .as_array() + .unwrap() + .iter() + .all(|path| Path::new(path.as_str().unwrap()).is_absolute()) + ); + assert!(response["environment"].as_array().unwrap().is_empty()); + for request in [ + serde_json::json!({"version":2,"cuda_root":toolkit,"target":"compute_120"}), + serde_json::json!({"version":1,"cuda_root":"relative","target":"compute_120"}), + serde_json::json!({"version":1,"cuda_root":toolkit,"target":"invalid"}), + ] { + assert!(serve_inputs(checkout, request.to_string().as_bytes(), Vec::new()).is_err()); + } + assert!(serve_inputs(checkout, vec![b' '; MAX_HEADER + 1].as_slice(), Vec::new()).is_err()); + } + + #[test] + fn registration_binds_native_recipe() { + let value: serde_json::Value = serde_json::from_str( + ®istration( + Path::new("/provider"), + Operation::LoadFilled, + Algorithm::Transpose, + ) + .unwrap(), + ) + .unwrap(); + assert_eq!(value["version"], 1); + assert_eq!(value["operation"], "load_filled"); + assert_eq!(value["algorithm"], "transpose"); + assert!( + registration( + Path::new("relative"), + Operation::LoadFull, + Algorithm::Direct + ) + .is_err() + ); + } + + #[test] + fn 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":[{ + "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!(0), "block threads"), + ( + "/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/check_carrier_rejections.py b/rust/cuda-coop-codegen/tests/check_carrier_rejections.py new file mode 100644 index 000000000000..3a6b64af8802 --- /dev/null +++ b/rust/cuda-coop-codegen/tests/check_carrier_rejections.py @@ -0,0 +1,326 @@ +# 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", + ), + "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" + "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", + ), + "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);", + "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, + ), + "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);", + 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 = """ +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)] +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; +} +#[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 + 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" + ) + 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=case_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/check_collective_rejections.py b/rust/cuda-coop-codegen/tests/check_collective_rejections.py new file mode 100644 index 000000000000..7f422c751fca --- /dev/null +++ b/rust/cuda-coop-codegen/tests/check_collective_rejections.py @@ -0,0 +1,483 @@ +# 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 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( + 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", + ), + "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( + " 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) + ) + }, + 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": ( + { + "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", + ), + "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( + 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", "atomic_control"]: + 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", + ) +cases["reborrow_uniform"] = ( + { + "helper": cases["projected_reborrow"][0]["helper"].replace( + "let valid = pair.0;", "let valid = pair.1;" + ) + }, + 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, +) +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", + ), + "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() +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", + "atomic_control", + ]: + 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" + 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( + 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/check_storage_rejections.py b/rust/cuda-coop-codegen/tests/check_storage_rejections.py new file mode 100644 index 000000000000..dde9b2417f83 --- /dev/null +++ b/rust/cuda-coop-codegen/tests/check_storage_rejections.py @@ -0,0 +1,236 @@ +# 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" + 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=case_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) +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/check_validity_rejections.py b/rust/cuda-coop-codegen/tests/check_validity_rejections.py new file mode 100644 index 000000000000..97b9399c70ed --- /dev/null +++ b/rust/cuda-coop-codegen/tests/check_validity_rejections.py @@ -0,0 +1,120 @@ +# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. +# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +"""Compile private Direct validity controls and resource-epoch contrasts.""" + +import argparse +import json +import os +import subprocess +from pathlib import Path + +parser = argparse.ArgumentParser(description=__doc__) +parser.add_argument("output", type=Path) +args = parser.parse_args() +output = args.output.resolve() +record = json.loads((output / "command.json").read_text()) +source = output / "fixture/carrier/src/validity.rs" +original = source.read_text() +command = record["argv"].copy() +command[command.index("run")] = "build" +# build forwards binary selection to Cargo; run has its own --bin option. +command.insert(command.index("--bin"), "--") +env = os.environ.copy() +env.update( + CCCL_SOURCE=record.get("cccl_source", str(Path(__file__).resolve().parents[3])), + CUDA_OXIDE_BACKEND=record["backend"], + CUDA_HOME=record["cuda_root"], + CUDA_PATH=record["cuda_root"], +) +prefix = """// Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception +use coop_marker_fixture_api::{ThreadData, this_block, validity}; +use cuda_device::{CollectiveTiles, DisjointSlice, kernel, launch_contract, thread}; +#[kernel] +#[launch_contract(domain = 1, block = (128, 1, 1), dynamic_shared = inferred)] +pub fn probe(input: &[i32], mut output: DisjointSlice>, valid: usize) { + let group = this_block(); + let offset = thread::blockIdx_x() as usize * 512; + let mut values = ThreadData::<_, 4>::new(); +""" +suffix = "\n}\nfn main() {}\n" +cases = { + "varying_count": ( + "validity::load_masked(group, input, offset, thread::threadIdx_x() as usize, &mut values);", + "not proved block-uniform", + ), + "varying_default": ( + "validity::load_filled(group, input, offset, valid, thread::threadIdx_x() as i32, &mut values);", + "not proved block-uniform", + ), + "varying_offset": ( + "validity::load_masked(group, input, offset + thread::threadIdx_x() as usize, valid, &mut values);", + "not proved block-uniform", + ), + "missing_barrier": ( + "validity::load_filled(group, input, offset, valid, 0, &mut values);\n" + "validity::store_partial(group, &mut output, offset, valid, &mut values);\n" + "validity::reload_filled(group, &mut output, offset, valid, 0, &mut values);", + "requires a full-block barrier", + ), + "varying_participation": ( + "if thread::threadIdx_x() & 1 == 0 {\n" + "validity::load_masked(group, input, offset, valid, &mut values);\n}", + "lane-varying", + ), + "loop_without_completion": ( + "validity::load_filled(group, input, offset, valid, 0, &mut values);\n" + "let mut count = 2; while count > 0 {\n" + "validity::store_partial(group, &mut output, offset, valid, &mut values);\n" + "count -= 1; }", + "requires a full-block barrier", + ), + "explicit_completion": ( + "validity::load_filled(group, input, offset, valid, 0, &mut values);\n" + "validity::store_partial(group, &mut output, offset, valid, &mut values);\n" + "thread::sync_threads();\n" + "validity::reload_filled(group, &mut output, offset, valid, 0, &mut values);", + None, + ), + "reconverged_payload": ( + "let mut other = ThreadData::<_, 4>::new();\n" + "values.write(0, 7); other.write(0, 9);\n" + "let mut selected = if thread::threadIdx_x() & 1 == 0 { values } else { other };\n" + "validity::store_partial(group, &mut output, offset, 1, &mut selected);", + None, + ), +} +results = [] +try: + for name, (body, diagnostic) in cases.items(): + source.write_text(prefix + body + suffix) + result = subprocess.run( + command, + cwd=record["cwd"], + env=env, + text=True, + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + ) + (output / f"validity-{name}.log").write_text(result.stdout) + passed = ( + result.returncode == 0 + if diagnostic is None + else result.returncode != 0 and diagnostic in result.stdout + ) + results.append( + { + "case": name, + "expected": diagnostic, + "returncode": result.returncode, + "passed": passed, + } + ) + print(f"{'PASS' if passed else 'FAIL'} {name}", flush=True) + (output / "validity-rejections.json").write_text( + json.dumps(results, indent=2) + "\n" + ) +finally: + source.write_text(original) +raise SystemExit(0 if all(row["passed"] for row in results) else 1) diff --git a/rust/cuda-coop-codegen/tests/native.rs b/rust/cuda-coop-codegen/tests/native.rs new file mode 100644 index 000000000000..1c782d59df39 --- /dev/null +++ b/rust/cuda-coop-codegen/tests/native.rs @@ -0,0 +1,332 @@ +// 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, Dtype, 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()); + } + 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!( + 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> { + 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 (algorithm, threads) in [(Algorithm::Direct, 128), (Algorithm::Transpose, 33)] { + 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, threads, items, + )?); + } + } + } + } + // Partial warps, complete warps, maximum block size, CUB bank padding and + // packed-mask byte boundaries exercise different actual storage layouts. + for (threads, items) in [ + (1, 4), + (31, 4), + (32, 4), + (128, 4), + (130, 4), + (1024, 1), + (33, 8), + (33, 9), + ] { + for operation in [ + Operation::LoadMaskedValidity, + Operation::StoreFullValidity, + Operation::StorePartialValidity, + ] { + requests.push(Request::new( + operation, + Algorithm::Transpose, + threads, + items, + )?); + } + } + // The first shape has a partial physical warp crossing rows; the second + // spans four Z planes, one warp each, and matches an existing 1D total. + for block in [[11, 3, 1], [8, 4, 4]] { + for algorithm in [Algorithm::Direct, Algorithm::Transpose] { + for items in [1, 4] { + for operation in [ + Operation::LoadFullValidity, + Operation::LoadMaskedValidity, + Operation::LoadFilledValidity, + Operation::StoreFullValidity, + Operation::StorePartialValidity, + ] { + requests.push(Request::new_with_shape(operation, algorithm, block, 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, request) in first.specializations.iter().zip(&requests) { + assert!(identities.insert(entry.identity)); + eprintln!( + "{request:?}: CUB {:?}, scratch {:?}", + entry.cub_layout, entry.scratch + ); + if request.algorithm() == Algorithm::Direct { + 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); + } else { + assert!(!entry.scratch_free); + assert_eq!( + entry.mutates_payload, + matches!( + request.operation(), + Operation::StoreFullValidity | Operation::StorePartialValidity + ) + ); + assert_eq!(entry.scratch.alignment, entry.cub_layout.alignment); + if request.operation() == Operation::StorePartialValidity + && request.items_per_thread() > 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 && request.items_per_thread() == 4 { + assert_eq!(entry.cub_layout.size, 2064); + assert_eq!(entry.scratch.size, 2128); + } + } else { + assert_eq!(entry.scratch, entry.cub_layout); + } + } + } + Ok(()) +} + +#[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)?); + } + } + } + for algorithm in [Algorithm::Direct, Algorithm::Transpose] { + for operation in [ + Operation::LoadFull, + Operation::LoadFilled, + Operation::StoreFull, + Operation::StorePartial, + ] { + requests.push(Request::new_with_shape(operation, algorithm, [8, 4, 4], 4)?); + } + } + let start = Instant::now(); + let first = provider.compile(&requests, &target)?; + eprintln!( + "{}-specialization batch: {:?}, {} bytes", + requests.len(), + start.elapsed(), + first.ltoir.len() + ); + let second = provider.compile(&requests, &target)?; + assert_eq!(first.identity, second.identity); + assert_eq!(first.ltoir, second.ltoir, "repeated native output changed"); + assert_eq!(first.specializations.len(), requests.len()); + let mut identities = std::collections::HashSet::new(); + for ((entry, repeated), request) in first + .specializations + .iter() + .zip(&second.specializations) + .zip(&requests) + { + assert_eq!(entry.identity, repeated.identity); + assert_eq!(entry.symbol, repeated.symbol); + assert_eq!( + entry.mutates_payload, + request.algorithm() == Algorithm::Transpose + && matches!( + request.operation(), + Operation::StoreFull | Operation::StorePartial + ) + ); + assert!(identities.insert(entry.identity)); + eprintln!( + "{request:?}: CUB {:?}, scratch {:?}", + entry.cub_layout, entry.scratch + ); + if request.algorithm() == Algorithm::Direct { + assert!(entry.scratch_free); + assert_eq!( + entry.scratch, + Layout { + size: 0, + alignment: 1 + } + ); + assert_eq!( + entry.cub_layout, + Layout { + size: 1, + alignment: 1 + } + ); + } else { + assert!(!entry.scratch_free); + assert_eq!(entry.scratch, entry.cub_layout); + assert!(entry.scratch.size >= 128 * request.items_per_thread() as usize * 4); + } + } + // A specialization's symbol is independent of unrelated batch members. + let single = provider.compile(&requests[..1], &target)?; + assert_eq!( + single.specializations[0].identity, + first.specializations[0].identity + ); + assert_eq!( + single.specializations[0].symbol, + first.specializations[0].symbol + ); + assert_ne!(single.identity, first.identity); + assert!( + provider + .compile(&[], &target) + .unwrap_err() + .to_string() + .contains("at least one") + ); + assert!( + provider + .compile(&[requests[0]; 2], &target) + .unwrap_err() + .to_string() + .contains("duplicate") + ); + assert!( + provider + .compile(&requests[..1], "compute_120;injected") + .unwrap_err() + .to_string() + .contains("target must") + ); + assert!( + provider + .compile(&requests[..1], "compute_999999") + .unwrap_err() + .to_string() + .contains("nvrtcCompileProgram") + ); + Ok(()) +} diff --git a/rust/cuda-coop-codegen/tests/oxide-fixture/Cargo.toml b/rust/cuda-coop-codegen/tests/oxide-fixture/Cargo.toml new file mode 100644 index 000000000000..3278c198b444 --- /dev/null +++ b/rust/cuda-coop-codegen/tests/oxide-fixture/Cargo.toml @@ -0,0 +1,14 @@ +# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. +# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +[workspace] +members = ["api", "helper", "consumer", "carrier"] +default-members = ["consumer"] +resolver = "3" + +[workspace.dependencies] +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 new file mode 100644 index 000000000000..9adb3c9bfeca --- /dev/null +++ b/rust/cuda-coop-codegen/tests/oxide-fixture/README.md @@ -0,0 +1,147 @@ +# 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 (seven with `--root-context`). No consumer supplies an LTO artifact list. + +From the CCCL checkout, with the agreed Oxide source and backend already built: + +```sh +python3 rust/cuda-coop-codegen/tests/run_oxide_fixture.py \ + --oxide-root /path/to/cuda-rust \ + --cargo-oxide /path/to/cargo-oxide \ + --backend /path/to/librustc_codegen_cuda.so \ + --cuda-root /path/to/cuda \ + --arch sm_120 \ + --output /path/to/new-evidence-directory +``` + +The runner copies the fixture into the new output directory and patches only +that copy to the supplied source checkout. It seeds dependency resolutions from +that checkout's `vecadd` lockfile. The selected nightly must match Oxide's pin. +The log and command receipt are retained beside the copied consumer. + +The 480 cases cover i32 Direct/Transpose Load and Store, items 1/4, one/three +blocks, full/default-filled Load, full/partial Store, post-Store payload values, +and Store → explicit block synchronization → Load of the same region. Every +block uses uniform origins, count, default and mode; per-thread payload arrays +are initialized before use. There is no no-default partial-validity claim. + +The fixture checks compiler-derived scratch sizes and exact block geometry, +rejects wrong shapes, thread counts and explicit byte counts, and runs with +`auto_shared`. Its helpers deliberately retain function boundaries to +exercise discovery; register/stack results are not an inlining baseline. The +materialized path uses `cargo oxide run --materialize-cubin` and the ordinary +embedded module loader. Convergent declarations require the coordinated O03a +compiler change. + +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, +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 +(`#[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. 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. + +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. +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 \ + /path/to/carrier-evidence-directory +``` + +This command checks missing and conditional barriers, loop backedges, divergent +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. + +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 packed validity bits; 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 initialized slots 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. A focused i32/items9 case crosses a bitmap-byte boundary +and preserves scattered initialized slots outside the Store prefix. The binary +accepts `hole`, `count`, `bounds`, `ownership` and `different-count`; each must +fail in a separate process. The facade is private test infrastructure, not a +published safe collective API. + +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. 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. + +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/api/Cargo.toml b/rust/cuda-coop-codegen/tests/oxide-fixture/api/Cargo.toml new file mode 100644 index 000000000000..ba58cd8b591d --- /dev/null +++ b/rust/cuda-coop-codegen/tests/oxide-fixture/api/Cargo.toml @@ -0,0 +1,20 @@ +# 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 = "../../.." } + +[features] +checked = [] +carrier = [] +storage = [] +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 new file mode 100644 index 000000000000..6f0a40201f83 --- /dev/null +++ b/rust/cuda-coop-codegen/tests/oxide-fixture/api/build.rs @@ -0,0 +1,51 @@ +// Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +use cuda_coop_codegen::{Algorithm, Operation, oxide}; + +fn main() { + let argument = std::env::args().nth(1); + if matches!( + argument.as_deref(), + Some("--oxide-library-request-v1" | "--oxide-library-inputs-v1") + ) { + let checkout = std::path::Path::new(env!("CCCL_SOURCE")); + let result = if argument.as_deref() == Some("--oxide-library-inputs-v1") { + oxide::serve_inputs(checkout, std::io::stdin().lock(), std::io::stdout().lock()) + } else { + oxide::serve(checkout, std::io::stdin().lock(), std::io::stdout().lock()) + }; + if let Err(error) = result { + eprintln!("{error}"); + std::process::exit(1); + } + return; + } + println!("cargo::rerun-if-changed=build.rs"); + println!("cargo::rerun-if-env-changed=CCCL_SOURCE"); + let executable = std::env::current_exe().unwrap(); + for (algorithm, tag) in [ + (Algorithm::Direct, "DIRECT"), + (Algorithm::Transpose, "TRANSPOSE"), + ] { + for (operation, name) in [ + (Operation::LoadFull, "LOAD_FULL"), + (Operation::LoadFilled, "LOAD_FILLED"), + (Operation::StoreFull, "STORE_FULL"), + (Operation::StorePartial, "STORE_PARTIAL"), + ] { + let registration = oxide::registration(&executable, operation, algorithm).unwrap(); + println!("cargo::rustc-env=COOP_{tag}_{name}={registration}"); + } + } + 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 new file mode 100644 index 000000000000..d23d219467e5 --- /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(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; + assert!(ITEMS > 0 && threads <= usize::MAX / ITEMS); + assert!(valid <= threads * ITEMS && valid <= i32::MAX as usize); + assert!(offset <= len && valid <= len - offset); +} + +#[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; + 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 new file mode 100644 index 000000000000..c50c6be888df --- /dev/null +++ b/rust/cuda-coop-codegen/tests/oxide-fixture/api/src/lib.rs @@ -0,0 +1,236 @@ +// 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] + +#[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" + )); + + // Fault injection and flag observation belong only to this private fixture. + pub fn flag(values: &ThreadData, index: usize) -> u8 { + (values.initialized[index / 8] >> (index % 8)) & 1 + } + pub fn clear_slot(values: &mut ThreadData, index: usize) { + values.initialized[index / 8] &= !(1 << (index % 8)); + } +} +#[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 { + cuda_device::library::registration!($variable) + } + }; +} +registration!(direct_load_full, "COOP_DIRECT_LOAD_FULL"); +registration!(direct_load_filled, "COOP_DIRECT_LOAD_FILLED"); +registration!(direct_store_full, "COOP_DIRECT_STORE_FULL"); +registration!(direct_store_partial, "COOP_DIRECT_STORE_PARTIAL"); +registration!(transpose_load_full, "COOP_TRANSPOSE_LOAD_FULL"); +registration!(transpose_load_filled, "COOP_TRANSPOSE_LOAD_FILLED"); +registration!(transpose_store_full, "COOP_TRANSPOSE_STORE_FULL"); +registration!(transpose_store_partial, "COOP_TRANSPOSE_STORE_PARTIAL"); +#[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 +/// 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), + ); + } + } +} + +#[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}; + +#[cfg(feature = "storage")] +pub mod storage; diff --git a/rust/cuda-coop-codegen/tests/oxide-fixture/api/src/storage.rs b/rust/cuda-coop-codegen/tests/oxide-fixture/api/src/storage.rs new file mode 100644 index 000000000000..6b5478c358d5 --- /dev/null +++ b/rust/cuda-coop-codegen/tests/oxide-fixture/api/src/storage.rs @@ -0,0 +1,108 @@ +// Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +//! Private storage integration using the unchanged R01 numeric provider. +use cuda_device::{CollectiveTiles, DisjointSlice, library}; + +/// # Safety +/// All lanes agree on source, offset, valid count and storage origin. The source +/// contains the full accessed block range; the filled operation initializes all +/// lane payload slots. The root is exactly one 128-thread block per owned tile. +#[inline(never)] +pub unsafe fn load( + source: &[i32], + offset: usize, + valid: usize, + values: &mut [i32; ITEMS], + storage: Storage, +) { + assert!(valid <= 128 * ITEMS && offset <= source.len() && valid <= source.len() - offset); + unsafe { + if TRANSPOSE { + library::checked_block_storage_call_v1::( + super::transpose_load_filled, + source, + storage, + (source.as_ptr().add(offset), valid as i32, -7_i32), + (values.as_mut_ptr(),), + ); + } else { + library::checked_block_storage_call_v1::( + super::direct_load_filled, + source, + storage, + (source.as_ptr().add(offset), valid as i32, -7_i32), + (values.as_mut_ptr(),), + ); + } + } +} + +/// # Safety +/// All lanes agree on target, offset, valid count and storage origin. Offset is +/// this block's canonical 512-element tile start; payload values are initialized. +#[inline(never)] +pub unsafe fn store( + target: &mut DisjointSlice>, + offset: usize, + valid: usize, + values: &mut [i32; ITEMS], + storage: Storage, +) { + assert!(valid <= 128 * ITEMS && offset <= target.len() && valid <= target.len() - offset); + assert!(offset == cuda_device::thread::blockIdx_x() as usize * 512); + let pointer = target.as_mut_ptr(); + unsafe { + if TRANSPOSE { + library::checked_block_storage_call_v1::( + super::transpose_store_partial, + target, + storage, + (pointer.add(offset), valid as i32), + (values.as_mut_ptr(),), + ); + } else { + library::checked_block_storage_call_v1::( + super::direct_store_partial, + target, + storage, + (pointer.add(offset), valid as i32), + (values.as_mut_ptr(),), + ); + } + } +} + +/// # Safety +/// The carrier and range obey `store`'s ownership contract. A prior write must +/// be completed by an explicit or genuinely required planned block barrier. +#[inline(never)] +pub unsafe fn reload( + source: &mut DisjointSlice>, + offset: usize, + valid: usize, + values: &mut [i32; ITEMS], + storage: Storage, +) { + assert!(valid <= 128 * ITEMS && offset <= source.len() && valid <= source.len() - offset); + let pointer = source.as_mut_ptr(); + unsafe { + if TRANSPOSE { + library::checked_block_storage_call_v1::( + super::transpose_load_filled, + &*source, + storage, + (pointer.add(offset).cast_const(), valid as i32, -7_i32), + (values.as_mut_ptr(),), + ); + } else { + library::checked_block_storage_call_v1::( + super::direct_load_filled, + &*source, + storage, + (pointer.add(offset).cast_const(), valid as i32, -7_i32), + (values.as_mut_ptr(),), + ); + } + } +} diff --git a/rust/cuda-coop-codegen/tests/oxide-fixture/api/src/validity.rs b/rust/cuda-coop-codegen/tests/oxide-fixture/api/src/validity.rs new file mode 100644 index 000000000000..3118fb4bbaf6 --- /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: &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_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 { + 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: &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_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 { + 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 new file mode 100644 index 000000000000..0f96fdca727d --- /dev/null +++ b/rust/cuda-coop-codegen/tests/oxide-fixture/carrier/Cargo.toml @@ -0,0 +1,24 @@ +# 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 +default-run = "coop-carrier-fixture" + +[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"] } + +[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/main.rs b/rust/cuda-coop-codegen/tests/oxide-fixture/carrier/src/main.rs new file mode 100644 index 000000000000..394c0da136c9 --- /dev/null +++ b/rust/cuda-coop-codegen/tests/oxide-fixture/carrier/src/main.rs @@ -0,0 +1,200 @@ +// 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 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); + 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. + } + + #[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, LaunchConfig2D}; + + 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])?; + 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] { + 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 { + 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); + + 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(()) +} diff --git a/rust/cuda-coop-codegen/tests/oxide-fixture/carrier/src/validity.rs b/rust/cuda-coop-codegen/tests/oxide-fixture/carrier/src/validity.rs new file mode 100644 index 000000000000..47d9f25b345d --- /dev/null +++ b/rust/cuda-coop-codegen/tests/oxide-fixture/carrier/src/validity.rs @@ -0,0 +1,556 @@ +// 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, + 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 * CAPACITY + 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 and selected slots outside it are initialized. + let mut other = ThreadData::<_, ITEMS_PER_THREAD>::new(); + let mut item = 0; + while item < ITEMS_PER_THREAD { + if lane * ITEMS_PER_THREAD + item < valid + || (mode == 7 && (lane + item) % 3 == 1) + { + values.write(item, default); + other.write(item, default); + } + item += 1; + } + values = if lane % 2 == 0 { values } else { other }; + } + } + if flag != 1 { + thread_data::clear_slot(&mut values, 0); + } + 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. + } + + // Cross a bitmap-byte boundary without expanding the ten-type matrix. + #[kernel] + #[launch_contract(domain = 1, block = (128, 1, 1), dynamic_shared = inferred)] + pub fn copy_i32_boundary_128( + input: &[i32], + mut output: DisjointSlice>, + valid: usize, + shift: usize, + default: i32, + mode: u32, + flag: u8, + ) { + exercise::<_, ITEMS_PER_THREAD, 16>(input, &mut output, valid, shift, default, mode, flag); + } + + #[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, + flag: u8, + ) { + exercise::<_, ITEMS_PER_THREAD, 4>(input, &mut output, valid, shift, default, mode, 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, + flag: u8, + ) { + exercise::<_, ITEMS_PER_THREAD, 4>(input, &mut output, valid, shift, default, mode, 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, + flag: u8, + ) { + exercise::<_, ITEMS_PER_THREAD, 4>(input, &mut output, valid, shift, default, mode, 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, + flag: u8, + ) { + exercise::<_, ITEMS_PER_THREAD, 4>(input, &mut output, valid, shift, default, mode, 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, + flag: u8, + ) { + exercise::<_, ITEMS_PER_THREAD, 4>(input, &mut output, valid, shift, default, mode, 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, + flag: u8, + ) { + exercise::<_, ITEMS_PER_THREAD, 4>(input, &mut output, valid, shift, default, mode, 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, + flag: u8, + ) { + exercise::<_, ITEMS_PER_THREAD, 4>(input, &mut output, valid, shift, default, mode, 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, + flag: u8, + ) { + exercise::<_, ITEMS_PER_THREAD, 4>(input, &mut output, valid, shift, default, mode, 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, + flag: u8, + ) { + exercise::<_, ITEMS_PER_THREAD, 4>(input, &mut output, valid, shift, default, mode, 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, + flag: u8, + ) { + exercise::<_, ITEMS_PER_THREAD, 4>(input, &mut output, valid, shift, default, mode, 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, + flag: u8, + ) { + exercise::<_, ITEMS_PER_THREAD, 4>(input, &mut output, valid, shift, default, mode, 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) => { + run!($t, $prepare, $launch, $block, $items, $default, 4) + }; + ($t:ty, $prepare:ident, $launch:ident, $block:literal, $items:literal, $default:expr, $capacity:literal) => {{ + 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, 7] { + if mode == 7 && $items != 9 { + continue; + } + 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 * $capacity + 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 * $capacity); + 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 == 7 || (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() { + "hole" => (512, 0, 0, 0), + "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 + ); + run!( + i32, + prepare_copy_i32_boundary_128, + copy_i32_boundary_128_async_owned, + 128, + 9, + -17, + 16 + ); + println!( + "PASS: {cases} Direct validity cases, ten numeric types, exact flags, short tails, owned async" + ); + Ok(()) +} diff --git a/rust/cuda-coop-codegen/tests/oxide-fixture/compile_cases/atomic_control.rs b/rust/cuda-coop-codegen/tests/oxide-fixture/compile_cases/atomic_control.rs new file mode 100644 index 000000000000..9646c1c74a30 --- /dev/null +++ b/rust/cuda-coop-codegen/tests/oxide-fixture/compile_cases/atomic_control.rs @@ -0,0 +1,34 @@ +// Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +use coop_marker_fixture_helper::direct_load_filled; +use cuda_device::atomic::{AtomicOrdering, DeviceAtomicU32}; +use cuda_device::{kernel, launch_contract}; +use cuda_host::cuda_module; + +#[cuda_module] +mod kernels { + use super::*; + + #[kernel] + #[launch_contract(block = (128, 1, 1), domain = 3, dynamic_shared = inferred)] + pub unsafe fn probe(input: *const i32, counter: *const DeviceAtomicU32, flag: bool) { + // A panic-only host placeholder still returns a varying device value. + // Merging it with a constant branch must not prove uniformity. + let count = if flag { + unsafe { (*counter).fetch_add(1, AtomicOrdering::Relaxed) as i32 } + } else { + 0 + }; + let mut value = 0_i32; + unsafe { + cuda_device::library::checked_block_call_v1::( + direct_load_filled, + (input, count, -7), + (&mut value as *mut i32,), + ); + } + } +} + +fn main() {} diff --git a/rust/cuda-coop-codegen/tests/oxide-fixture/compile_cases/native_mutation.rs b/rust/cuda-coop-codegen/tests/oxide-fixture/compile_cases/native_mutation.rs new file mode 100644 index 000000000000..a0e57f30b2b4 --- /dev/null +++ b/rust/cuda-coop-codegen/tests/oxide-fixture/compile_cases/native_mutation.rs @@ -0,0 +1,43 @@ +// Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception +use coop_marker_fixture_helper::{direct_load_filled, direct_load_full}; +use cuda_device::{kernel, launch_contract}; +use cuda_host::cuda_module; +#[inline(never)] +fn twice( + value: &&i32, + _a: i32, + _b: i32, + _c: i32, + _d: i32, + _e: i32, + _f: i32, + _g: i32, + _h: i32, +) -> i32 { + **value +} +#[cuda_module] +mod kernels { + use super::*; + #[kernel] + #[launch_contract(block = (128,1,1), domain = 3, dynamic_shared = inferred)] + pub unsafe fn probe(input: *const i32) { + let mut count = 1_i32; + let mut payload = 0_i32; + unsafe { + cuda_device::library::checked_block_call_v1::( + direct_load_full, + (input,), + (&mut count as *mut i32,), + ); + cuda_device::thread::sync_threads(); + cuda_device::library::checked_block_call_v1::( + direct_load_filled, + (input, count, -7_i32), + (&mut payload as *mut i32,), + ); + } + } +} +fn main() {} diff --git a/rust/cuda-coop-codegen/tests/oxide-fixture/compile_cases/nested_reference.rs b/rust/cuda-coop-codegen/tests/oxide-fixture/compile_cases/nested_reference.rs new file mode 100644 index 000000000000..c2df4f7d7382 --- /dev/null +++ b/rust/cuda-coop-codegen/tests/oxide-fixture/compile_cases/nested_reference.rs @@ -0,0 +1,39 @@ +// Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception +use coop_marker_fixture_helper::{direct_load_filled, direct_load_full}; +use cuda_device::{kernel, launch_contract}; +use cuda_host::cuda_module; +#[inline(never)] +fn twice( + value: &&i32, + _a: i32, + _b: i32, + _c: i32, + _d: i32, + _e: i32, + _f: i32, + _g: i32, + _h: i32, +) -> i32 { + **value +} +#[cuda_module] +mod kernels { + use super::*; + #[kernel] + #[launch_contract(block = (128,1,1), domain = 3, dynamic_shared = inferred)] + pub unsafe fn probe(input: *const i32) { + let value = cuda_device::thread::threadIdx_x() as i32; + let reference = &value; + let count = twice(&reference, 0, 0, 0, 0, 0, 0, 0, 0); + let mut payload = 0_i32; + unsafe { + cuda_device::library::checked_block_call_v1::( + direct_load_filled, + (input, count, -7_i32), + (&mut payload as *mut i32,), + ); + } + } +} +fn main() {} diff --git a/rust/cuda-coop-codegen/tests/oxide-fixture/consumer/Cargo.toml b/rust/cuda-coop-codegen/tests/oxide-fixture/consumer/Cargo.toml new file mode 100644 index 000000000000..1d496140d001 --- /dev/null +++ b/rust/cuda-coop-codegen/tests/oxide-fixture/consumer/Cargo.toml @@ -0,0 +1,21 @@ +# 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 +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 new file mode 100644 index 000000000000..3f2b53b5f283 --- /dev/null +++ b/rust/cuda-coop-codegen/tests/oxide-fixture/consumer/src/main.rs @@ -0,0 +1,353 @@ +// 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::*; + #[kernel] + #[launch_bounds(256)] + #[launch_contract(block = (128, 1, 1), domain = 3, dynamic_shared = inferred)] + pub unsafe fn direct_1( + input: *const i32, + output: *mut i32, + post: *mut i32, + valid: i32, + mode: u32, + ) { + unsafe { + coop_marker_fixture_helper::roundtrip::<1, false>(input, output, post, valid, mode); + } + } + #[kernel] + #[launch_bounds(256)] + #[launch_contract(block = (128, 1, 1), domain = 3, dynamic_shared = inferred)] + pub unsafe fn direct_4( + input: *const i32, + output: *mut i32, + post: *mut i32, + valid: i32, + mode: u32, + ) { + unsafe { + coop_marker_fixture_helper::roundtrip::<4, false>(input, output, post, valid, mode); + } + } + #[kernel] + #[launch_bounds(256)] + #[launch_contract(block = (128, 1, 1), domain = 3, dynamic_shared = inferred)] + pub unsafe fn transpose_1( + input: *const i32, + output: *mut i32, + post: *mut i32, + valid: i32, + mode: u32, + ) { + unsafe { + coop_marker_fixture_helper::roundtrip::<1, true>(input, output, post, valid, mode); + } + } + #[kernel] + #[launch_bounds(256)] + #[launch_contract(block = (128, 1, 1), domain = 3, dynamic_shared = inferred)] + pub unsafe fn transpose_4( + input: *const i32, + output: *mut i32, + post: *mut i32, + valid: i32, + mode: u32, + ) { + unsafe { + coop_marker_fixture_helper::roundtrip::<4, true>(input, output, post, valid, mode); + } + } + #[kernel] + #[launch_bounds(256)] + #[launch_contract(block = (128, 1, 1), domain = 3, dynamic_shared = inferred)] + pub unsafe fn shared_helper_root( + input: *const i32, + output: *mut i32, + post: *mut i32, + valid: i32, + mode: u32, + ) { + unsafe { + coop_marker_fixture_helper::shared_roundtrip(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::shared_roundtrip(input, output, post, valid, mode); + } + } +} + +#[cfg(not(feature = "storage"))] +fn main() -> Result<(), Box> { + use cuda_core::{CudaContext, DeviceBuffer}; + use cuda_core::{LaunchConfig3D, LaunchContractError}; + use oxide_artifacts::{ArtifactPayloadKind, read_artifact_bundles_from_object_bytes}; + let executable = std::fs::read(std::env::current_exe()?)?; + let bundles = read_artifact_bundles_from_object_bytes(&executable)?; + let bundle = bundles + .iter() + .find(|b| { + b.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. + let module = unsafe { kernels::load(&context)? }; + macro_rules! check_plan { + ($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(($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 [ + ($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 { .. }) + )); + } + 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), ($threads, 1, 1), bytes)), + Err(LaunchContractError::DynamicSharedMemoryExactMismatch { .. }) + )); + } + } + }}; + } + 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, 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 = threads * 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 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 = + 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; + unsafe { + match kernel { + 0 => module.direct_1( + &stream, + &module.prepare_direct_1(config)?, + input_ptr, + output_ptr, + post_ptr, + valid as i32, + mode, + )?, + 1 => module.direct_4( + &stream, + &module.prepare_direct_4(config)?, + input_ptr, + output_ptr, + post_ptr, + valid as i32, + mode, + )?, + 2 => module.transpose_1( + &stream, + &module.prepare_transpose_1(config)?, + input_ptr, + output_ptr, + post_ptr, + valid as i32, + mode, + )?, + 3 => module.transpose_4( + &stream, + &module.prepare_transpose_4(config)?, + input_ptr, + output_ptr, + post_ptr, + valid as i32, + mode, + )?, + 4 => module.shared_helper_root( + &stream, + &module.prepare_shared_helper_root(config)?, + input_ptr, + output_ptr, + post_ptr, + 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 + .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 * threads + 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, + if cfg!(feature = "root-context") { + 672 + } else { + 480 + } + ); + println!( + "PASS {cases} native marker GPU cases with inferred scratch through shared generic helpers; {} CUBIN bytes", + cubin.len() + ); + Ok(()) +} + +#[cfg(feature = "storage")] +mod storage; +#[cfg(feature = "storage")] +fn main() -> Result<(), Box> { + storage::run() +} diff --git a/rust/cuda-coop-codegen/tests/oxide-fixture/consumer/src/storage.rs b/rust/cuda-coop-codegen/tests/oxide-fixture/consumer/src/storage.rs new file mode 100644 index 000000000000..c1ab62f93231 --- /dev/null +++ b/rust/cuda-coop-codegen/tests/oxide-fixture/consumer/src/storage.rs @@ -0,0 +1,176 @@ +// 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 + }; + 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, + actual - 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 if $mode == 8 { + (1 + (valid & 3)) 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); + 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/Cargo.toml b/rust/cuda-coop-codegen/tests/oxide-fixture/helper/Cargo.toml new file mode 100644 index 000000000000..bce66cd3101c --- /dev/null +++ b/rust/cuda-coop-codegen/tests/oxide-fixture/helper/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 = "coop-marker-fixture-helper" +version = "0.0.0" +edition = "2024" +publish = false + +[dependencies] +cuda-device.workspace = true +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/checked.rs b/rust/cuda-coop-codegen/tests/oxide-fixture/helper/src/checked.rs new file mode 100644 index 000000000000..9dcf388dc50d --- /dev/null +++ b/rust/cuda-coop-codegen/tests/oxide-fixture/helper/src/checked.rs @@ -0,0 +1,108 @@ +// 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 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 +/// fill value is uniform, and each thread owns its separate initialized payload. +#[cuda_device::device] +pub unsafe fn roundtrip( + input: *const i32, + output: *mut i32, + post_store: *mut i32, + 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 * block_threads * ITEMS_PER_THREAD; + let valid = (valid - origin as i32) + .max(0) + .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 + } 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, + ); + if mode & 16 != 0 { + payload[trap_index] = 0; + } + cuda_device::thread::sync_threads(); + iteration += 1; + } + let mut i = 0; + while i < ITEMS_PER_THREAD { + payload[i] += CONTEXT_INCREMENT; + 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 new file mode 100644 index 000000000000..77f9296ed30d --- /dev/null +++ b/rust/cuda-coop-codegen/tests/oxide-fixture/helper/src/lib.rs @@ -0,0 +1,89 @@ +// Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +#![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 +/// 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. +#[cfg(not(feature = "checked"))] +#[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); + } + } + } +} + +#[cfg(feature = "checked")] +mod checked; +#[cfg(feature = "checked")] +pub use checked::roundtrip; + +#[cfg(feature = "storage")] +pub mod storage; diff --git a/rust/cuda-coop-codegen/tests/oxide-fixture/helper/src/storage.rs b/rust/cuda-coop-codegen/tests/oxide-fixture/helper/src/storage.rs new file mode 100644 index 000000000000..b51cad4c48e5 --- /dev/null +++ b/rust/cuda-coop-codegen/tests/oxide-fixture/helper/src/storage.rs @@ -0,0 +1,143 @@ +// 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); + } + 7 => { + 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); + } + 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); + } + } + } +} diff --git a/rust/cuda-coop-codegen/tests/run_oxide_fixture.py b/rust/cuda-coop-codegen/tests/run_oxide_fixture.py new file mode 100644 index 000000000000..78825e62869f --- /dev/null +++ b/rust/cuda-coop-codegen/tests/run_oxide_fixture.py @@ -0,0 +1,116 @@ +# 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") +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", +) +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", +) +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] +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", "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. +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", + "build" if args.compile_only else "run", + "--arch", + args.arch, + "--materialize-cubin", +] +if args.validity: + command.extend(["--features", "validity", "--bin", "coop-validity-fixture"]) +elif args.storage: + command.extend(["--features", "storage"]) +elif args.root_context: + command.extend(["--features", "root-context"]) +elif args.checked: + command.extend(["--features", "checked"]) +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 + ) +(output / "command.json").write_text( + json.dumps( + { + "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, + }, + indent=2, + ) + + "\n" +) +raise SystemExit(result.returncode) diff --git a/rust/cuda-coop/Cargo.toml b/rust/cuda-coop/Cargo.toml new file mode 100644 index 000000000000..0f3c80862caf --- /dev/null +++ b/rust/cuda-coop/Cargo.toml @@ -0,0 +1,18 @@ +# 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 +publish = false + +[dependencies] +cuda-device = { git = "https://github.com/tpn/cuda-rust.git", rev = "204e4e6f721bd111bf5e0e00caa69dc672a4d378" } + +[build-dependencies] +cuda-coop-codegen = { path = "../cuda-coop-codegen" } + +[lints] +workspace = true diff --git a/rust/cuda-coop/build.rs b/rust/cuda-coop/build.rs new file mode 100644 index 000000000000..64adc8bb40e8 --- /dev/null +++ b/rust/cuda-coop/build.rs @@ -0,0 +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 block Load and Store. +use cuda_coop_codegen::{Algorithm, Operation, oxide}; + +fn main() { + 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/.gitignore b/rust/cuda-coop/examples/load_store/.gitignore new file mode 100644 index 000000000000..74bffa8af4a4 --- /dev/null +++ b/rust/cuda-coop/examples/load_store/.gitignore @@ -0,0 +1,7 @@ +# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. +# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception +/target/ +/.oxide-artifacts/ +/*.ll +/*.options +/*.target diff --git a/rust/cuda-coop/examples/load_store/Cargo.lock b/rust/cuda-coop/examples/load_store/Cargo.lock new file mode 100644 index 000000000000..161a9f9544e0 --- /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=204e4e6f721bd111bf5e0e00caa69dc672a4d378#204e4e6f721bd111bf5e0e00caa69dc672a4d378" +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=204e4e6f721bd111bf5e0e00caa69dc672a4d378#204e4e6f721bd111bf5e0e00caa69dc672a4d378" +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=204e4e6f721bd111bf5e0e00caa69dc672a4d378#204e4e6f721bd111bf5e0e00caa69dc672a4d378" +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=204e4e6f721bd111bf5e0e00caa69dc672a4d378#204e4e6f721bd111bf5e0e00caa69dc672a4d378" +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=204e4e6f721bd111bf5e0e00caa69dc672a4d378#204e4e6f721bd111bf5e0e00caa69dc672a4d378" +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=204e4e6f721bd111bf5e0e00caa69dc672a4d378#204e4e6f721bd111bf5e0e00caa69dc672a4d378" +dependencies = [ + "cuda-macros", +] + +[[package]] +name = "cuda-host" +version = "0.2.1" +source = "git+https://github.com/tpn/cuda-rust.git?rev=204e4e6f721bd111bf5e0e00caa69dc672a4d378#204e4e6f721bd111bf5e0e00caa69dc672a4d378" +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=204e4e6f721bd111bf5e0e00caa69dc672a4d378#204e4e6f721bd111bf5e0e00caa69dc672a4d378" +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=204e4e6f721bd111bf5e0e00caa69dc672a4d378#204e4e6f721bd111bf5e0e00caa69dc672a4d378" + +[[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=204e4e6f721bd111bf5e0e00caa69dc672a4d378#204e4e6f721bd111bf5e0e00caa69dc672a4d378" +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=204e4e6f721bd111bf5e0e00caa69dc672a4d378#204e4e6f721bd111bf5e0e00caa69dc672a4d378" +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=204e4e6f721bd111bf5e0e00caa69dc672a4d378#204e4e6f721bd111bf5e0e00caa69dc672a4d378" +dependencies = [ + "object", +] + +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +source = "git+https://github.com/tpn/cuda-rust.git?rev=204e4e6f721bd111bf5e0e00caa69dc672a4d378#204e4e6f721bd111bf5e0e00caa69dc672a4d378" +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=204e4e6f721bd111bf5e0e00caa69dc672a4d378#204e4e6f721bd111bf5e0e00caa69dc672a4d378" + +[[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=204e4e6f721bd111bf5e0e00caa69dc672a4d378#204e4e6f721bd111bf5e0e00caa69dc672a4d378" + +[[package]] +name = "rustc-hash" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "08d43f7aa6b08d49f382cde6a7982047c3426db949b1424bc4b7ec9ae12c6ce2" + +[[package]] +name = "rustc-hash" +version = "2.1.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6b1e7f9a428571be2dc5bc0505c13fb6bf936822b894ec87abf8a08a4e51742d" + +[[package]] +name = "rustix" +version = "0.38.44" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fdb5bc1ae2baa591800df16c9ca78619bf65c0488b41b96ccec5d11220d8c154" +dependencies = [ + "bitflags", + "errno", + "libc", + "linux-raw-sys", + "windows-sys 0.59.0", +] + +[[package]] +name = "scopeguard" +version = "1.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "94143f37725109f92c262ed2cf5e59bce7498c01bcc1502d7b9afe439a4e9f49" + +[[package]] +name = "serde" +version = "1.0.229" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4148590afebada386688f18773da617792bf2ef03ffc1e4cbd2b1d45b023e0ba" +dependencies = [ + "serde_core", + "serde_derive", +] + +[[package]] +name = "serde_core" +version = "1.0.229" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "67dca2c9c51e58a4791a4b1ed58308b39c64224d349a935ab5039aa360942a48" +dependencies = [ + "serde_derive", +] + +[[package]] +name = "serde_derive" +version = "1.0.229" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e7a5d71263a5a7d47b41f6b3f06ba276f10cc18b0931f1799f710578e2309348" +dependencies = [ + "proc-macro2", + "quote", + "syn 3.0.9", +] + +[[package]] +name = "serde_json" +version = "1.0.154" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e7e9cc8b1b85264074fbcc02a88680c4096b1e47df8f739dceb03bf482f04bd6" +dependencies = [ + "itoa", + "memchr", + "serde", + "serde_core", + "zmij", +] + +[[package]] +name = "sha2" +version = "0.10.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a7507d819769d01a365ab707794a4084392c824f54a7a6a7862f8c3d0892b283" +dependencies = [ + "cfg-if", + "cpufeatures", + "digest", +] + +[[package]] +name = "shlex" +version = "1.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0fda2ff0d084019ba4d7c6f371c95d8fd75ce3524c3cb8fb653a3023f6323e64" + +[[package]] +name = "slab" +version = "0.4.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0c790de23124f9ab44544d7ac05d60440adc586479ce501c1d6d7da3cd8c9cf5" + +[[package]] +name = "smallvec" +version = "1.16.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5b3dc8af474f516a851ff4bd12db780f948b9250ad37211e4eec0bccea54e01b" + +[[package]] +name = "syn" +version = "2.0.119" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "872831b642d1a07999a962a351ed35b955ea2cfc8f3862091e2a240a84f17297" +dependencies = [ + "proc-macro2", + "quote", + "unicode-ident", +] + +[[package]] +name = "syn" +version = "3.0.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d78c8dee4c7bf0e14673097256fed6142ce9d3b85a408189d07482442145823b" +dependencies = [ + "proc-macro2", + "quote", + "unicode-ident", +] + +[[package]] +name = "thiserror" +version = "1.0.69" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b6aaf5339b578ea85b50e080feb250a3e8ae8cfcdff9a461c9ec2904bc923f52" +dependencies = [ + "thiserror-impl 1.0.69", +] + +[[package]] +name = "thiserror" +version = "2.0.21" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "09e52cb86a36cede5cb101bf8908837b3e4c6e5e59fe7fd85c23fb56200d189e" +dependencies = [ + "thiserror-impl 2.0.21", +] + +[[package]] +name = "thiserror-impl" +version = "1.0.69" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4fee6c4efc90059e10f81e6d42c60a18f76588c3d74cb83a0b242a2b6c7504c1" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", +] + +[[package]] +name = "thiserror-impl" +version = "2.0.21" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fe5197923287db20a58125f0bc85c062f7f2c892de97b18c356f9efb14b28524" +dependencies = [ + "proc-macro2", + "quote", + "syn 3.0.9", +] + +[[package]] +name = "typenum" +version = "1.20.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b6f5e870be6c3b371b77fe0ee0bafb859fa4964b4404c27de1d380043c4dda20" + +[[package]] +name = "unicode-ident" +version = "1.0.27" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a2c754d6c33795a1c324727428e5a7dedb5b06195f9890bdbcba760d3e246563" + +[[package]] +name = "version_check" +version = "0.9.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0b928f33d975fc6ad9f86c8f283853ad26bdd5b10b7f1542aa2fa15e2289105a" + +[[package]] +name = "which" +version = "4.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "87ba24419a2078cd2b0f2ede2691b6c66d8e47836da3b6db8265ebad47afbfc7" +dependencies = [ + "either", + "home", + "once_cell", + "rustix", +] + +[[package]] +name = "windows-link" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f0805222e57f7521d6a62e36fa9163bc891acd422f971defe97d64e70d0a4fe5" + +[[package]] +name = "windows-sys" +version = "0.59.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1e38bc4d79ed67fd075bcc251a1c39b32a1776bbe92e5bef1f0bf1f8c531853b" +dependencies = [ + "windows-targets", +] + +[[package]] +name = "windows-sys" +version = "0.61.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ae137229bcbd6cdf0f7b80a31df61766145077ddf49416a728b02cb3921ff3fc" +dependencies = [ + "windows-link", +] + +[[package]] +name = "windows-targets" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9b724f72796e036ab90c1021d4780d4d3d648aca59e491e6b98e725b84e99973" +dependencies = [ + "windows_aarch64_gnullvm", + "windows_aarch64_msvc", + "windows_i686_gnu", + "windows_i686_gnullvm", + "windows_i686_msvc", + "windows_x86_64_gnu", + "windows_x86_64_gnullvm", + "windows_x86_64_msvc", +] + +[[package]] +name = "windows_aarch64_gnullvm" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "32a4622180e7a0ec044bb555404c800bc9fd9ec262ec147edd5989ccd0c02cd3" + +[[package]] +name = "windows_aarch64_msvc" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "09ec2a7bb152e2252b53fa7803150007879548bc709c039df7627cabbd05d469" + +[[package]] +name = "windows_i686_gnu" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8e9b5ad5ab802e97eb8e295ac6720e509ee4c243f69d781394014ebfe8bbfa0b" + +[[package]] +name = "windows_i686_gnullvm" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0eee52d38c090b3caa76c563b86c3a4bd71ef1a819287c19d586d7334ae8ed66" + +[[package]] +name = "windows_i686_msvc" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "240948bc05c5e7c6dabba28bf89d89ffce3e303022809e73deaefe4f6ec56c66" + +[[package]] +name = "windows_x86_64_gnu" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "147a5c80aabfbf0c7d901cb5895d1de30ef2907eb21fbbab29ca94c5b08b1a78" + +[[package]] +name = "windows_x86_64_gnullvm" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "24d5b23dc417412679681396f2b49f3de8c1473deb516bd34410872eff51ed0d" + +[[package]] +name = "windows_x86_64_msvc" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "589f6da84c646204747d1270a2a5661ea66ed1cced2631d546fdfb155959f9ec" + +[[package]] +name = "zerocopy" +version = "0.8.63" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e5fe1f8f1b06191a00962174c61aa5005e0bb391a6d80d07e24d115c01a92ed8" +dependencies = [ + "zerocopy-derive", +] + +[[package]] +name = "zerocopy-derive" +version = "0.8.63" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "863ad3ac83293fb4d740aedbfdc9240dd8d1a50c1099acd76ce80ce7c7230c7f" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", +] + +[[package]] +name = "zmij" +version = "1.0.23" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "29666d0abbfad1e3dc4dcf6144730dd3a3ab225bbbdac83319345b1b44ccfc1b" diff --git a/rust/cuda-coop/examples/load_store/Cargo.toml b/rust/cuda-coop/examples/load_store/Cargo.toml new file mode 100644 index 000000000000..258c8f4142bd --- /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 = "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/examples/load_store/helper/Cargo.toml b/rust/cuda-coop/examples/load_store/helper/Cargo.toml new file mode 100644 index 000000000000..1c6b59094851 --- /dev/null +++ b/rust/cuda-coop/examples/load_store/helper/Cargo.toml @@ -0,0 +1,11 @@ +# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. +# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +[package] +name = "coop-load-store-helper" +version = "0.0.0" +edition = "2024" +publish = false + +[dependencies] +cuda-coop = { path = "../../.." } diff --git a/rust/cuda-coop/examples/load_store/helper/src/lib.rs b/rust/cuda-coop/examples/load_store/helper/src/lib.rs new file mode 100644 index 000000000000..b587972831d8 --- /dev/null +++ b/rust/cuda-coop/examples/load_store/helper/src/lib.rs @@ -0,0 +1,58 @@ +// 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, +) { + 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( + 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 new file mode 100644 index 000000000000..0f41cc7fdfa7 --- /dev/null +++ b/rust/cuda-coop/examples/load_store/src/main.rs @@ -0,0 +1,96 @@ +// 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, + ); + } + #[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> { + 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, $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::<$items>(LaunchConfig1D::auto_shared(3, 128))?; + let (_, target) = module + .$launch::<$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, 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/examples/payload.rs b/rust/cuda-coop/examples/payload.rs new file mode 100644 index 000000000000..a7625da6c1ba --- /dev/null +++ b/rust/cuda-coop/examples/payload.rs @@ -0,0 +1,22 @@ +// Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +//! Construct and inspect a per-thread payload on the host. + +use cuda_coop as coop; + +fn payload(value: i16) -> coop::ThreadData { + let mut items = coop::ThreadData::<_, ITEMS_PER_THREAD>::new(); + items.write(0, value); + items +} + +fn main() { + let mut items = payload::<4>(7); + assert_eq!(items.get(0), Some(7_i16)); + assert_eq!(items.get(1), None); + items[0] += 2; + items.write(3, 15); + assert_eq!(items.get(0), Some(9)); + assert_eq!(items.get(3), Some(15)); +} diff --git a/rust/cuda-coop/src/alignment.rs b/rust/cuda-coop/src/alignment.rs new file mode 100644 index 000000000000..63e8f2987607 --- /dev/null +++ b/rust/cuda-coop/src/alignment.rs @@ -0,0 +1,65 @@ +// Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +/// Type-level byte alignment for [`crate::ThreadData`]. +/// +/// `BYTES` must be a power of two from 1 through 536,870,912 (2²⁹), inclusive. +/// These are Rust's supported explicit alignments. The payload's actual +/// alignment is at least this value and the element type's natural alignment. +/// Larger alignments can increase the payload's size through padding. +pub struct Alignment; + +/// A supported [`Alignment`] selector. +/// +/// This trait is sealed: implementations are supplied for all powers of two +/// from 1 through 2²⁹. Use this public bound when the alignment is generic: +/// +/// ``` +/// use cuda_coop as coop; +/// +/// fn payload(value: u32) +/// -> coop::ThreadData +/// where +/// coop::Alignment: coop::SupportedAlignment, +/// { +/// let mut items = coop::ThreadData::<_, ITEMS_PER_THREAD, ALIGN>::new(); +/// items.write(0, value); +/// items +/// } +/// +/// assert_eq!(payload::<4, 32>(11).get(0), Some(11)); +/// ``` +pub trait SupportedAlignment: sealed::Alignment {} + +pub(crate) mod sealed { + pub trait Alignment { + type Marker: Copy; + } + + macro_rules! alignments { + ($($marker:ident: $bytes:literal),* $(,)?) => { + $( + #[repr(align($bytes))] + #[derive(Clone, Copy)] + pub struct $marker; + + impl Alignment for super::Alignment<$bytes> { + type Marker = $marker; + } + + impl super::SupportedAlignment for super::Alignment<$bytes> {} + )* + }; + } + + alignments! { + A1: 1, A2: 2, A4: 4, A8: 8, A16: 16, A32: 32, + A64: 64, A128: 128, A256: 256, A512: 512, A1024: 1024, + A2048: 2048, A4096: 4096, A8192: 8192, A16384: 16384, + A32768: 32768, A65536: 65536, A131072: 131072, A262144: 262144, + A524288: 524288, A1048576: 1048576, A2097152: 2097152, + A4194304: 4194304, A8388608: 8388608, A16777216: 16777216, + A33554432: 33554432, A67108864: 67108864, A134217728: 134217728, + A268435456: 268435456, A536870912: 536870912, + } +} diff --git a/rust/cuda-coop/src/group.rs b/rust/cuda-coop/src/group.rs new file mode 100644 index 000000000000..97004caff679 --- /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 a thread excluded from a mapped group. + /// + /// ```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,E0277 +/// let group = cuda_coop::this_warp().group_by::<8>(); +/// group.rank_at::(); +/// ``` +/// ```compile_fail,E0599 +/// 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,E0277 +/// 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,E0080 + /// cuda_coop::this_warp().group_by::<0>(); + /// ``` + /// ```compile_fail,E0080 + /// 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,E0080 + /// 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 new file mode 100644 index 000000000000..c7536f88cf7e --- /dev/null +++ b/rust/cuda-coop/src/lib.rs @@ -0,0 +1,56 @@ +// Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +//! 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. +//! [`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); +//! ``` +//! +//! 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. [`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] + +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::{ + Algorithm, Direct, LoadOptions, LoadSource, StoreOptions, StoreValue, Transpose, load, store, +}; +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/src/load_store.rs b/rust/cuda-coop/src/load_store.rs new file mode 100644 index 000000000000..5960d27edbe0 --- /dev/null +++ b/rust/cuda-coop/src/load_store.rs @@ -0,0 +1,404 @@ +// 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 core::marker::PhantomData; +use cuda_device::{CollectiveTiles, DisjointSlice}; + +mod boundary; + +/// 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. +/// +/// ```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; + +/// 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 +/// 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 { + algorithm: PhantomData, + valid_items: Option, + oob_default: Option, + offset: usize, +} + +impl Default for LoadOptions { + fn default() -> Self { + Self { + algorithm: PhantomData, + valid_items: None, + oob_default: None, + offset: 0, + } + } +} + +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. + 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 [`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)] +pub struct StoreOptions { + algorithm: PhantomData, + valid_items: Option, + offset: usize, +} + +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 + /// 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. [`Direct`] uses 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. +/// 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, A: Algorithm>( + group: ThreadGroup, + source: S, + output: &mut ThreadData, + options: LoadOptions, +) where + Alignment: SupportedAlignment, +{ + source.load(group, output, options); +} + +/// 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`] 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. 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 +/// +/// 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, A: Algorithm>( + 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 [`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 +/// positions as described by [`store`]. +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, + ) { + 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 { + #[inline(always)] + fn store( + self, + group: ThreadGroup, + destination: &mut DisjointSlice<'_, T, CollectiveTiles>, + options: StoreOptions, + ) { + let mut value = ThreadData::<_, 1>::new(); + value.write(0, self); + (&mut value).store(group, destination, options); + } +} + +mod sealed { + pub trait LoadSource {} + 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 new file mode 100644 index 000000000000..88e8a239e439 --- /dev/null +++ b/rust/cuda-coop/src/load_store/boundary.rs @@ -0,0 +1,219 @@ +// 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 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"); + +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 + * 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 = if offset < period { + offset + } else { + 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< + T: Numeric, + Memory, + const ITEMS: usize, + const ALIGN: usize, + A: Algorithm, +>( + 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(); + // 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, 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!(), + } + } +} + +#[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, + ), + } + } +} + +#[inline(always)] +pub(super) fn store_transpose< + T: Numeric, + const ITEMS: usize, + const ALIGN: usize, + const CAPACITY: usize, +>( + 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/src/thread_data.rs b/rust/cuda-coop/src/thread_data.rs new file mode 100644 index 000000000000..76203c847736 --- /dev/null +++ b/rust/cuda-coop/src/thread_data.rs @@ -0,0 +1,287 @@ +// 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], + // 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], +} + +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 / 8] |= 1 << (index % 8); + } + + /// Copy an initialized value, or return `None` if unwritten or out of bounds. + #[inline(always)] + pub fn get(&self, index: usize) -> Option { + if index >= ITEMS { + return None; + } + if (self.initialized[index / 8] & (1 << (index % 8))) != 0 { + // SAFETY: write and the private native boundary set this bit only + // 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 / 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 + } + } +} + +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); + 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() } + } +} + +impl IndexMut + for ThreadData +where + Alignment: SupportedAlignment, +{ + #[inline] + fn index_mut(&mut self, index: usize) -> &mut T { + assert!(index < ITEMS); + 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() } + } +} + +/// 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, +{ +} + +#[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 new file mode 100644 index 000000000000..a155164ad459 --- /dev/null +++ b/rust/cuda-coop/src/thread_data/block_io.rs @@ -0,0 +1,26 @@ +// Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +// 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; + ( + self.values.as_mut_ptr().cast::(), + self.initialized.as_mut_ptr(), + ) + } +} diff --git a/rust/cuda-coop/tests/check_load_store.py b/rust/cuda-coop/tests/check_load_store.py new file mode 100644 index 000000000000..e10ab4d21d19 --- /dev/null +++ b/rust/cuda-coop/tests/check_load_store.py @@ -0,0 +1,203 @@ +# 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( + "--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() +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( + 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 +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"), + "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", + ), + "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, + ), +} +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_varying_participation=( + "if thread::threadIdx_x()&1==0 {" + transpose_load + "}", + "lane-varying", + ), + 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", + "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/device_queries/.gitignore b/rust/cuda-coop/tests/device_queries/.gitignore new file mode 100644 index 000000000000..74bffa8af4a4 --- /dev/null +++ b/rust/cuda-coop/tests/device_queries/.gitignore @@ -0,0 +1,7 @@ +# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. +# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception +/target/ +/.oxide-artifacts/ +/*.ll +/*.options +/*.target diff --git a/rust/cuda-coop/tests/device_queries/Cargo.lock b/rust/cuda-coop/tests/device_queries/Cargo.lock new file mode 100644 index 000000000000..2026e4b3bedf --- /dev/null +++ b/rust/cuda-coop/tests/device_queries/Cargo.lock @@ -0,0 +1,785 @@ +# This file is automatically @generated by Cargo. +# It is not intended for manual editing. +version = 4 + +[[package]] +name = "aho-corasick" +version = "1.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c982642fa9e8606056828ee9a8505737230110bb1099153c79efe865c59d12ba" +dependencies = [ + "memchr", +] + +[[package]] +name = "anyhow" +version = "1.0.104" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "330a5ed07fa54e4702c9d6c4174f74427fc0ef6e214bbd677ae50a5099946470" + +[[package]] +name = "bindgen" +version = "0.69.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "271383c67ccabffb7381723dea0672a673f292304fcb45c01cc648c7a8d58088" +dependencies = [ + "bitflags", + "cexpr", + "clang-sys", + "itertools", + "lazy_static", + "lazycell", + "log", + "prettyplease", + "proc-macro2", + "quote", + "regex", + "rustc-hash", + "shlex", + "syn 2.0.119", + "which", +] + +[[package]] +name = "bitflags" +version = "2.13.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3ded4057c258ba199e2d26386d3af3780957ecaee6c4ef4041c6b4b8b97c0b06" + +[[package]] +name = "block-buffer" +version = "0.10.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3078c7629b62d3f0439517fa394996acacc5cbc91c5a20d8c658e77abd503a71" +dependencies = [ + "generic-array", +] + +[[package]] +name = "cexpr" +version = "0.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6fac387a98bb7c37292057cffc56d62ecb629900026402633ae9160df93a8766" +dependencies = [ + "nom", +] + +[[package]] +name = "cfg-if" +version = "1.0.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4e7648175b45a9a48536d676f68d918270699102aa8dab5496df06904c914600" + +[[package]] +name = "clang-sys" +version = "1.9.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "157a8ba7b480713b56f4c09fd13fc3e0a22a5dfab8097ba61cbc5feef950788a" +dependencies = [ + "glob", + "libc", + "libloading", +] + +[[package]] +name = "cpufeatures" +version = "0.2.17" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "59ed5838eebb26a2bb2e58f6d5b5316989ae9d08bab10e0e6d103e656d1b0280" +dependencies = [ + "libc", +] + +[[package]] +name = "crc32fast" +version = "1.5.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "01a7799fd6b852db0e61728dde9a204c423b44d689dbd432522543614b490e78" +dependencies = [ + "cfg-if", +] + +[[package]] +name = "crunchy" +version = "0.2.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "460fbee9c2c2f33933d720630a6a0bac33ba7053db5344fac858d4b8952d77d5" + +[[package]] +name = "crypto-common" +version = "0.1.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "78c8292055d1c1df0cce5d180393dc8cce0abec0a7102adb6c7b1eef6016d60a" +dependencies = [ + "generic-array", + "typenum", +] + +[[package]] +name = "cuda-artifact-finalizer" +version = "0.2.1" +source = "git+https://github.com/tpn/cuda-rust.git?rev=8c4a42e9c95b2a6bdc14a7531904d8469f0e7ca6#8c4a42e9c95b2a6bdc14a7531904d8469f0e7ca6" +dependencies = [ + "libnvvm-sys", + "nvjitlink-sys", + "serde", + "sha2", + "thiserror", +] + +[[package]] +name = "cuda-bindings" +version = "0.4.0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=8c4a42e9c95b2a6bdc14a7531904d8469f0e7ca6#8c4a42e9c95b2a6bdc14a7531904d8469f0e7ca6" +dependencies = [ + "bindgen", + "libloading", + "prettyplease", + "proc-macro2", + "quote", + "syn 2.0.119", +] + +[[package]] +name = "cuda-coop" +version = "0.0.0" +dependencies = [ + "cuda-device", +] + +[[package]] +name = "cuda-core" +version = "0.4.0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=8c4a42e9c95b2a6bdc14a7531904d8469f0e7ca6#8c4a42e9c95b2a6bdc14a7531904d8469f0e7ca6" +dependencies = [ + "anyhow", + "cuda-bindings", + "cuda-core-derive", + "half", + "oxide-artifacts", +] + +[[package]] +name = "cuda-core-derive" +version = "0.4.0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=8c4a42e9c95b2a6bdc14a7531904d8469f0e7ca6#8c4a42e9c95b2a6bdc14a7531904d8469f0e7ca6" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", +] + +[[package]] +name = "cuda-device" +version = "0.2.1" +source = "git+https://github.com/tpn/cuda-rust.git?rev=8c4a42e9c95b2a6bdc14a7531904d8469f0e7ca6#8c4a42e9c95b2a6bdc14a7531904d8469f0e7ca6" +dependencies = [ + "cuda-macros", +] + +[[package]] +name = "cuda-host" +version = "0.2.1" +source = "git+https://github.com/tpn/cuda-rust.git?rev=8c4a42e9c95b2a6bdc14a7531904d8469f0e7ca6#8c4a42e9c95b2a6bdc14a7531904d8469f0e7ca6" +dependencies = [ + "cuda-artifact-finalizer", + "cuda-core", + "cuda-macros", + "half", + "libc", + "oxide-artifacts", + "ptx-parse", + "sha2", + "thiserror", +] + +[[package]] +name = "cuda-macros" +version = "0.2.1" +source = "git+https://github.com/tpn/cuda-rust.git?rev=8c4a42e9c95b2a6bdc14a7531904d8469f0e7ca6#8c4a42e9c95b2a6bdc14a7531904d8469f0e7ca6" +dependencies = [ + "oxide-tile-export", + "proc-macro2", + "quote", + "reserved-oxide-symbols", + "syn 2.0.119", +] + +[[package]] +name = "cuda-target-spec" +version = "0.1.0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=8c4a42e9c95b2a6bdc14a7531904d8469f0e7ca6#8c4a42e9c95b2a6bdc14a7531904d8469f0e7ca6" + +[[package]] +name = "device_queries" +version = "0.0.0" +dependencies = [ + "cuda-coop", + "cuda-core", + "cuda-device", + "cuda-host", +] + +[[package]] +name = "digest" +version = "0.10.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9ed9a281f7bc9b7576e61468ba615a66a5c8cfdff42420a70aa82701a3b1e292" +dependencies = [ + "block-buffer", + "crypto-common", +] + +[[package]] +name = "either" +version = "1.19.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0e9c71c2167ca323c882b99918929403426e2373ea17242ff5653e0d5e1058be" + +[[package]] +name = "equivalent" +version = "1.0.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "00d174d5400e5e8fd687ad1049e2f578285fa914201b1af7e8b112a4546bd826" + +[[package]] +name = "errno" +version = "0.3.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "39cab71617ae0d63f51a36d69f866391735b51691dbda63cf6f96d042b63efeb" +dependencies = [ + "libc", + "windows-sys 0.61.2", +] + +[[package]] +name = "foldhash" +version = "0.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d9c4f5dac5e15c24eb999c26181a6ca40b39fe946cbe4c263c7209467bc83af2" + +[[package]] +name = "generic-array" +version = "0.14.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "85649ca51fd72272d7821adaf274ad91c288277713d9c18820d8499a7ff69e9a" +dependencies = [ + "typenum", + "version_check", +] + +[[package]] +name = "glob" +version = "0.3.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e4eba85ea1d0a966a983acd07deee566e67395d2d96b6fb39e62b5a833f1eb0b" + +[[package]] +name = "half" +version = "2.7.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" +dependencies = [ + "cfg-if", + "crunchy", + "zerocopy", +] + +[[package]] +name = "hashbrown" +version = "0.15.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9229cfe53dfd69f0609a49f65461bd93001ea1ef889cd5529dd176593f5338a1" +dependencies = [ + "foldhash", +] + +[[package]] +name = "hashbrown" +version = "0.17.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ed5909b6e89a2db4456e54cd5f673791d7eca6732202bbf2a9cc504fe2f9b84a" + +[[package]] +name = "home" +version = "0.5.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cc627f471c528ff0c4a49e1d5e60450c8f6461dd6d10ba9dcd3a61d3dff7728d" +dependencies = [ + "windows-sys 0.61.2", +] + +[[package]] +name = "indexmap" +version = "2.14.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cc4e190f5d26ca7051642629da2c52fc03bde85a03197c99408dcd291734c855" +dependencies = [ + "equivalent", + "hashbrown 0.17.1", +] + +[[package]] +name = "itertools" +version = "0.12.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ba291022dbbd398a455acf126c1e341954079855bc60dfdda641363bd6922569" +dependencies = [ + "either", +] + +[[package]] +name = "itoa" +version = "1.0.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8f42a60cbdf9a97f5d2305f08a87dc4e09308d1276d28c869c684d7777685682" + +[[package]] +name = "lazy_static" +version = "1.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "20870f649af7073d53e38067b2a84312175d56ea15217e1b15bc83506ec50afb" + +[[package]] +name = "lazycell" +version = "1.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "830d08ce1d1d941e6b30645f1a0eb5643013d835ce3779a5fc208261dbe10f55" + +[[package]] +name = "libc" +version = "0.2.191" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "65cd87c9a6d60536580b316ad985726104730a5151c99845f2490428ef91240a" + +[[package]] +name = "libloading" +version = "0.8.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d7c4b02199fee7c5d21a5ae7d8cfa79a6ef5bb2fc834d6e9058e89c825efdc55" +dependencies = [ + "cfg-if", + "windows-link", +] + +[[package]] +name = "libnvvm-sys" +version = "0.2.1" +source = "git+https://github.com/tpn/cuda-rust.git?rev=8c4a42e9c95b2a6bdc14a7531904d8469f0e7ca6#8c4a42e9c95b2a6bdc14a7531904d8469f0e7ca6" +dependencies = [ + "cuda-target-spec", + "libloading", + "thiserror", +] + +[[package]] +name = "linux-raw-sys" +version = "0.4.15" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d26c52dbd32dccf2d10cac7725f8eae5296885fb5703b261f7d0a0739ec807ab" + +[[package]] +name = "log" +version = "0.4.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9f8bd3e56ce4dfc153cf470fffbfa98c7620958b312ca5c3a4b8d5181fd13c6" + +[[package]] +name = "memchr" +version = "2.8.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cf8baf1c55e62ffcace7a9f06f4bd9cd3f0c4beb022d3b367256b91b87513d98" + +[[package]] +name = "minimal-lexical" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "68354c5c6bd36d73ff3feceb05efa59b6acb7626617f4962be322a825e61f79a" + +[[package]] +name = "nom" +version = "7.1.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d273983c5a657a70a3e8f2a01329822f3b8c8172b73826411a55751e404a0a4a" +dependencies = [ + "memchr", + "minimal-lexical", +] + +[[package]] +name = "nvjitlink-sys" +version = "0.2.1" +source = "git+https://github.com/tpn/cuda-rust.git?rev=8c4a42e9c95b2a6bdc14a7531904d8469f0e7ca6#8c4a42e9c95b2a6bdc14a7531904d8469f0e7ca6" +dependencies = [ + "libloading", + "thiserror", +] + +[[package]] +name = "object" +version = "0.36.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "62948e14d923ea95ea2c7c86c71013138b66525b86bdc08d2dcc262bdb497b87" +dependencies = [ + "crc32fast", + "hashbrown 0.15.5", + "indexmap", + "memchr", +] + +[[package]] +name = "once_cell" +version = "1.21.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9f7c3e4beb33f85d45ae3e3a1792185706c8e16d043238c593331cc7cd313b50" + +[[package]] +name = "oxide-artifacts" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "74d9a941c0c43d2dceaf25f5c9ca4c9f9cceb28dfa5e4e3cfd14e476f003e7c9" +dependencies = [ + "object", +] + +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +source = "git+https://github.com/tpn/cuda-rust.git?rev=8c4a42e9c95b2a6bdc14a7531904d8469f0e7ca6#8c4a42e9c95b2a6bdc14a7531904d8469f0e7ca6" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + +[[package]] +name = "prettyplease" +version = "0.2.37" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "479ca8adacdd7ce8f1fb39ce9ecccbfe93a3f1344b3d0d97f20bc0196208f62b" +dependencies = [ + "proc-macro2", + "syn 2.0.119", +] + +[[package]] +name = "proc-macro2" +version = "1.0.107" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "985e7ec9bb745e6ce6535b544d84d6cd6f7ad8bd711c398938ae983b91a766d9" +dependencies = [ + "unicode-ident", +] + +[[package]] +name = "ptx-parse" +version = "0.1.0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=8c4a42e9c95b2a6bdc14a7531904d8469f0e7ca6#8c4a42e9c95b2a6bdc14a7531904d8469f0e7ca6" + +[[package]] +name = "quote" +version = "1.0.47" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1fbf4db142a473a8d80c26bbf18454ed458bf8d26c8219c331daecfdbd079001" +dependencies = [ + "proc-macro2", +] + +[[package]] +name = "regex" +version = "1.13.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f020237b6c8eed93db2e2cb53c00c60a8e1bc73da7d073199a1180401450218d" +dependencies = [ + "aho-corasick", + "memchr", + "regex-automata", + "regex-syntax", +] + +[[package]] +name = "regex-automata" +version = "0.4.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ad8553b9b26413251cbf30e620595c7a41b3887f03da04579c0e6b0d6a06b4b2" +dependencies = [ + "aho-corasick", + "memchr", + "regex-syntax", +] + +[[package]] +name = "regex-syntax" +version = "0.8.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d6f6ff9a378485b298a5286656da665ba74413d36db0979633275d2e708145d4" + +[[package]] +name = "reserved-oxide-symbols" +version = "0.2.1" +source = "git+https://github.com/tpn/cuda-rust.git?rev=8c4a42e9c95b2a6bdc14a7531904d8469f0e7ca6#8c4a42e9c95b2a6bdc14a7531904d8469f0e7ca6" + +[[package]] +name = "rustc-hash" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "08d43f7aa6b08d49f382cde6a7982047c3426db949b1424bc4b7ec9ae12c6ce2" + +[[package]] +name = "rustix" +version = "0.38.44" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fdb5bc1ae2baa591800df16c9ca78619bf65c0488b41b96ccec5d11220d8c154" +dependencies = [ + "bitflags", + "errno", + "libc", + "linux-raw-sys", + "windows-sys 0.59.0", +] + +[[package]] +name = "serde" +version = "1.0.229" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4148590afebada386688f18773da617792bf2ef03ffc1e4cbd2b1d45b023e0ba" +dependencies = [ + "serde_core", + "serde_derive", +] + +[[package]] +name = "serde_core" +version = "1.0.229" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "67dca2c9c51e58a4791a4b1ed58308b39c64224d349a935ab5039aa360942a48" +dependencies = [ + "serde_derive", +] + +[[package]] +name = "serde_derive" +version = "1.0.229" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e7a5d71263a5a7d47b41f6b3f06ba276f10cc18b0931f1799f710578e2309348" +dependencies = [ + "proc-macro2", + "quote", + "syn 3.0.9", +] + +[[package]] +name = "serde_json" +version = "1.0.154" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e7e9cc8b1b85264074fbcc02a88680c4096b1e47df8f739dceb03bf482f04bd6" +dependencies = [ + "itoa", + "memchr", + "serde", + "serde_core", + "zmij", +] + +[[package]] +name = "sha2" +version = "0.10.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a7507d819769d01a365ab707794a4084392c824f54a7a6a7862f8c3d0892b283" +dependencies = [ + "cfg-if", + "cpufeatures", + "digest", +] + +[[package]] +name = "shlex" +version = "1.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0fda2ff0d084019ba4d7c6f371c95d8fd75ce3524c3cb8fb653a3023f6323e64" + +[[package]] +name = "syn" +version = "2.0.119" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "872831b642d1a07999a962a351ed35b955ea2cfc8f3862091e2a240a84f17297" +dependencies = [ + "proc-macro2", + "quote", + "unicode-ident", +] + +[[package]] +name = "syn" +version = "3.0.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d78c8dee4c7bf0e14673097256fed6142ce9d3b85a408189d07482442145823b" +dependencies = [ + "proc-macro2", + "quote", + "unicode-ident", +] + +[[package]] +name = "thiserror" +version = "2.0.21" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "09e52cb86a36cede5cb101bf8908837b3e4c6e5e59fe7fd85c23fb56200d189e" +dependencies = [ + "thiserror-impl", +] + +[[package]] +name = "thiserror-impl" +version = "2.0.21" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fe5197923287db20a58125f0bc85c062f7f2c892de97b18c356f9efb14b28524" +dependencies = [ + "proc-macro2", + "quote", + "syn 3.0.9", +] + +[[package]] +name = "typenum" +version = "1.20.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b6f5e870be6c3b371b77fe0ee0bafb859fa4964b4404c27de1d380043c4dda20" + +[[package]] +name = "unicode-ident" +version = "1.0.27" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a2c754d6c33795a1c324727428e5a7dedb5b06195f9890bdbcba760d3e246563" + +[[package]] +name = "version_check" +version = "0.9.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0b928f33d975fc6ad9f86c8f283853ad26bdd5b10b7f1542aa2fa15e2289105a" + +[[package]] +name = "which" +version = "4.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "87ba24419a2078cd2b0f2ede2691b6c66d8e47836da3b6db8265ebad47afbfc7" +dependencies = [ + "either", + "home", + "once_cell", + "rustix", +] + +[[package]] +name = "windows-link" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f0805222e57f7521d6a62e36fa9163bc891acd422f971defe97d64e70d0a4fe5" + +[[package]] +name = "windows-sys" +version = "0.59.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1e38bc4d79ed67fd075bcc251a1c39b32a1776bbe92e5bef1f0bf1f8c531853b" +dependencies = [ + "windows-targets", +] + +[[package]] +name = "windows-sys" +version = "0.61.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ae137229bcbd6cdf0f7b80a31df61766145077ddf49416a728b02cb3921ff3fc" +dependencies = [ + "windows-link", +] + +[[package]] +name = "windows-targets" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9b724f72796e036ab90c1021d4780d4d3d648aca59e491e6b98e725b84e99973" +dependencies = [ + "windows_aarch64_gnullvm", + "windows_aarch64_msvc", + "windows_i686_gnu", + "windows_i686_gnullvm", + "windows_i686_msvc", + "windows_x86_64_gnu", + "windows_x86_64_gnullvm", + "windows_x86_64_msvc", +] + +[[package]] +name = "windows_aarch64_gnullvm" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "32a4622180e7a0ec044bb555404c800bc9fd9ec262ec147edd5989ccd0c02cd3" + +[[package]] +name = "windows_aarch64_msvc" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "09ec2a7bb152e2252b53fa7803150007879548bc709c039df7627cabbd05d469" + +[[package]] +name = "windows_i686_gnu" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8e9b5ad5ab802e97eb8e295ac6720e509ee4c243f69d781394014ebfe8bbfa0b" + +[[package]] +name = "windows_i686_gnullvm" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0eee52d38c090b3caa76c563b86c3a4bd71ef1a819287c19d586d7334ae8ed66" + +[[package]] +name = "windows_i686_msvc" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "240948bc05c5e7c6dabba28bf89d89ffce3e303022809e73deaefe4f6ec56c66" + +[[package]] +name = "windows_x86_64_gnu" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "147a5c80aabfbf0c7d901cb5895d1de30ef2907eb21fbbab29ca94c5b08b1a78" + +[[package]] +name = "windows_x86_64_gnullvm" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "24d5b23dc417412679681396f2b49f3de8c1473deb516bd34410872eff51ed0d" + +[[package]] +name = "windows_x86_64_msvc" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "589f6da84c646204747d1270a2a5661ea66ed1cced2631d546fdfb155959f9ec" + +[[package]] +name = "zerocopy" +version = "0.8.63" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e5fe1f8f1b06191a00962174c61aa5005e0bb391a6d80d07e24d115c01a92ed8" +dependencies = [ + "zerocopy-derive", +] + +[[package]] +name = "zerocopy-derive" +version = "0.8.63" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "863ad3ac83293fb4d740aedbfdc9240dd8d1a50c1099acd76ce80ce7c7230c7f" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", +] + +[[package]] +name = "zmij" +version = "1.0.23" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "29666d0abbfad1e3dc4dcf6144730dd3a3ab225bbbdac83319345b1b44ccfc1b" diff --git a/rust/cuda-coop/tests/device_queries/Cargo.toml b/rust/cuda-coop/tests/device_queries/Cargo.toml new file mode 100644 index 000000000000..7e29090be630 --- /dev/null +++ b/rust/cuda-coop/tests/device_queries/Cargo.toml @@ -0,0 +1,16 @@ +# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. +# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +[package] +name = "device_queries" +version = "0.0.0" +edition = "2024" +publish = false + +[workspace] + +[dependencies] +cuda-coop = { path = "../.." } +cuda-device = { git = "https://github.com/tpn/cuda-rust.git", rev = "8c4a42e9c95b2a6bdc14a7531904d8469f0e7ca6" } +cuda-host = { git = "https://github.com/tpn/cuda-rust.git", rev = "8c4a42e9c95b2a6bdc14a7531904d8469f0e7ca6" } +cuda-core = { git = "https://github.com/tpn/cuda-rust.git", rev = "8c4a42e9c95b2a6bdc14a7531904d8469f0e7ca6" } diff --git a/rust/cuda-coop/tests/device_queries/src/main.rs b/rust/cuda-coop/tests/device_queries/src/main.rs new file mode 100644 index 000000000000..b7188d1def73 --- /dev/null +++ b/rust/cuda-coop/tests/device_queries/src/main.rs @@ -0,0 +1,420 @@ +// 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, 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) => {{ + 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 = output_index() * 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 = 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] + 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(), + 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(output_index()).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 <= 4); + 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 or payload access must trap"); + assert!(err.to_string().contains("719"), "unexpected error: {err}"); + println!("PASS rejected invalid query/payload 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()); + } + 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(()) +} diff --git a/rust/cuda-coop/tests/load_store/.gitignore b/rust/cuda-coop/tests/load_store/.gitignore new file mode 100644 index 000000000000..74bffa8af4a4 --- /dev/null +++ b/rust/cuda-coop/tests/load_store/.gitignore @@ -0,0 +1,7 @@ +# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. +# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception +/target/ +/.oxide-artifacts/ +/*.ll +/*.options +/*.target diff --git a/rust/cuda-coop/tests/load_store/Cargo.lock b/rust/cuda-coop/tests/load_store/Cargo.lock new file mode 100644 index 000000000000..4e95c613b326 --- /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=204e4e6f721bd111bf5e0e00caa69dc672a4d378#204e4e6f721bd111bf5e0e00caa69dc672a4d378" +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=204e4e6f721bd111bf5e0e00caa69dc672a4d378#204e4e6f721bd111bf5e0e00caa69dc672a4d378" +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=204e4e6f721bd111bf5e0e00caa69dc672a4d378#204e4e6f721bd111bf5e0e00caa69dc672a4d378" +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=204e4e6f721bd111bf5e0e00caa69dc672a4d378#204e4e6f721bd111bf5e0e00caa69dc672a4d378" +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=204e4e6f721bd111bf5e0e00caa69dc672a4d378#204e4e6f721bd111bf5e0e00caa69dc672a4d378" +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=204e4e6f721bd111bf5e0e00caa69dc672a4d378#204e4e6f721bd111bf5e0e00caa69dc672a4d378" +dependencies = [ + "cuda-macros", +] + +[[package]] +name = "cuda-host" +version = "0.2.1" +source = "git+https://github.com/tpn/cuda-rust.git?rev=204e4e6f721bd111bf5e0e00caa69dc672a4d378#204e4e6f721bd111bf5e0e00caa69dc672a4d378" +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=204e4e6f721bd111bf5e0e00caa69dc672a4d378#204e4e6f721bd111bf5e0e00caa69dc672a4d378" +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=204e4e6f721bd111bf5e0e00caa69dc672a4d378#204e4e6f721bd111bf5e0e00caa69dc672a4d378" + +[[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=204e4e6f721bd111bf5e0e00caa69dc672a4d378#204e4e6f721bd111bf5e0e00caa69dc672a4d378" +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=204e4e6f721bd111bf5e0e00caa69dc672a4d378#204e4e6f721bd111bf5e0e00caa69dc672a4d378" +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=204e4e6f721bd111bf5e0e00caa69dc672a4d378#204e4e6f721bd111bf5e0e00caa69dc672a4d378" +dependencies = [ + "object", +] + +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +source = "git+https://github.com/tpn/cuda-rust.git?rev=204e4e6f721bd111bf5e0e00caa69dc672a4d378#204e4e6f721bd111bf5e0e00caa69dc672a4d378" +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=204e4e6f721bd111bf5e0e00caa69dc672a4d378#204e4e6f721bd111bf5e0e00caa69dc672a4d378" + +[[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=204e4e6f721bd111bf5e0e00caa69dc672a4d378#204e4e6f721bd111bf5e0e00caa69dc672a4d378" + +[[package]] +name = "rustc-hash" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "08d43f7aa6b08d49f382cde6a7982047c3426db949b1424bc4b7ec9ae12c6ce2" + +[[package]] +name = "rustc-hash" +version = "2.1.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6b1e7f9a428571be2dc5bc0505c13fb6bf936822b894ec87abf8a08a4e51742d" + +[[package]] +name = "rustix" +version = "0.38.44" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fdb5bc1ae2baa591800df16c9ca78619bf65c0488b41b96ccec5d11220d8c154" +dependencies = [ + "bitflags", + "errno", + "libc", + "linux-raw-sys", + "windows-sys 0.59.0", +] + +[[package]] +name = "scopeguard" +version = "1.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "94143f37725109f92c262ed2cf5e59bce7498c01bcc1502d7b9afe439a4e9f49" + +[[package]] +name = "serde" +version = "1.0.229" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4148590afebada386688f18773da617792bf2ef03ffc1e4cbd2b1d45b023e0ba" +dependencies = [ + "serde_core", + "serde_derive", +] + +[[package]] +name = "serde_core" +version = "1.0.229" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "67dca2c9c51e58a4791a4b1ed58308b39c64224d349a935ab5039aa360942a48" +dependencies = [ + "serde_derive", +] + +[[package]] +name = "serde_derive" +version = "1.0.229" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e7a5d71263a5a7d47b41f6b3f06ba276f10cc18b0931f1799f710578e2309348" +dependencies = [ + "proc-macro2", + "quote", + "syn 3.0.9", +] + +[[package]] +name = "serde_json" +version = "1.0.154" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e7e9cc8b1b85264074fbcc02a88680c4096b1e47df8f739dceb03bf482f04bd6" +dependencies = [ + "itoa", + "memchr", + "serde", + "serde_core", + "zmij", +] + +[[package]] +name = "sha2" +version = "0.10.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a7507d819769d01a365ab707794a4084392c824f54a7a6a7862f8c3d0892b283" +dependencies = [ + "cfg-if", + "cpufeatures", + "digest", +] + +[[package]] +name = "shlex" +version = "1.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0fda2ff0d084019ba4d7c6f371c95d8fd75ce3524c3cb8fb653a3023f6323e64" + +[[package]] +name = "slab" +version = "0.4.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0c790de23124f9ab44544d7ac05d60440adc586479ce501c1d6d7da3cd8c9cf5" + +[[package]] +name = "smallvec" +version = "1.16.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5b3dc8af474f516a851ff4bd12db780f948b9250ad37211e4eec0bccea54e01b" + +[[package]] +name = "syn" +version = "2.0.119" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "872831b642d1a07999a962a351ed35b955ea2cfc8f3862091e2a240a84f17297" +dependencies = [ + "proc-macro2", + "quote", + "unicode-ident", +] + +[[package]] +name = "syn" +version = "3.0.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d78c8dee4c7bf0e14673097256fed6142ce9d3b85a408189d07482442145823b" +dependencies = [ + "proc-macro2", + "quote", + "unicode-ident", +] + +[[package]] +name = "thiserror" +version = "1.0.69" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b6aaf5339b578ea85b50e080feb250a3e8ae8cfcdff9a461c9ec2904bc923f52" +dependencies = [ + "thiserror-impl 1.0.69", +] + +[[package]] +name = "thiserror" +version = "2.0.21" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "09e52cb86a36cede5cb101bf8908837b3e4c6e5e59fe7fd85c23fb56200d189e" +dependencies = [ + "thiserror-impl 2.0.21", +] + +[[package]] +name = "thiserror-impl" +version = "1.0.69" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4fee6c4efc90059e10f81e6d42c60a18f76588c3d74cb83a0b242a2b6c7504c1" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", +] + +[[package]] +name = "thiserror-impl" +version = "2.0.21" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fe5197923287db20a58125f0bc85c062f7f2c892de97b18c356f9efb14b28524" +dependencies = [ + "proc-macro2", + "quote", + "syn 3.0.9", +] + +[[package]] +name = "typenum" +version = "1.20.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b6f5e870be6c3b371b77fe0ee0bafb859fa4964b4404c27de1d380043c4dda20" + +[[package]] +name = "unicode-ident" +version = "1.0.27" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a2c754d6c33795a1c324727428e5a7dedb5b06195f9890bdbcba760d3e246563" + +[[package]] +name = "version_check" +version = "0.9.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0b928f33d975fc6ad9f86c8f283853ad26bdd5b10b7f1542aa2fa15e2289105a" + +[[package]] +name = "which" +version = "4.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "87ba24419a2078cd2b0f2ede2691b6c66d8e47836da3b6db8265ebad47afbfc7" +dependencies = [ + "either", + "home", + "once_cell", + "rustix", +] + +[[package]] +name = "windows-link" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f0805222e57f7521d6a62e36fa9163bc891acd422f971defe97d64e70d0a4fe5" + +[[package]] +name = "windows-sys" +version = "0.59.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1e38bc4d79ed67fd075bcc251a1c39b32a1776bbe92e5bef1f0bf1f8c531853b" +dependencies = [ + "windows-targets", +] + +[[package]] +name = "windows-sys" +version = "0.61.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ae137229bcbd6cdf0f7b80a31df61766145077ddf49416a728b02cb3921ff3fc" +dependencies = [ + "windows-link", +] + +[[package]] +name = "windows-targets" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9b724f72796e036ab90c1021d4780d4d3d648aca59e491e6b98e725b84e99973" +dependencies = [ + "windows_aarch64_gnullvm", + "windows_aarch64_msvc", + "windows_i686_gnu", + "windows_i686_gnullvm", + "windows_i686_msvc", + "windows_x86_64_gnu", + "windows_x86_64_gnullvm", + "windows_x86_64_msvc", +] + +[[package]] +name = "windows_aarch64_gnullvm" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "32a4622180e7a0ec044bb555404c800bc9fd9ec262ec147edd5989ccd0c02cd3" + +[[package]] +name = "windows_aarch64_msvc" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "09ec2a7bb152e2252b53fa7803150007879548bc709c039df7627cabbd05d469" + +[[package]] +name = "windows_i686_gnu" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8e9b5ad5ab802e97eb8e295ac6720e509ee4c243f69d781394014ebfe8bbfa0b" + +[[package]] +name = "windows_i686_gnullvm" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0eee52d38c090b3caa76c563b86c3a4bd71ef1a819287c19d586d7334ae8ed66" + +[[package]] +name = "windows_i686_msvc" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "240948bc05c5e7c6dabba28bf89d89ffce3e303022809e73deaefe4f6ec56c66" + +[[package]] +name = "windows_x86_64_gnu" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "147a5c80aabfbf0c7d901cb5895d1de30ef2907eb21fbbab29ca94c5b08b1a78" + +[[package]] +name = "windows_x86_64_gnullvm" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "24d5b23dc417412679681396f2b49f3de8c1473deb516bd34410872eff51ed0d" + +[[package]] +name = "windows_x86_64_msvc" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "589f6da84c646204747d1270a2a5661ea66ed1cced2631d546fdfb155959f9ec" + +[[package]] +name = "zerocopy" +version = "0.8.63" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e5fe1f8f1b06191a00962174c61aa5005e0bb391a6d80d07e24d115c01a92ed8" +dependencies = [ + "zerocopy-derive", +] + +[[package]] +name = "zerocopy-derive" +version = "0.8.63" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "863ad3ac83293fb4d740aedbfdc9240dd8d1a50c1099acd76ce80ce7c7230c7f" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", +] + +[[package]] +name = "zmij" +version = "1.0.23" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "29666d0abbfad1e3dc4dcf6144730dd3a3ab225bbbdac83319345b1b44ccfc1b" diff --git a/rust/cuda-coop/tests/load_store/Cargo.toml b/rust/cuda-coop/tests/load_store/Cargo.toml new file mode 100644 index 000000000000..2b2d125d2aa3 --- /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 = "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/src/main.rs b/rust/cuda-coop/tests/load_store/src/main.rs new file mode 100644 index 000000000000..f93056d411d1 --- /dev/null +++ b/rust/cuda-coop/tests/load_store/src/main.rs @@ -0,0 +1,722 @@ +// Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. +// 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; + +#[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 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( + 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 = 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(); + 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> { + 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, LaunchConfig2D, LaunchConfig3D}; + 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) => { + 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>(($config)(groups))?; + assert_eq!(prepared.__raw_config().shared_mem_bytes, 0); + 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 { + 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 + ); + 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 + .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(()) +} diff --git a/rust/cuda-coop/tests/load_store/src/transpose.rs b/rust/cuda-coop/tests/load_store/src/transpose.rs new file mode 100644 index 000000000000..6fc2cb9fbadb --- /dev/null +++ b/rust/cuda-coop/tests/load_store/src/transpose.rs @@ -0,0 +1,798 @@ +// 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 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(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=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( + 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, LaunchConfig2D, LaunchConfig3D}; + 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) => { + 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>(($config)(groups))?; + assert!(plan.__raw_config().shared_mem_bytes > 0); + 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 { + 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 + ); + 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 + .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(()) +} diff --git a/rust/cuda-coop/tests/payload.rs b/rust/cuda-coop/tests/payload.rs new file mode 100644 index 000000000000..88e8c24b4a9e --- /dev/null +++ b/rust/cuda-coop/tests/payload.rs @@ -0,0 +1,178 @@ +// 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 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(); + 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 + ); +} diff --git a/rust/rust-toolchain.toml b/rust/rust-toolchain.toml new file mode 100644 index 000000000000..06d6cb6baf70 --- /dev/null +++ b/rust/rust-toolchain.toml @@ -0,0 +1,6 @@ +# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. +# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception +[toolchain] +channel = "nightly-2026-10-11" +components = ["clippy", "rustfmt"] +profile = "minimal"