From 53da4923a2523814081c3a37f84a9082e34620b8 Mon Sep 17 00:00:00 2001 From: Trent Nelson Date: Sat, 24 Oct 2026 17:37:21 -0700 Subject: [PATCH 1/3] Add Rust cooperative foundation and complete Load/Store Provide the Rust payload, group and scratch contracts, native CUB provider and Oxide integration with complete Block and Warp I/O. Signed-off-by: Trent Nelson --- .github/workflows/rust-host.yml | 43 + docs/index.rst | 3 + docs/rust/groups.rst | 107 ++ docs/rust/index.rst | 121 ++ docs/rust/load_store.rst | 365 ++++++ docs/rust/payload.rst | 77 ++ docs/rust/provider.rst | 163 +++ rust/.cargo/config.toml | 7 + rust/.gitignore | 1 + rust/AGENTS.md | 63 + rust/Cargo.lock | 284 +++++ rust/Cargo.toml | 18 + rust/cuda-coop-codegen/Cargo.toml | 19 + rust/cuda-coop-codegen/examples/specialize.rs | 35 + rust/cuda-coop-codegen/src/block_io.cuh | 67 ++ rust/cuda-coop-codegen/src/headers.rs | 95 ++ rust/cuda-coop-codegen/src/lib.rs | 988 ++++++++++++++++ rust/cuda-coop-codegen/src/nvrtc.rs | 307 +++++ rust/cuda-coop-codegen/src/oxide.rs | 678 +++++++++++ .../tests/check_carrier_rejections.py | 326 ++++++ .../tests/check_collective_rejections.py | 483 ++++++++ .../tests/check_storage_rejections.py | 346 ++++++ .../tests/check_validity_rejections.py | 120 ++ rust/cuda-coop-codegen/tests/native.rs | 572 +++++++++ .../tests/oxide-fixture/Cargo.toml | 14 + .../tests/oxide-fixture/README.md | 147 +++ .../tests/oxide-fixture/api/Cargo.toml | 20 + .../tests/oxide-fixture/api/build.rs | 51 + .../tests/oxide-fixture/api/src/carrier.rs | 105 ++ .../tests/oxide-fixture/api/src/lib.rs | 236 ++++ .../tests/oxide-fixture/api/src/storage.rs | 108 ++ .../tests/oxide-fixture/api/src/validity.rs | 174 +++ .../tests/oxide-fixture/carrier/Cargo.toml | 24 + .../tests/oxide-fixture/carrier/src/main.rs | 200 ++++ .../oxide-fixture/carrier/src/validity.rs | 556 +++++++++ .../compile_cases/atomic_control.rs | 34 + .../compile_cases/native_mutation.rs | 43 + .../compile_cases/nested_reference.rs | 39 + .../tests/oxide-fixture/consumer/Cargo.toml | 21 + .../tests/oxide-fixture/consumer/src/main.rs | 353 ++++++ .../oxide-fixture/consumer/src/storage.rs | 176 +++ .../tests/oxide-fixture/helper/Cargo.toml | 16 + .../tests/oxide-fixture/helper/src/checked.rs | 108 ++ .../tests/oxide-fixture/helper/src/lib.rs | 89 ++ .../tests/oxide-fixture/helper/src/storage.rs | 143 +++ .../tests/run_oxide_fixture.py | 116 ++ rust/cuda-coop/Cargo.toml | 18 + rust/cuda-coop/build.rs | 38 + rust/cuda-coop/examples/load_store/.gitignore | 7 + rust/cuda-coop/examples/load_store/Cargo.lock | 1019 +++++++++++++++++ rust/cuda-coop/examples/load_store/Cargo.toml | 19 + .../examples/load_store/helper/Cargo.toml | 11 + .../examples/load_store/helper/src/lib.rs | 116 ++ .../cuda-coop/examples/load_store/src/main.rs | 244 ++++ rust/cuda-coop/examples/payload.rs | 22 + rust/cuda-coop/src/alignment.rs | 65 ++ rust/cuda-coop/src/group.rs | 584 ++++++++++ rust/cuda-coop/src/lib.rs | 66 ++ rust/cuda-coop/src/load_store.rs | 760 ++++++++++++ rust/cuda-coop/src/load_store/boundary.rs | 440 +++++++ rust/cuda-coop/src/temp_storage.rs | 222 ++++ rust/cuda-coop/src/thread_data.rs | 287 +++++ rust/cuda-coop/src/thread_data/block_io.rs | 26 + rust/cuda-coop/tests/check_load_store.py | 902 +++++++++++++++ .../cuda-coop/tests/device_queries/.gitignore | 7 + .../cuda-coop/tests/device_queries/Cargo.lock | 798 +++++++++++++ .../cuda-coop/tests/device_queries/Cargo.toml | 16 + .../tests/device_queries/src/main.rs | 420 +++++++ rust/cuda-coop/tests/load_store/.gitignore | 7 + rust/cuda-coop/tests/load_store/Cargo.lock | 1011 ++++++++++++++++ rust/cuda-coop/tests/load_store/Cargo.toml | 17 + .../tests/load_store/src/block_warp.rs | 366 ++++++ .../tests/load_store/src/logical_plain.rs | 228 ++++ .../tests/load_store/src/logical_transpose.rs | 513 +++++++++ .../tests/load_store/src/logical_warp.rs | 504 ++++++++ rust/cuda-coop/tests/load_store/src/main.rs | 766 +++++++++++++ .../cuda-coop/tests/load_store/src/storage.rs | 357 ++++++ .../cuda-coop/tests/load_store/src/striped.rs | 415 +++++++ .../tests/load_store/src/transpose.rs | 798 +++++++++++++ .../tests/load_store/src/vectorize.rs | 380 ++++++ rust/cuda-coop/tests/load_store/src/warp.rs | 334 ++++++ rust/cuda-coop/tests/payload.rs | 178 +++ rust/rust-toolchain.toml | 6 + 83 files changed, 20028 insertions(+) create mode 100644 .github/workflows/rust-host.yml create mode 100644 docs/rust/groups.rst create mode 100644 docs/rust/index.rst create mode 100644 docs/rust/load_store.rst create mode 100644 docs/rust/payload.rst create mode 100644 docs/rust/provider.rst create mode 100644 rust/.cargo/config.toml create mode 100644 rust/.gitignore create mode 100644 rust/AGENTS.md create mode 100644 rust/Cargo.lock create mode 100644 rust/Cargo.toml create mode 100644 rust/cuda-coop-codegen/Cargo.toml create mode 100644 rust/cuda-coop-codegen/examples/specialize.rs create mode 100644 rust/cuda-coop-codegen/src/block_io.cuh create mode 100644 rust/cuda-coop-codegen/src/headers.rs create mode 100644 rust/cuda-coop-codegen/src/lib.rs create mode 100644 rust/cuda-coop-codegen/src/nvrtc.rs create mode 100644 rust/cuda-coop-codegen/src/oxide.rs create mode 100644 rust/cuda-coop-codegen/tests/check_carrier_rejections.py create mode 100644 rust/cuda-coop-codegen/tests/check_collective_rejections.py create mode 100644 rust/cuda-coop-codegen/tests/check_storage_rejections.py create mode 100644 rust/cuda-coop-codegen/tests/check_validity_rejections.py create mode 100644 rust/cuda-coop-codegen/tests/native.rs create mode 100644 rust/cuda-coop-codegen/tests/oxide-fixture/Cargo.toml create mode 100644 rust/cuda-coop-codegen/tests/oxide-fixture/README.md create mode 100644 rust/cuda-coop-codegen/tests/oxide-fixture/api/Cargo.toml create mode 100644 rust/cuda-coop-codegen/tests/oxide-fixture/api/build.rs create mode 100644 rust/cuda-coop-codegen/tests/oxide-fixture/api/src/carrier.rs create mode 100644 rust/cuda-coop-codegen/tests/oxide-fixture/api/src/lib.rs create mode 100644 rust/cuda-coop-codegen/tests/oxide-fixture/api/src/storage.rs create mode 100644 rust/cuda-coop-codegen/tests/oxide-fixture/api/src/validity.rs create mode 100644 rust/cuda-coop-codegen/tests/oxide-fixture/carrier/Cargo.toml create mode 100644 rust/cuda-coop-codegen/tests/oxide-fixture/carrier/src/main.rs create mode 100644 rust/cuda-coop-codegen/tests/oxide-fixture/carrier/src/validity.rs create mode 100644 rust/cuda-coop-codegen/tests/oxide-fixture/compile_cases/atomic_control.rs create mode 100644 rust/cuda-coop-codegen/tests/oxide-fixture/compile_cases/native_mutation.rs create mode 100644 rust/cuda-coop-codegen/tests/oxide-fixture/compile_cases/nested_reference.rs create mode 100644 rust/cuda-coop-codegen/tests/oxide-fixture/consumer/Cargo.toml create mode 100644 rust/cuda-coop-codegen/tests/oxide-fixture/consumer/src/main.rs create mode 100644 rust/cuda-coop-codegen/tests/oxide-fixture/consumer/src/storage.rs create mode 100644 rust/cuda-coop-codegen/tests/oxide-fixture/helper/Cargo.toml create mode 100644 rust/cuda-coop-codegen/tests/oxide-fixture/helper/src/checked.rs create mode 100644 rust/cuda-coop-codegen/tests/oxide-fixture/helper/src/lib.rs create mode 100644 rust/cuda-coop-codegen/tests/oxide-fixture/helper/src/storage.rs create mode 100644 rust/cuda-coop-codegen/tests/run_oxide_fixture.py create mode 100644 rust/cuda-coop/Cargo.toml create mode 100644 rust/cuda-coop/build.rs create mode 100644 rust/cuda-coop/examples/load_store/.gitignore create mode 100644 rust/cuda-coop/examples/load_store/Cargo.lock create mode 100644 rust/cuda-coop/examples/load_store/Cargo.toml create mode 100644 rust/cuda-coop/examples/load_store/helper/Cargo.toml create mode 100644 rust/cuda-coop/examples/load_store/helper/src/lib.rs create mode 100644 rust/cuda-coop/examples/load_store/src/main.rs create mode 100644 rust/cuda-coop/examples/payload.rs create mode 100644 rust/cuda-coop/src/alignment.rs create mode 100644 rust/cuda-coop/src/group.rs create mode 100644 rust/cuda-coop/src/lib.rs create mode 100644 rust/cuda-coop/src/load_store.rs create mode 100644 rust/cuda-coop/src/load_store/boundary.rs create mode 100644 rust/cuda-coop/src/temp_storage.rs create mode 100644 rust/cuda-coop/src/thread_data.rs create mode 100644 rust/cuda-coop/src/thread_data/block_io.rs create mode 100644 rust/cuda-coop/tests/check_load_store.py create mode 100644 rust/cuda-coop/tests/device_queries/.gitignore create mode 100644 rust/cuda-coop/tests/device_queries/Cargo.lock create mode 100644 rust/cuda-coop/tests/device_queries/Cargo.toml create mode 100644 rust/cuda-coop/tests/device_queries/src/main.rs create mode 100644 rust/cuda-coop/tests/load_store/.gitignore create mode 100644 rust/cuda-coop/tests/load_store/Cargo.lock create mode 100644 rust/cuda-coop/tests/load_store/Cargo.toml create mode 100644 rust/cuda-coop/tests/load_store/src/block_warp.rs create mode 100644 rust/cuda-coop/tests/load_store/src/logical_plain.rs create mode 100644 rust/cuda-coop/tests/load_store/src/logical_transpose.rs create mode 100644 rust/cuda-coop/tests/load_store/src/logical_warp.rs create mode 100644 rust/cuda-coop/tests/load_store/src/main.rs create mode 100644 rust/cuda-coop/tests/load_store/src/storage.rs create mode 100644 rust/cuda-coop/tests/load_store/src/striped.rs create mode 100644 rust/cuda-coop/tests/load_store/src/transpose.rs create mode 100644 rust/cuda-coop/tests/load_store/src/vectorize.rs create mode 100644 rust/cuda-coop/tests/load_store/src/warp.rs create mode 100644 rust/cuda-coop/tests/payload.rs create mode 100644 rust/rust-toolchain.toml diff --git a/.github/workflows/rust-host.yml b/.github/workflows/rust-host.yml new file mode 100644 index 000000000000..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..6b88f1c1ba07 --- /dev/null +++ b/docs/rust/groups.rst @@ -0,0 +1,107 @@ +.. 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 +``1.99.1`` toolchain. The device dependency is pinned in Cargo; there +is no backend-selection feature. Plain Cargo commands need the workspace's +unstable-feature configuration; see :doc:`index` for standalone crate setup. +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. On known targets below sm_90, cluster queries +describe one block per cluster. +On sm_90 and newer they describe the launch's cluster; an ordinary, unclustered +launch still has one block per cluster. An unknown target retains the sm_90 +minimum for these queries. 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. The Load/Store API +accepts power-of-two lane groups of widths 1, 2, 4, 8, 16 and 32 with Direct, Striped, Vectorize or +Transpose and implicit storage. Other mapped groups remain query-only. Canonical Warp-in-Block +and logical-group-in-Warp ranks preserve group uniformity in checked operations; +ordinary lane arithmetic does not receive that guarantee. + +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..d037aa8adbf5 --- /dev/null +++ b/docs/rust/index.rst @@ -0,0 +1,121 @@ +.. Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. +.. SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +``cuda.coop`` in Rust +===================== + +Write group-first cooperative operations in an ordinary generic Rust kernel. +This kernel copies partial block tiles through a reusable helper in another +crate. The item count stays a const generic from the host launch through the +kernel and helper: + +.. literalinclude:: ../../rust/cuda-coop/examples/load_store/src/main.rs + :language: rust + :lines: 12-35 + :dedent: 4 + +The helper uses the following calls for a partial tile. ``ThreadData`` infers its +element type from ``input``. Load fills the payload in place; Store writes the +valid prefix and leaves the rest of the destination unchanged: + +.. literalinclude:: ../../rust/cuda-coop/examples/load_store/helper/src/lib.rs + :language: rust + :lines: 55-73 + :dedent: 8 + +The complete :doc:`load_store` example includes the helper, imports, host +allocation and launch. It runs Direct, Striped, Vectorize, Transpose, +WarpTranspose and WarpTransposeTimesliced with one and four items per thread. +These listings come directly from that executable source. + +Start with a kernel +------------------- + +Use ``use cuda_coop as coop;``. A ``ThreadData::<_, ITEMS_PER_THREAD>`` holds +each thread's values and tracks which slots have been initialized. The group +argument selects the participating threads. ``DisjointSlice`` with +``CollectiveTiles`` supplies collective output ownership +without giving every lane an overlapping mutable Rust slice. + +For a block tile, calculate its size as ``block.count() * ITEMS_PER_THREAD``, +its element offset from ``block.rank_at::()``, and its valid count +as the remaining input length capped at the tile size. The kernel above shows +that idiom, including a short final tile. Counts and offsets describe the whole +group's tile, and must be uniform across that group. + +The current public operations are Load and Store, with the supported groups, +algorithms and numeric types described in :doc:`load_store`. Additional +primitive families are under development. The crate is unpublished and its +interface can change. + +Build and run +------------- + +Use the Oxide revision selected by the Cargo manifests and the pinned Rust +1.99.1 toolchain. The example's initial native configuration is Linux, CUDA 13.3 and an +NVIDIA GPU. From a complete CCCL checkout with the matching Oxide 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 + +Run plain Cargo commands from within ``rust/`` so they use its configuration. +For a standalone crate, ``cargo oxide`` enables the unstable Rust features +required by ``cuda-coop`` and ``cuda-device``. Plain ``cargo build``, ``test``, +``clippy`` and ``doc`` also need this setting in that crate's +``.cargo/config.toml``: + +.. code-block:: toml + + [env] + RUSTC_BOOTSTRAP = "1" + +Choose the architecture for your GPU. The kernel's launch contract declares +its exact block shape. Preparing a launch checks that shape and uses the +compiler's inferred shared-memory requirement. The consumer needs no C++ shim, +provider build script, manual linker inputs or scratch-byte calculation. + +How compilation works +--------------------- + +``cargo oxide`` discovers the concrete Load/Store uses, including those inside +generic helpers from another crate. The dependency's host provider specializes +CUB for their types, item counts, algorithms and launch geometry. NVRTC returns +device LTO-IR and the actual C++ temporary-storage layouts. Oxide checks +participation and resource use, plans shared storage and reuse synchronization, +and links the device code into the final CUBIN. + +The same build supplies launch geometry and resource facts to the typed host +launcher. Borrow checking remains useful, but collective participation and +scratch completion also require the compiler's checks. See :doc:`provider` for +the request format, native ABI and provider development workflow. + +Generic Rust API +---------------- + +The public API and downstream helpers use Rust types, traits and const +generics. Element types are inferred from typed producers, and ordinary calls +infer the algorithm from their options. Supported element types and algorithms +are deliberately constrained by the API's traits. + +Each concrete native specialization gets a generated device C ABI wrapper at +build time. This is an implementation boundary, not a finite set of functions +that users select or declare. The user's source stays generic Rust. + +Read next +--------- + +.. toctree:: + :maxdepth: 1 + + payload + groups + load_store + provider + +Generate the API reference with ``cargo doc --workspace --no-deps`` from +``rust/``; ``cuda_coop`` contains the public item documentation. The Rust +workspace is separate from CCCL's CMake build, and existing C++ users do not +need a Rust toolchain. diff --git a/docs/rust/load_store.rst b/docs/rust/load_store.rst new file mode 100644 index 000000000000..32e7b88e9c34 --- /dev/null +++ b/docs/rust/load_store.rst @@ -0,0 +1,365 @@ +.. Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. +.. SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +Block and Warp Load/Store +========================= + +``cuda_coop::load`` fills an existing ``ThreadData`` and returns unit. +``cuda_coop::store`` writes its values and returns unit. Both take the group +first. Complete one-, two- and three-dimensional blocks support Direct, Striped, +Vectorize, Transpose, WarpTranspose and WarpTransposeTimesliced. The two warp +transpose algorithms require a total block thread count divisible by 32. +Striped Load gives each thread elements separated by the block size; the other +algorithms produce blocked payloads, where each thread owns consecutive elements. +Physical Warp and logical lane groups of widths 1, 2, 4, 8, 16 and 32 support +Direct, Striped, Vectorize and Transpose with implicit storage when the exact +block size is a multiple of 32. WarpTranspose and WarpTransposeTimesliced are +Block algorithms. + +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 each of the six algorithms at item counts +1 and 4, including a short last tile. It also uses explicit automatic scratch in +a repeated helper call. The consumer has no build script, native shim, provider, +or scratch-size argument. +Its ``cuda-coop`` dependency supplies the native specialization during the normal +Oxide build. + +Use the Oxide revision selected by the manifests and ``1.99.1``. +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, Striped and Vectorize Load/Store and includes the planned shared storage +for the three transposing algorithms. + +.. literalinclude:: ../../rust/cuda-coop/examples/load_store/src/main.rs + :language: rust + :start-after: // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +The reusable helper preserves the source element type. Its algorithm-generic +form borrows the payload mutably because the transposing Stores permute it. + +.. 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 group'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 +participating group; sibling groups may choose different values. + +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``; a Striped slot addresses +``offset + t + X*Y*Z * i``. The exact shape is part of the kernel +launch contract; another shape with the same thread count is rejected. +Use ``LaunchConfig2D`` or ``LaunchConfig3D`` for the corresponding launch domain. +Declare the matching kernel contract, for example +``#[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, Striped and Vectorize Store accept ``&ThreadData``, a shared reborrow of ``&mut ThreadData``, +or a scalar numeric value. A scalar contributes one item per thread. All three preserve every payload value +and validity bit, including +initialized slots outside their write domain. + +Physical and logical Warp Load/Store +------------------------------------ + +Pass ``this_warp()`` through the same group-first API. No block dimensions or +extra payload type annotations are needed at the call site: + +.. code-block:: rust + + let warp = coop::this_warp(); + let mut values = coop::ThreadData::<_, ITEMS_PER_THREAD>::new(); + coop::load(warp, input, &mut values, + coop::LoadOptions::default().offset(block_origin).valid_items(count)); + coop::store(warp, output, &values, + coop::StoreOptions::default().offset(block_origin).valid_items(count)); + +Use ``this_warp().group_by::<8>()`` for an eight-lane group. Both exhaustive +and non-exhaustive forms support the six power-of-two widths; every lane is a +member for these widths. Let ``W`` be the selected width (32 for ``this_warp()``), +``t = x + X * (y + Y * z)``, ``q = t / W``, and ``lane = t % W``. +A blocked Warp slot ``i`` addresses +``offset + q * W * ITEMS_PER_THREAD + lane * ITEMS_PER_THREAD + i``; +a Striped slot addresses +``offset + q * W * ITEMS_PER_THREAD + lane + W * i``. +The operation adds this Warp origin; the caller still supplies the block +origin. The exact launch contract must contain only complete physical warps. +Direct, Striped and Vectorize use no shared scratch and accept shared payload +borrows for Store. Striped selects the striped mapping above. Full Vectorize +can use wide accesses when the actual offset-adjusted pointer is aligned; it +falls back to Direct otherwise. Counted Vectorize calls use Direct accesses. +Transpose uses disjoint shared storage for each group. The compiler reserves +one padded instance per group from the provider's actual layout; launch metadata +includes every instance. Its mutable Store permutes both values and validity. +Select it with ``.algorithm(coop::Transpose)`` on the same Load/Store options. +Explicit ``TempStorage`` is currently supported for Block calls only. + +Planned library calls cannot share a root with an unconstrained native dynamic +shared-memory view. This also applies to zero-scratch Block Direct calls with +an explicit numeric launch contract. Declare a fixed native view with +``DynamicSharedArray::get_fixed`` when composing it with planned storage. + +A collective carrier keeps its declared ``CAPACITY`` even when an operation +uses a different item count. Each group owns ``W * CAPACITY`` elements of its +block region. For example, with 64 threads, ``CAPACITY=4``, and one item per +thread, warp one owns ``[128,256)``. Caller offset zero gives the common +operation origin 32 and traps for a nonempty Store; caller offset 96 gives +origin 128 and is legal. Warp zero can keep offset zero. Supply the appropriate +block origin as well. The operation neither changes the offset nor repacks the +carrier to fit unrelated operations. Read-only slice loads have no carrier +grain restriction. With logical ``W=8``, ``CAPACITY=4`` and a scalar Store, +group one owns ``[32,64)``: offset zero gives origin 8 and traps, while offset 24 +gives origin 32. Carrier ownership does not change with an operation's item count. + +A group descriptor alone does not prove participation. Checked calls require +matching group-uniform controls and supported structured control flow. Ordinary +lane arithmetic is not inferred to be Warp-uniform merely because it divides +by a group width. Canonical ``group.rank_at::()`` for a logical +lane group and ``this_block().rank_at::()`` supply checked rank +queries. + +Conflicting accesses within a Warp tile can be completed by the matching +``cuda_device::cooperative_groups::WarpTile::sync()``. Construct it with +``this_thread_block().tiled_partition::()`` and import that module's +``ThreadGroup`` trait. Every lane in the synchronization tile must participate; +a narrower tile does not complete a wider access. A containing tile can +complete narrower accesses. Transitions involving Block accesses still need +``cuda_device::thread::sync_threads()``. Raw mask barriers are unsupported in checked +collective roots, and Warp synchronization does not release Block scratch. Direct adds no +implicit Warp barrier. + +Reusing dirty implicit Transpose storage inserts a matching Warp barrier before +the next call and completes conflicting accesses covered by that tile. First +use, scratch-free calls, and an operation in a different storage pool do not +provide this completion. A canonical containing barrier also completes narrower +implicit pools; a narrower barrier leaves wider and Block storage pending. +These barriers never end live explicit reservations or imply completion after +the final call. + +The public logical group width is ``usize`` and the canonical WarpTile width is +``u32``. An ordinary generic helper can take the group and tile as separate +arguments, or leave synchronization to its caller; no generic const-expression +feature or width cast is required. The compiler checks participation and whether +the chosen completion scope suffices for the accesses. + +Generic helpers can state ``G: coop::LoadStoreGroup`` alongside +``A: coop::Algorithm`` and ``S: coop::TempStorageLike``. Calls infer this sealed +bound from their group, algorithm and options; Block helpers retain their +existing algorithm and storage choices. + +Selecting an algorithm +---------------------- + +Select the algorithm through the trailing options. Selection is a Rust type, +so each call compiles one algorithm; ordinary defaults continue to select Direct. + +.. 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)); + +Striped uses CUB's direct striped accesses without shared scratch. Select +``.algorithm(coop::Striped)`` on both operations to load and store a striped +payload. Its Store accepts a shared payload borrow and preserves the complete +payload, including a zero-count Store and initialized slots outside its prefix. + +.. code-block:: rust + + coop::load(block, input, &mut values, + coop::LoadOptions::default().algorithm(coop::Striped)); + coop::store(block, output, &values, + coop::StoreOptions::default().algorithm(coop::Striped)); + +Vectorize preserves blocked payload order and uses no shared scratch. Full +Load and Store can use wide accesses when the numeric type, item count and +actual offset-adjusted pointer permit them. CUB falls back to Direct access +when the pointer is not sufficiently aligned; ordinary scalar alignment is +enough to call the operation. A counted overload uses Direct access even when +its count equals the tile size. Filled Load is also counted. The example helper +uses full calls for complete tiles and counted calls for its final short tile. +Its Store accepts a shared payload borrow and preserves values and validity. + +Transpose uses CUB's coalesced accesses and shared exchange within the selected group. Its Store +requires a mutable payload borrow and permutes every value and initialized bit +from blocked into striped order. For a group 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 every algorithm. + +WarpTranspose and WarpTransposeTimesliced use the same blocked Load layout. +Their Store requires a mutable payload and permutes every slot within its +physical warp: ``(t, i)`` receives old blocked slot +``q = (t / 32) * (32 * I) + 32 * i + t % 32``. Both values and validity follow +that mapping, including zero-count and repeated Stores. The timesliced variant +reuses one warp's CUB exchange storage across successive warps; its counted +Store retains a separate validity plane for all warps. Both algorithms still +require the entire block to participate, and their internal warp barriers do +not complete a conflicting full-block memory access. For shape ``[X,Y,Z]``, +``X*Y*Z`` must be divisible by 32; X alone need not be. + +The compiler allocates the provider's complete shared storage, including the +validity plane when needed, and inserts completion before reusing managed +scratch. This does not add an unconditional trailing barrier. ``auto_shared`` +uses the inferred total for the kernel. + +Explicit temporary storage +-------------------------- + +Construct ``TempStorage`` inside the kernel and pass a mutable borrow through +the options. Default storage shares one slice between its uses and requires +an explicit full-block barrier before reuse. Omitting ``temp_storage`` requests +compiler-managed reuse synchronization instead. + +.. code-block:: rust + + let mut scratch = coop::TempStorage::new().auto_sync(); + coop::load(block, input, &mut values, + coop::LoadOptions::default().algorithm(coop::Transpose) + .temp_storage(&mut scratch)); + coop::store(block, output, &mut values, + coop::StoreOptions::default().algorithm(coop::Transpose) + .temp_storage(&mut scratch)); + +The consuming ``auto_sync()`` builder inserts a trailing block barrier after +each operation that actually uses scratch. Direct, Striped and Vectorize +Load/Store receive no barrier or allocation from this policy. Options carrying a mutable storage borrow are +consumed by the operation; construct fresh options with an ordinary reborrow +for each call, including inside loops. + +``exclusive()`` gives separate static primitive and reservation sites separate +slices. Repeated execution of one site still reuses its slice. Each static +constructor site denotes one descriptor, including in a loop; different +constructor sites receive independent storage. Independent scratch does not +complete conflicting accesses to the same global-memory carrier. + +``with_capacity::()`` specifies a positive capacity for the complete +descriptor plan, including exclusive slices and padding. An undersized request +is a compilation error. ``with_alignment::()`` selects a positive +power-of-two minimum alignment for the allocation base; each exclusive slice +also satisfies its own operation or reservation requirement. Omitted capacity +and alignment are inferred. Consuming builders retain the original descriptor +identity and must not change its policy after use. No destructor synchronizes. + +Typed reservations +------------------ + +Descriptors without auto-sync support ``reserve::()`` and +``reserve_aligned::()``. ``N`` is a positive element count, and the +aligned form requests a positive power-of-two minimum for that view. Only the +ten numeric types listed above are supported. Select the element type in the reservation, for example +``scratch.reserve::()``. + +Reservations return non-Copy tokens with ``as_mut_ptr()`` and consuming +``release()`` methods. They expose no initialized Rust value, reference, slice, +or safe indexing operation. The token does not hold a descriptor-long mutable +borrow, allowing simultaneous exclusive reservations. + +All element access is unsafe. The caller must establish bounds, initialization +before reads, aliasing, non-conflicting accesses, and phase synchronization. +Saved pointers must not be accessed after release, during overlapping reuse, +or after kernel exit. Complete asynchronous accesses before ending their view. +The compiler checks planned origins and live overlap; it does not prove arbitrary +unsafe pointer code. + +To reuse a shared region, release its reservation, then synchronize the full +block before the next overlapping operation. Release does not emit a barrier or +complete global-memory accesses. A barrier before release is insufficient, and +dropping or forgetting the token does not release it. Typed reservations cannot +be combined with automatic synchronization on their descriptor. + +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`` and the six algorithm selectors +specify their full generic signatures and +supported argument forms. Native correctness and artifact checks are separate +from host compilation. Performance qualification for masked payload access is +ongoing; the current implementation makes no negligible-overhead claim. diff --git a/docs/rust/payload.rst b/docs/rust/payload.rst new file mode 100644 index 000000000000..2991f1c19515 --- /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 ``1.99.1`` toolchain. See +:doc:`groups` for physical and mapped group queries, and :doc:`load_store` for +collective Load/Store and scratch storage. + +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/docs/rust/provider.rst b/docs/rust/provider.rst new file mode 100644 index 000000000000..344b94f5df9e --- /dev/null +++ b/docs/rust/provider.rst @@ -0,0 +1,163 @@ +.. Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. +.. SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +Provider 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, block Load/Store with +Direct, Striped, Vectorize, Transpose, WarpTranspose and WarpTransposeTimesliced, +explicit block storage, and Direct, Striped, Vectorize and Transpose on physical +and logical Warp groups; see :doc:`payload`, :doc:`groups`, and :doc:`load_store`. +Its private ``cuda-coop-codegen`` crate specializes CUB block and warp operations +and returns linkable device LTO-IR with exact temporary-storage requirements. It supports +compiler integration development. Additional collective algorithms and Warp +capabilities remain under development. The crate is unpublished and its interface +can change. + +Device code uses the pinned Oxide dependency and ``1.99.1``. +The host provider supports Rust 1.85 independently; validate it with +``RUSTC_BOOTSTRAP=-1 cargo +1.85.0 test -p cuda-coop-codegen --locked`` +from ``rust/``; this disables unstable features for the provider check. + +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 carries exact ``[X,Y,Z]`` block 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. Validity-aware Direct, Striped, Vectorize and Transpose +requests also support physical warps and aligned logical warps of 1, 2, 4, 8, 16 +or 32 lanes when the block contains complete physical warps. Each group receives +a separate tile. Direct, Striped and Vectorize use no shared scratch; Transpose +uses compiler-planned shared scratch for each group. 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 the selected arrangement + - 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 + +Direct, Striped and Vectorize access no shared scratch. Direct and Vectorize +use blocked payloads; Striped uses striped payloads. Full Vectorize operations +use wide accesses when the type, item count and adjusted pointer permit them; +misaligned pointers and counted overloads use Direct accesses. Transpose Store +redistributes blocked payloads across lanes into striped order and uses the +queried CUB storage layout. WarpTranspose and WarpTransposeTimesliced exchange +within physical warps and require a complete block whose thread count is +divisible by 32. The timesliced form reuses one warp's CUB exchange buffer. +Compiler integrations enforce collective +participation, memory ownership, valid bounds and completion before scratch reuse. + +The validity-aware recipes support all six algorithms and all ten numeric +payload types: ``i8``, ``u8``, ``i16``, ``u16``, ``i32``, ``u32``, ``i64``, +``u64``, ``f32`` and ``f64``. ``Request::new_typed`` selects the exact type. +Separate packed validity bits track initialized slots. Full and default-filled +Load initialize every slot; masked Load resets slots outside its selected +prefix to unwritten. Store checks initialization throughout its write domain. +Direct, Striped and Vectorize preserve payloads and bits; the three transposing +algorithms permute both, including slots outside the prefix. Their complete +scratch layouts include a validity plane when needed. For counted warp-transpose +Store with more than one item per thread, that plane covers every warp even +when the CUB exchange buffer is timesliced. A failing Store traps and is not transactional +across lanes. Counted Store forms destination pointers only within its valid +prefix, including when the allocation is shorter than a complete tile. + +The public API uses this representation without exposing an initialized slice +or a public native layout. + +Compile a provider +------------------ + +Run Cargo from within ``rust/`` using the pinned toolchain in +``rust/rust-toolchain.toml``. The workspace +configuration enables the unstable Rust features required by Oxide and +``cuda-coop`` on that compiler. Cargo finds this configuration from the working +directory; ``--manifest-path`` alone does not select it. 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 algorithm forms at +item counts 1 and 4 plus odd-count Vectorize cases, compare repeated LTO output, +and check exact layout decoding +and specialization identity. They do not execute GPU kernels. + +Compiler integration and API reference +-------------------------------------- + +``Provider::from_checkout`` snapshots headers from the explicitly selected CCCL +checkout and toolkit, then loads NVRTC from that toolkit's ``lib64`` directory. +Those native libraries are trusted executable compiler inputs. No Python wheel, +working-directory header lookup, or runtime query kernel is involved. Reuse the +provider and batch requests for a common target into one NVRTC compilation. + +The same NVRTC program produces both LTO-IR and each CUB ``TempStorage`` layout. +The result distinguishes scratch-free specializations from an empty C++ type's +nonzero ``sizeof``. Input identities include the request, header contents, +provider source, loaded NVRTC/builtins bytes, version, target, and compiler options. +Symbols include a specialization digest. Request planning must assign each +specialization to exactly one batch per final link: a batch is one indivisible +LTO-IR buffer, so partially overlapping batches would duplicate strong definitions. + +Rustdoc is the API reference, including the generated ABI's pointer, validity, +mutation, and collective requirements. Open ``rust/target/doc/cuda_coop_codegen/index.html`` +after the documentation command above, or run ``cargo doc --no-deps --open`` from +``rust/``. Device discovery, lowering, linking, launch enforcement, and Cargo +invalidation of external compiler inputs belong to the compiler integration. + +The Oxide adapter's ``oxide::serve_inputs`` answers +``--oxide-library-inputs-v1`` with the selected header roots and NVRTC library +paths, without loading CUDA. A provider build script dispatches this query +alongside ``--oxide-library-request-v1``. Device registration getters use +``cuda_device::library::registration!("REGISTRATION_ENV_NAME")`` so the frontend +tracks those files, directory contents and selected toolkit before emitting +Cargo dependency information. Host checks and rustdoc do not query CUDA inputs. +Cargo's timestamp freshness rules apply; replacements with unchanged or older +timestamps may require an explicit rebuild. diff --git a/rust/.cargo/config.toml b/rust/.cargo/config.toml new file mode 100644 index 000000000000..0ab90f87e900 --- /dev/null +++ b/rust/.cargo/config.toml @@ -0,0 +1,7 @@ +# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. +# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +[env] +# cuda-device uses f16. +# Enable these features in processes Cargo launches from this tree. +RUSTC_BOOTSTRAP = "1" 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..8861258fbb13 --- /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=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "cuda-macros", +] + +[[package]] +name = "cuda-macros" +version = "0.2.1" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +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=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +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=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" + +[[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..44aaca2295dd --- /dev/null +++ b/rust/cuda-coop-codegen/src/block_io.cuh @@ -0,0 +1,67 @@ +// 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; + // Warp-striped CUB Load constructs a cell from an initialized input value. + __device__ opaque_cell(T value) + { + cuda::std::memcpy(bytes, &value, sizeof(T)); + } + __device__ opaque_cell(const opaque_cell& other) + { + cuda::std::memcpy(bytes, other.bytes, sizeof(T)); + } + __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..006776db9414 --- /dev/null +++ b/rust/cuda-coop-codegen/src/lib.rs @@ -0,0 +1,988 @@ +// 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 and Warp 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`], [`Algorithm::Striped`], +//! [`Algorithm::Vectorize`], [`Algorithm::Transpose`], [`Algorithm::WarpTranspose`] +//! or [`Algorithm::WarpTransposeTimesliced`]. The two warp-transpose algorithms +//! require a total block thread count divisible by 32. +//! The raw i32 recipes support all six algorithms. Their Store payloads must +//! be fully initialized. The validity-aware recipes support all ten +//! [`Dtype`] variants with separate value storage and packed per-slot validity bits. +//! Masked Load resets validity outside its prefix; Direct/Striped/Vectorize Store preserves both +//! arrays and checks initialization only in its write domain. Transposing Store +//! replaces each lane's blocked payload and validity with striped items across +//! the block or selected Warp, including when its write count is zero. Full Vectorize operations +//! use numeric local arrays after full initialization or complete Store preflight; +//! counted Vectorize overloads retain the opaque Direct path. +//! Physical and logical Warp requests support Direct, Striped, Vectorize and +//! Transpose at power-of-two widths from 1 through 32, within complete physical +//! warps. Their scratch layout describes one group, not the complete block. +//! +//! [`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 and per-thread item arrangement. +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub enum Algorithm { + /// Consecutive items per thread, without shared scratch or a reuse barrier. + Direct, + /// Thread `t` accesses item `i` at `t + group_threads * i`, without scratch. + Striped, + /// Blocked access with aligned vector operations for full tiles. + /// Misaligned pointers and counted overloads use Direct accesses. + Vectorize, + /// Coalesced access with shared-memory transposition to/from blocked layout. + Transpose, + /// Coalesced access with a separate shared exchange for each physical warp. + /// Requires a block thread count divisible by 32. + WarpTranspose, + /// Warp transposition that reuses one physical warp's CUB exchange storage. + /// Requires a block thread count divisible by 32. + WarpTransposeTimesliced, +} +impl Algorithm { + fn scratch_free(self) -> bool { + matches!(self, Self::Direct | Self::Striped | Self::Vectorize) + } + + fn token(self) -> &'static str { + match self { + Self::Direct => "DIRECT", + Self::Striped => "STRIPED", + Self::Vectorize => "VECTORIZE", + Self::Transpose => "TRANSPOSE", + Self::WarpTranspose => "WARP_TRANSPOSE", + Self::WarpTransposeTimesliced => "WARP_TRANSPOSE_TIMESLICED", + } + } +} + +/// 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. + /// Transposing algorithms redistribute the payload; all resulting slots remain initialized. + StoreFullValidity, + /// Counted Store after checking its write domain. + /// Transposing algorithms redistribute 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, + warp_width: Option, +} +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 every 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. Warp-transpose algorithms require a total thread count + /// divisible by 32; the X extent need not be divisible by 32. + pub fn new_typed_with_shape( + dtype: Dtype, + operation: Operation, + 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", + )); + } + if matches!( + algorithm, + Algorithm::WarpTranspose | Algorithm::WarpTransposeTimesliced + ) && block_threads.unwrap() % 32 != 0 + { + return Err(Error::message( + "block warp transpose requires a block thread count divisible by 32", + )); + } + Ok(Self { + operation, + algorithm, + dtype, + block, + items_per_thread, + warp_width: None, + }) + } + /// Select a physical or logical Warp specialization within this block. + /// + /// # Errors + /// Rejects non-power-of-two widths and widths outside 1..=32, partial physical warps, + /// algorithms other than Direct, Striped, Vectorize and Transpose, + /// and legacy raw recipes. The exact root block shape is retained. + pub fn with_warp_width(mut self, width: u32) -> Result { + if width == 0 + || width > 32 + || !width.is_power_of_two() + || self.block_threads() % 32 != 0 + || matches!( + self.algorithm, + Algorithm::WarpTranspose | Algorithm::WarpTransposeTimesliced + ) + || !self.operation.validity() + { + return Err(Error::message( + "Warp recipes require power-of-two width 1..=32, complete physical warps, Direct/Striped/Vectorize/Transpose and a validity-aware operation", + )); + } + self.warp_width = Some(width); + Ok(self) + } + /// 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 selected group executes the provider in converged control +/// flow. The group is the exact block, or one aligned Warp of the selected width. +/// 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 group tile base pointer in every participating lane; the +/// provider applies no block-origin or subgroup-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 group-uniform, with `0 <= valid <= tile_size`; generated +/// code traps on invalid counts. A filled Load's default is also group-uniform. +/// Scratch for transposing algorithms satisfies [`Self::scratch`] and cannot alias another live +/// region. Concurrent groups require disjoint scratch regions. The caller +/// synchronizes participating members before reuse, including loop iterations; +/// the provider does not add a trailing reuse barrier. Direct, Striped and Vectorize ignore scratch, +/// so a null pointer is permitted for that argument only. All providers return +/// 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/Striped/Vectorize Store receives const pointers and borrows them +/// read-only. Transposing Stores exclusively borrow both arrays. Destination +/// `(lane, item)` receives source `(q / items_per_thread, q % items_per_thread)`: +/// Transpose uses `q = group_threads * item + lane`; both Block warp variants use +/// `q = (lane / 32) * (32 * items_per_thread) + 32 * item + lane % 32`. +/// Lane rank is x-fastest within the block, or modulo the selected Warp width. +/// 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. The counted transposing forms transport +/// arbitrary validity in a separate shared bitmap included in [`Self::scratch`]. +/// A full Store checks the complete tile and therefore needs no shared bitmap. +/// Failure traps; it does not promise rollback of stores by other lanes. +#[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, Striped and Vectorize use zero bytes despite their empty C++ types' sizeof=1. + pub scratch: Layout, + /// Whether this specialization accesses no shared scratch or scratch-reuse barrier. + pub scratch_free: bool, + /// Whether Store redistributes payloads across lanes into striped order. + /// True for transposing 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.scratch_free(), + mutates_payload: !request.operation.is_load() && !request.algorithm.scratch_free(), + }); + } + 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.scratch_free(); + let warp_transpose = matches!( + request.algorithm, + Algorithm::WarpTranspose | Algorithm::WarpTransposeTimesliced + ); + // CUB's full vectorized overload requires the same numeric T for the + // global pointer and local array. Counted overloads keep opaque cells. + let numeric_full = request.algorithm == Algorithm::Vectorize + && matches!( + op, + Operation::LoadFullValidity | Operation::StoreFullValidity + ); + 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 + .warp_width + .unwrap_or_else(|| request.block_threads()); + let (group, primitive_type, rank) = if let Some(width) = request.warp_width { + let (primitive, algorithm) = if op.is_load() { + ("WarpLoad", "WARP_LOAD") + } else { + ("WarpStore", "WARP_STORE") + }; + ( + "warp", + format!( + "cub::{primitive}", + request.algorithm.token() + ), + format!("cub::RowMajorTid({block_x}, {block_y}, {block_z}) % {width}"), + ) + } else { + ( + "block", + format!( + "cub::{primitive}", + request.algorithm.token() + ), + format!("cub::RowMajorTid({block_x}, {block_y}, {block_z})"), + ) + }; + let bitmap_bytes = items.div_ceil(8); + let tile = threads * items; + let source_index = if warp_transpose { + format!("(tid / 32) * (32 * {items}) + 32 * i + (tid % 32)") + } else { + format!("{threads} * i + tid") + }; + let mask_plane = op == Operation::StorePartialValidity && transpose && items > 1; + let storage = if mask_plane { + let plane_bytes = u64::from(threads.div_ceil(32)) * u64::from(items) * 4; + let layout_assertions = if request.warp_width.is_some() { + format!( + "static_assert(offsetof(scratch_storage, masks) >= sizeof(primitive::TempStorage));\n\ + static_assert(offsetof(scratch_storage, masks) % alignof(unsigned int) == 0);\n\ + static_assert(offsetof(scratch_storage, masks) + {plane_bytes} <= sizeof(scratch_storage));\n\ + static_assert(alignof(scratch_storage) >= alignof(primitive::TempStorage));" + ) + } else { + "static_assert(offsetof(scratch_storage, masks) == sizeof(primitive::TempStorage));\n\ + static_assert(alignof(scratch_storage) == alignof(primitive::TempStorage));" + .to_owned() + }; + 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); +{layout_assertions}"# + ) + } 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.scratch_free() { + "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 cell_bytes = if numeric_full { + "&values[i]" + } else { + "values[i].bytes" + }; + let domain = if matches!( + op, + Operation::LoadMaskedValidity | Operation::StorePartialValidity + ) { + if request.algorithm == Algorithm::Striped { + format!("tid + {threads} * i < valid") + } else { + 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, {cell_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. Nonmutating algorithms only + // need their write domain, checked before any numeric copy below. + let copy_domain = if transpose { "true" } else { &domain }; + let prepare = if let Some(width) = request.warp_width.filter(|_| mask_plane) { + let mask = u32::MAX >> (32 - width); + format!( + r#"unsigned char* const mask_plane = scratch + offsetof(scratch_storage, masks); + const unsigned int first_lane = cuda::ptx::get_sreg_laneid() & (32u - {width}); + const unsigned int member = 0x{mask:08x}u << first_lane; +#pragma unroll + for (int i = 0; i < {items}; ++i) {{ + const unsigned int bits = __ballot_sync(member, (flags[i / 8] & (1u << (i % 8))) != 0) >> first_lane; + if (tid == 0) {{ + void* const address = mask_plane + sizeof(unsigned int) * i; + ::new (address) unsigned int(bits); + }} + }}"# + ) + } else 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 { + let completion = if request.warp_width.is_some() { + "// Complete this group's plane writes before any plane read.\n __syncwarp(member);" + } else { + "// CUB's exchange barrier orders the initialized plane writes above." + }; + format!( + r#"{completion} +#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 = {source_index}; + 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({cell_bytes}, bytes + sizeof(value_type) * i, sizeof(value_type)); + }} + operation.Store({memory}, values{operands});{restore}"# + ) + }; + ( + if numeric_full { + "value_type" + } else { + "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 = {primitive_type}; +{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 = {rank}; + {check} + {constructor} + {body} +}} +}} +"# + ) +} 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..367375864e51 --- /dev/null +++ b/rust/cuda-coop-codegen/src/oxide.rs @@ -0,0 +1,678 @@ +// 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 raw i32 and typed validity recipes for all six block algorithms +//! documented by [`Operation`]. +//! Device calls retain all participation, validity and scratch obligations of +//! [`crate::Specialization`]. No allocation or synchronization policy is supplied. + +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, + #[serde(default)] + warp_width: Option, +} + +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, + "striped" => Algorithm::Striped, + "vectorize" => Algorithm::Vectorize, + "transpose" => Algorithm::Transpose, + "warp_transpose" => Algorithm::WarpTranspose, + "warp_transpose_timesliced" => Algorithm::WarpTransposeTimesliced, + _ => { + return Err(Error::message(format!( + "unsupported algorithm {:?}", + self.algorithm + ))); + } + }; + let request = Request::new_typed_with_shape( + self.dtype.dtype()?, + operation, + algorithm, + self.block, + self.items_per_thread, + )?; + match self.warp_width { + Some(width) => request.with_warp_width(width), + None => Ok(request), + } + } +} + +#[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().scratch_free(); + 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::Striped => "striped", + Algorithm::Vectorize => "vectorize", + Algorithm::Transpose => "transpose", + Algorithm::WarpTranspose => "warp_transpose", + Algorithm::WarpTransposeTimesliced => "warp_transpose_timesliced", + }, + }) + .map_err(|error| Error::message(format!("cannot encode provider registration: {error}"))) +} + +#[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, + warp_width: None, + } + .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", + "striped", + "vectorize", + "transpose", + "warp_transpose", + "warp_transpose_timesliced", + ] { + let request = SpecializationRequest { + operation: "load_filled_validity_bits_v1".into(), + algorithm: algorithm.into(), + dtype: scalar, + block: [128, 1, 1], + items_per_thread: 4, + warp_width: None, + } + .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::Striped, + Algorithm::Vectorize, + Algorithm::Transpose, + Algorithm::WarpTranspose, + Algorithm::WarpTransposeTimesliced, + ] { + let request = Request::new_typed(Dtype::I16, operation, algorithm, 128, 4).unwrap(); + let mut expected = vec![ + Argument::MutPointer(Scalar::Signed(16)), + if algorithm.scratch_free() { + Argument::ConstPointer(Scalar::Signed(16)) + } else { + Argument::MutPointer(Scalar::Signed(16)) + }, + if algorithm.scratch_free() { + 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, + warp_width: None, + }; + 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..6d4bdbfc7d57 --- /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 barrier covering both accesses", + ), + "loop_backedge": ( + "let mut remaining = 2;\nwhile remaining > 0 {\n" + "carrier::store(&mut output, offset, valid, &mut values);\n" + "remaining -= 1;\n}", + "requires a barrier covering both accesses", + ), + "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 barrier covering both accesses", + ), + "conditional_access": ( + "if valid > 0 { carrier::store(&mut output, offset, valid, &mut values); }\n" + "carrier::reload(&mut output, offset, valid, &mut values);", + "requires a barrier covering both accesses", + ), + "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 barrier covering both accesses", + ), + "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 barrier covering both accesses", + ), + "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..3a3fcdc9579d --- /dev/null +++ b/rust/cuda-coop-codegen/tests/check_storage_rejections.py @@ -0,0 +1,346 @@ +# 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, + ), + } +) +wrapper = """ +struct Wrapper { raw: library::StorageDescriptor } +#[inline(never)] +fn use_raw(raw: &mut library::StorageDescriptor) { + unsafe { + let view = library::storage_reserve_v1::(raw); + library::storage_release_v1(view); + } +} +#[inline(never)] +fn use_wrapper(wrapper: &mut Wrapper) { use_raw(&mut wrapper.raw); } +let mut wrapper = Wrapper { raw: library::storage_new_v1::<0, 0, false, false>() }; +""" +cases.update( + { + "wrapper_raw_borrow": (wrapper + "use_raw(&mut wrapper.raw);", None), + "wrapper_local_field": ( + wrapper + + "let view=library::storage_reserve_v1::(&mut wrapper.raw);" + + release, + None, + ), + "wrapper_helper": (wrapper + "use_wrapper(&mut wrapper);", None), + "wrapper_helper_reuse": ( + wrapper + "use_wrapper(&mut wrapper);use_wrapper(&mut wrapper);", + "explicit storage reuse requires", + ), + "wrapper_helper_reuse_sync": ( + wrapper + + "use_wrapper(&mut wrapper);" + + sync + + "use_wrapper(&mut wrapper);", + None, + ), + "wrapper_returned_field": ( + wrapper + + """ +#[inline(never)] +fn raw(wrapper: &mut Wrapper) -> &mut library::StorageDescriptor { &mut wrapper.raw } +use_raw(raw(&mut wrapper)); +""", + None, + ), + "wrapper_nested": ( + wrapper + + """ +struct Outer { inner: Wrapper } +#[inline(never)] +fn use_outer(outer: &mut Outer) { use_wrapper(&mut outer.inner); } +let mut outer = Outer { inner: wrapper }; +use_outer(&mut outer); +""", + None, + ), + "wrapper_ambiguous": ( + wrapper + + "let mut other=Wrapper { raw: library::storage_new_v1::<0,0,false,false>() };" + + "use_wrapper(if valid==0 {&mut wrapper} else {&mut other});", + "no unambiguous descriptor origin", + ), + "wrapper_callee_ambiguous": ( + wrapper + + """ +#[inline(never)] +fn pick<'a>(a: &'a mut Wrapper, b: &'a mut Wrapper, count: usize) + -> &'a mut library::StorageDescriptor +{ + if count == 0 { &mut a.raw } else { &mut b.raw } +} +let mut other=Wrapper { raw: library::storage_new_v1::<0,0,false,false>() }; +use_raw(pick(&mut wrapper, &mut other, valid)); +""", + "no unambiguous descriptor origin", + ), + "wrapper_replacement": ( + wrapper + + """ +#[inline(never)] +fn replace(wrapper: &mut Wrapper, other: Wrapper) { *wrapper = other; } +let other=Wrapper { raw: library::storage_new_v1::<0,0,false,false>() }; +replace(&mut wrapper, other); +use_wrapper(&mut wrapper); +""", + "mutation through a reference to a resource descriptor", + ), + "wrapper_enum_replacement": ( + wrapper + + """ +enum Storage { Present(Wrapper), Absent } +#[inline(never)] +fn replace(storage: &mut Storage) { *storage = Storage::Absent; } +let mut storage = Storage::Present(wrapper); +replace(&mut storage); +""", + "mutation through a reference to a resource descriptor", + ), + "control_scalar_invalidation": ( + new + + """ +struct Scalar { value: usize } +#[inline(never)] +fn change(value: &mut Scalar) { value.value=thread::threadIdx_x() as usize; } +let mut control=Scalar { value: valid }; +change(&mut control); +""" + + load.replace("offset,valid", "offset,control.value"), + "lane-varying control flow", + ), + } +) +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..28d29413a9e0 --- /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 barrier covering both accesses", + ), + "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 barrier covering both accesses", + ), + "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..d29a5dd6ad7c --- /dev/null +++ b/rust/cuda-coop-codegen/tests/native.rs @@ -0,0 +1,572 @@ +// 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 algorithm in [Algorithm::WarpTranspose, Algorithm::WarpTransposeTimesliced] { + assert!( + Request::new(Operation::LoadFull, algorithm, 33, 4) + .unwrap_err() + .to_string() + .contains("divisible by 32") + ); + assert!(Request::new_with_shape(Operation::LoadFull, algorithm, [11, 3, 1], 4).is_err()); + assert!(Request::new_with_shape(Operation::LoadFull, algorithm, [8, 4, 2], 4).is_ok()); + } + for block in [ + [0, 4, 4], + [8, 0, 4], + [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() + ); + let warp = Request::new(Operation::LoadMaskedValidity, Algorithm::Direct, 64, 4).unwrap(); + for width in [0, 3, 6, 33, 64] { + assert!(warp.with_warp_width(width).is_err()); + } + for width in [1, 2, 4, 8, 16, 32] { + assert!(warp.with_warp_width(width).is_ok()); + assert!( + Request::new(Operation::LoadMaskedValidity, Algorithm::Direct, 48, 4) + .unwrap() + .with_warp_width(width) + .is_err() + ); + } + for algorithm in [Algorithm::WarpTranspose, Algorithm::WarpTransposeTimesliced] { + assert!( + Request::new(Operation::LoadMaskedValidity, algorithm, 64, 4) + .unwrap() + .with_warp_width(8) + .is_err() + ); + } + for algorithm in [ + Algorithm::Direct, + Algorithm::Striped, + Algorithm::Vectorize, + Algorithm::Transpose, + ] { + assert!( + Request::new(Operation::LoadMaskedValidity, algorithm, 64, 4) + .unwrap() + .with_warp_width(8) + .is_ok() + ); + } + assert!( + Request::new(Operation::LoadFull, Algorithm::Direct, 64, 4) + .unwrap() + .with_warp_width(8) + .is_err() + ); +} + +#[test] +#[ignore = "requires CUDA_HOME and a Linux CUDA toolkit with NVRTC LTO support"] +fn native_warp_algorithms_and_padded_masks() -> 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 cases = Vec::new(); + for width in [1, 8, 32] { + for dtype in [Dtype::I8, Dtype::I32, Dtype::F64] { + for items in [1, 4] { + cases.push((width, dtype, items, Algorithm::Transpose)); + } + } + } + for width in [2, 4, 16] { + cases.push((width, Dtype::I32, 4, Algorithm::Transpose)); + } + cases.push((8, Dtype::I8, 9, Algorithm::Transpose)); + for algorithm in [Algorithm::Striped, Algorithm::Vectorize] { + for width in [1, 8] { + cases.push((width, Dtype::I32, 4, algorithm)); + } + } + let mut requests = Vec::new(); + for &(width, dtype, items, algorithm) in &cases { + for operation in [ + Operation::LoadFullValidity, + Operation::LoadMaskedValidity, + Operation::LoadFilledValidity, + Operation::StoreFullValidity, + Operation::StorePartialValidity, + ] { + requests.push( + Request::new_typed_with_shape(dtype, operation, algorithm, [8, 4, 2], items)? + .with_warp_width(width)?, + ); + } + } + let artifact = provider.compile(&requests, &target)?; + for ((width, dtype, items, algorithm), entries) in cases + .into_iter() + .zip(artifact.specializations.chunks_exact(5)) + { + for (index, entry) in entries.iter().enumerate() { + eprintln!( + "Warp{width} {dtype:?} {algorithm:?} N{items} operation{index}: CUB {:?}, total {:?}", + entry.cub_layout, entry.scratch + ); + if algorithm != Algorithm::Transpose { + assert_eq!( + entry.scratch, + Layout { + size: 0, + alignment: 1 + } + ); + assert!(!entry.mutates_payload); + } else { + assert_eq!(entry.mutates_payload, index >= 3); + assert!(entry.scratch.alignment >= entry.cub_layout.alignment); + if index == 4 && items > 1 { + assert!(entry.scratch.alignment >= 4); + assert!(entry.scratch.size >= entry.cub_layout.size + 4 * items as usize); + if (width, dtype, items) == (8, Dtype::I8, 4) { + // CUB has 33 byte-aligned cells; the actual combined C++ + // struct pads before its 16-byte unsigned-int mask plane. + assert_eq!( + entry.cub_layout, + Layout { + size: 33, + alignment: 1 + } + ); + assert_eq!( + entry.scratch, + Layout { + size: 52, + alignment: 4 + } + ); + } + } else { + assert_eq!(entry.scratch, entry.cub_layout); + } + } + } + } + Ok(()) +} + +#[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::Striped, 33), + (Algorithm::Vectorize, 33), + (Algorithm::Transpose, 33), + (Algorithm::WarpTranspose, 64), + (Algorithm::WarpTransposeTimesliced, 64), + ] { + 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, + )?); + } + } + } + } + // The same five typed Direct recipes support each canonical logical width + // inside a multidimensional block; their actual queried scratch remains zero. + for width in [1, 2, 4, 8, 16, 32] { + for operation in [ + Operation::LoadFullValidity, + Operation::LoadMaskedValidity, + Operation::LoadFilledValidity, + Operation::StoreFullValidity, + Operation::StorePartialValidity, + ] { + requests.push( + Request::new_with_shape(operation, Algorithm::Direct, [8, 4, 2], 4)? + .with_warp_width(width)?, + ); + } + } + // 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::Striped, + Algorithm::Vectorize, + 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)?); + } + } + } + } + // Odd item counts retain CUB's scalar fallback with the same validity ABI. + for operation in [Operation::LoadFullValidity, Operation::StoreFullValidity] { + requests.push(Request::new(operation, Algorithm::Vectorize, 33, 3)?); + } + for algorithm in [Algorithm::WarpTranspose, Algorithm::WarpTransposeTimesliced] { + 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, + [8, 4, 2], + 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 matches!( + request.algorithm(), + Algorithm::Direct | Algorithm::Striped | Algorithm::Vectorize + ) { + 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.dtype() == Dtype::I32 && request.block_threads() == 64 { + let (load, store) = match (request.algorithm(), request.items_per_thread()) { + (Algorithm::WarpTranspose, 1) => (256, 272), + (Algorithm::WarpTranspose, 4) => (1024, 1040), + (Algorithm::WarpTransposeTimesliced, 1) => (128, 144), + (Algorithm::WarpTransposeTimesliced, 4) => (512, 528), + _ => unreachable!(), + }; + let is_store = matches!( + request.operation(), + Operation::StoreFullValidity | Operation::StorePartialValidity + ); + assert_eq!( + entry.cub_layout, + Layout { + size: if is_store { store } else { load }, + alignment: 16 + } + ); + } + if request.operation() == Operation::StorePartialValidity + && request.items_per_thread() > 1 + { + 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::Striped, + Algorithm::Vectorize, + Algorithm::Transpose, + Algorithm::WarpTranspose, + Algorithm::WarpTransposeTimesliced, + ] { + 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::Striped, + Algorithm::Vectorize, + Algorithm::Transpose, + Algorithm::WarpTranspose, + Algorithm::WarpTransposeTimesliced, + ] { + 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, + !matches!( + request.algorithm(), + Algorithm::Direct | Algorithm::Striped | Algorithm::Vectorize + ) && matches!( + request.operation(), + Operation::StoreFull | Operation::StorePartial + ) + ); + assert!(identities.insert(entry.identity)); + eprintln!( + "{request:?}: CUB {:?}, scratch {:?}", + entry.cub_layout, entry.scratch + ); + if matches!( + request.algorithm(), + Algorithm::Direct | Algorithm::Striped | Algorithm::Vectorize + ) { + 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); + let exchange_threads = if request.algorithm() == Algorithm::WarpTransposeTimesliced { + 32 + } else { + 128 + }; + assert!( + entry.scratch.size >= exchange_threads * request.items_per_thread() as usize * 4 + ); + } + } + // A specialization's symbol is independent of unrelated batch members. + 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..95fdabccc6ab --- /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 toolchain 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..45682aa419bd --- /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 = "96caed46c776eb49226137eac849e795b83b7ce0" } + +[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..1052620aed55 --- /dev/null +++ b/rust/cuda-coop/build.rs @@ -0,0 +1,38 @@ +// 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"), + ] { + for (algorithm, prefix) in [ + (Algorithm::Striped, "STRIPED"), + (Algorithm::Vectorize, "VECTORIZE"), + (Algorithm::Transpose, "TRANSPOSE"), + (Algorithm::WarpTranspose, "WARP_TRANSPOSE"), + ( + Algorithm::WarpTransposeTimesliced, + "WARP_TRANSPOSE_TIMESLICED", + ), + ] { + let value = oxide::registration(&executable, operation, algorithm) + .expect("cuda.coop algorithm registration"); + println!("cargo::rustc-env=COOP_{prefix}_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..98b133dcd9dd --- /dev/null +++ b/rust/cuda-coop/examples/load_store/Cargo.lock @@ -0,0 +1,1019 @@ +# 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=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "libnvvm-sys", + "nvjitlink-sys", + "serde", + "sha2", + "thiserror 2.0.21", +] + +[[package]] +name = "cuda-async" +version = "0.5.0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +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.5.0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +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.5.0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "anyhow", + "cuda-bindings", + "cuda-core-derive", + "half", + "oxide-artifacts", +] + +[[package]] +name = "cuda-core-derive" +version = "0.5.0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +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=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "cuda-macros", +] + +[[package]] +name = "cuda-host" +version = "0.2.1" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "cuda-artifact-finalizer", + "cuda-async", + "cuda-core", + "cuda-device", + "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=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +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=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" + +[[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=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +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=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +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=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "object", +] + +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +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=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" + +[[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=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" + +[[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..45b0a1df3f9b --- /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 = "96caed46c776eb49226137eac849e795b83b7ce0" } +cuda-host = { git = "https://github.com/tpn/cuda-rust.git", rev = "96caed46c776eb49226137eac849e795b83b7ce0", features = ["async"] } +cuda-core = { git = "https://github.com/tpn/cuda-rust.git", rev = "96caed46c776eb49226137eac849e795b83b7ce0" } +cuda-async = { git = "https://github.com/tpn/cuda-rust.git", rev = "96caed46c776eb49226137eac849e795b83b7ce0" } 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..76d573ce03c6 --- /dev/null +++ b/rust/cuda-coop/examples/load_store/helper/src/lib.rs @@ -0,0 +1,116 @@ +// 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(); + // Full calls let Vectorize use wide accesses for complete tiles. Its counted + // calls retain CUB's guarded Direct path for the final partial tile. + if valid == block.count() as usize * ITEMS_PER_THREAD { + coop::load( + block, + input, + &mut values, + coop::LoadOptions::from(algorithm).offset(offset), + ); + coop::store( + block, + output, + &mut values, + coop::StoreOptions::from(algorithm).offset(offset), + ); + } else { + coop::load( + block, + input, + &mut values, + coop::LoadOptions::default() + .valid_items(valid) + .oob_default(default) + .offset(offset) + .algorithm(algorithm), + ); + coop::store( + block, + output, + &mut values, + coop::StoreOptions::default() + .valid_items(valid) + .offset(offset) + .algorithm(algorithm), + ); + } +} + +/// Copy a tile using caller-selected explicit storage and synchronization policy. +pub fn copy_with_storage< + T: coop::Numeric, + const ITEMS_PER_THREAD: usize, + const C: usize, + const ALIGN: usize, + const AUTO_SYNC: bool, + const EXCLUSIVE: bool, +>( + block: coop::ThreadGroup, + input: &[T], + output: &mut coop::DisjointSlice<'_, T, coop::CollectiveTiles>, + offset: usize, + valid: usize, + default: T, + scratch: &mut coop::TempStorage, +) { + 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(coop::Transpose) + .temp_storage(scratch), + ); + coop::store( + block, + output, + &mut values, + coop::StoreOptions::default() + .valid_items(valid) + .offset(offset) + .algorithm(coop::Transpose) + .temp_storage(scratch), + ); +} 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..3255aeecde4c --- /dev/null +++ b/rust/cuda-coop/examples/load_store/src/main.rs @@ -0,0 +1,244 @@ +// 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, + ); + } + #[kernel] + #[launch_contract(domain = 1, block = (128, 1, 1), dynamic_shared = inferred)] + pub fn striped( + input: &[i16], + mut output: coop::DisjointSlice>, + ) { + let block = coop::this_block(); + let tile = block.count() as usize * ITEMS_PER_THREAD; + let offset = block.rank_at::() as usize * tile; + let remaining = if offset <= input.len() { + input.len() - offset + } else { + 0 + }; + let valid = remaining.min(tile); + coop_load_store_helper::copy_with_algorithm::<_, ITEMS_PER_THREAD, _>( + block, + input, + &mut output, + offset, + valid, + -7_i16, + coop::Striped, + ); + } + #[kernel] + #[launch_contract(domain = 1, block = (128, 1, 1), dynamic_shared = inferred)] + pub fn vectorize( + input: &[i16], + mut output: coop::DisjointSlice>, + ) { + let block = coop::this_block(); + let tile = block.count() as usize * ITEMS_PER_THREAD; + let offset = block.rank_at::() as usize * tile; + let remaining = if offset <= input.len() { + input.len() - offset + } else { + 0 + }; + let valid = remaining.min(tile); + coop_load_store_helper::copy_with_algorithm::<_, ITEMS_PER_THREAD, _>( + block, + input, + &mut output, + offset, + valid, + -7_i16, + coop::Vectorize, + ); + } + #[kernel] + #[launch_contract(domain = 1, block = (128, 1, 1), dynamic_shared = inferred)] + pub fn warp_transpose( + input: &[i16], + mut output: coop::DisjointSlice>, + ) { + let block = coop::this_block(); + let tile = block.count() as usize * ITEMS_PER_THREAD; + let offset = block.rank_at::() as usize * tile; + let remaining = if offset <= input.len() { + input.len() - offset + } else { + 0 + }; + let valid = remaining.min(tile); + coop_load_store_helper::copy_with_algorithm::<_, ITEMS_PER_THREAD, _>( + block, + input, + &mut output, + offset, + valid, + -7_i16, + coop::WarpTranspose, + ); + } + #[kernel] + #[launch_contract(domain = 1, block = (128, 1, 1), dynamic_shared = inferred)] + pub fn warp_transpose_timesliced( + input: &[i16], + mut output: coop::DisjointSlice>, + ) { + let block = coop::this_block(); + let tile = block.count() as usize * ITEMS_PER_THREAD; + let offset = block.rank_at::() as usize * tile; + let remaining = if offset <= input.len() { + input.len() - offset + } else { + 0 + }; + let valid = remaining.min(tile); + coop_load_store_helper::copy_with_algorithm::<_, ITEMS_PER_THREAD, _>( + block, + input, + &mut output, + offset, + valid, + -7_i16, + coop::WarpTransposeTimesliced, + ); + } + #[kernel] + #[launch_contract(domain = 1, block = (128, 1, 1), dynamic_shared = inferred)] + pub fn explicit_storage( + 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); + let mut scratch = coop::TempStorage::new().auto_sync(); + let mut iteration = 0; + while iteration < 2 { + coop_load_store_helper::copy_with_storage( + block, + input, + &mut output, + offset, + valid, + -7_i16, + &mut scratch, + ); + iteration += 1; + } + } +} + +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_striped, striped_async_owned); + run!(4, prepare_striped, striped_async_owned); + run!(1, prepare_vectorize, vectorize_async_owned); + run!(4, prepare_vectorize, vectorize_async_owned); + run!(1, prepare_transpose, transpose_async_owned); + run!(4, prepare_transpose, transpose_async_owned); + run!(1, prepare_warp_transpose, warp_transpose_async_owned); + run!(4, prepare_warp_transpose, warp_transpose_async_owned); + run!( + 1, + prepare_warp_transpose_timesliced, + warp_transpose_timesliced_async_owned + ); + run!( + 4, + prepare_warp_transpose_timesliced, + warp_transpose_timesliced_async_owned + ); + run!(1, prepare_explicit_storage, explicit_storage_async_owned); + run!(4, prepare_explicit_storage, explicit_storage_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..2b5ad8bcfbca --- /dev/null +++ b/rust/cuda-coop/src/group.rs @@ -0,0 +1,584 @@ +// 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, library, 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 on known +/// pre-sm_90 targets describe one block per cluster; on newer targets they +/// describe the launch cluster. Unknown targets retain the sm_90 minimum. +#[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, + library::warp_group_index_v1::<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 POWER_OF_TWO: bool = N > 0 && N <= 32 && N.is_power_of_two(); + const COMPLETE: bool = N > 0 && N <= 32 && 32 % N == 0; + 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 { + if Self::COMPLETE { + 1 + } else { + ((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)); + 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)); + assert!(N as u32 <= warps); + assert!(!E || warps.is_multiple_of(N as u32)); + 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); + $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, + if ThreadsWithinWarp::::POWER_OF_TWO { + library::warp_group_index_v1::() % (32 / N as u32) + } else { + 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..99b223393877 --- /dev/null +++ b/rust/cuda-coop/src/lib.rs @@ -0,0 +1,66 @@ +// 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 `1.99.1` toolchain. +//! Plain Cargo commands also require `RUSTC_BOOTSTRAP=1`, set by this workspace's +//! `.cargo/config.toml`; standalone consumers must supply that setting themselves. +//! [`this_block`] and the other dimension-free group constructors describe +//! the executing thread's current launch. Queries grant no collective +//! participation or synchronization guarantee. [`load`] and [`store`] support +//! Direct, Striped, Vectorize and three transpose algorithms on complete blocks. +//! Physical Warp groups and logical widths 1, 2, 4, 8, 16 and 32 support Direct, +//! Striped, Vectorize and Transpose. These operations have compiler-checked +//! participation and collective memory ownership. The transpose algorithms use +//! compiler-planned shared scratch. [`TempStorage`] selects explicit Block reuse +//! and completion policies; [`StorageReservation`] exposes typed scratch through +//! raw pointers. +//! 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 temp_storage; +mod thread_data; + +pub use alignment::{Alignment, SupportedAlignment}; +pub use cuda_device::{CollectiveTiles, DisjointSlice}; +pub use load_store::{ + Algorithm, Direct, LoadOptions, LoadSource, LoadStoreGroup, StoreOptions, StoreValue, Striped, + Transpose, Vectorize, WarpTranspose, WarpTransposeTimesliced, load, store, +}; +pub use temp_storage::{StorageReservation, TempStorage, TempStorageLike}; +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..69a9a82696b7 --- /dev/null +++ b/rust/cuda-coop/src/load_store.rs @@ -0,0 +1,760 @@ +// Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +use crate::{ + Alignment, Block, GroupScope, Numeric, SupportedAlignment, TempStorage, TempStorageLike, + ThreadData, ThreadGroup, ThreadsWithinWarp, Warp, +}; +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; + +/// Direct striped access without shared scratch. +/// +/// Thread rank `t` and item `i` address `t + group.count() * i`. Load produces +/// this striped arrangement; Store consumes it and preserves the payload. +#[derive(Clone, Copy, Debug, Default)] +pub struct Striped; + +/// Blocked access using aligned vector operations for full tiles. +/// +/// CUB checks the actual pointer after the element offset and falls back to +/// Direct access when it is not sufficiently aligned. Counts, including a +/// full-tile count, select Direct accesses. No extra alignment precondition or +/// shared scratch is required. Store preserves every payload value and bit. +#[derive(Clone, Copy, Debug, Default)] +pub struct Vectorize; + +/// Coalesced access through CUB's shared-memory exchange within the group. +/// +/// 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::Transpose.into()); +/// coop::store(block, output, &mut values, coop::Transpose.into()); +/// } +/// ``` +#[derive(Clone, Copy, Debug, Default)] +pub struct Transpose; + +/// Coalesced blocked access through a shared-memory exchange within each warp. +/// +/// The complete block participates; its total thread count must be divisible +/// by 32. Load produces blocked payloads. Store mutates every value and validity +/// bit into warp-striped order, including when its write count is zero. See +/// [`store`] for the exact permutation and synchronization requirements. +#[derive(Clone, Copy, Debug, Default)] +pub struct WarpTranspose; + +/// [`WarpTranspose`] using one warp's CUB exchange storage in successive slices. +/// +/// Every thread in the complete block participates, including the block barriers +/// between slices. The block thread count must be divisible by 32. Load, Store +/// mutation and validity follow [`WarpTranspose`]. Counted Store additionally +/// reserves its validity plane for every warp; that plane is not timesliced. +#[derive(Clone, Copy, Debug, Default)] +pub struct WarpTransposeTimesliced; + +/// A sealed, compile-time algorithm selected through Load/Store options. +pub trait Algorithm: sealed::Algorithm + Copy {} +impl Algorithm for Direct {} +impl Algorithm for Striped {} +impl Algorithm for Vectorize {} +impl Algorithm for Transpose {} +impl Algorithm for WarpTranspose {} +impl Algorithm for WarpTransposeTimesliced {} +impl sealed::Algorithm for Direct { + const KIND: u8 = sealed::DIRECT; +} +impl sealed::Algorithm for Striped { + const KIND: u8 = sealed::STRIPED; +} +impl sealed::Algorithm for Vectorize { + const KIND: u8 = sealed::VECTORIZE; +} +impl sealed::Algorithm for Transpose { + const KIND: u8 = sealed::TRANSPOSE; +} +impl sealed::Algorithm for WarpTranspose { + const KIND: u8 = sealed::WARP_TRANSPOSE; +} +impl sealed::Algorithm for WarpTransposeTimesliced { + const KIND: u8 = sealed::WARP_TRANSPOSE_TIMESLICED; +} + +/// A sealed group supported by a Load/Store algorithm and storage choice. +/// +/// Blocks support the existing algorithms and explicit or implicit storage. +/// Physical [`Warp`] and power-of-two [`ThreadsWithinWarp`] widths 1..=32 +/// support [`Direct`], [`Striped`], [`Vectorize`] and [`Transpose`] with implicit +/// storage `()`. Generic +/// helpers can carry this bound; ordinary calls infer it from the group/options. +/// A helper that also synchronizes can accept a canonical +/// `cuda_device::cooperative_groups::WarpTile` separately. Its `N` uses +/// `u32`, while the group's width uses `usize`; no generic const conversion is +/// required. The compiler checks that the tile's completion scope covers the +/// accesses that follow. +pub trait LoadStoreGroup: GroupScope + sealed::Group {} +impl LoadStoreGroup for Block {} +impl LoadStoreGroup for Warp {} +impl LoadStoreGroup for Warp {} +impl LoadStoreGroup for Warp {} +impl LoadStoreGroup for Warp {} +impl sealed::Group for Block { + const WARP: bool = false; + const WIDTH: usize = 0; + #[inline(always)] + fn group_index() -> u32 { + 0 + } + #[inline(always)] + unsafe fn call< + T: Copy, + R: Fn() -> &'static str + Copy, + M, + S, + U: Copy, + P: Copy, + const I: usize, + const WRITE: bool, + >( + registration: R, + memory: M, + storage: S, + uniform: U, + payload: P, + ) { + // SAFETY: the owning boundary supplies the typed witness and checked bounds. + unsafe { + cuda_device::library::checked_block_storage_call_v1::( + registration, + memory, + storage, + uniform, + payload, + ) + } + } +} +macro_rules! warp_group { + ($scope:ty, $width:expr $(, $e:ident)?) => { + impl<$(const $e: bool)?> sealed::Group for $scope { + const WARP: bool = true; + const WIDTH: usize = $width; + #[inline(always)] + fn group_index() -> u32 { cuda_device::library::warp_group_index_v1::<$width>() } + #[inline(always)] + unsafe fn call &'static str + Copy, M, S, U: Copy, P: Copy, + const I: usize, const WRITE: bool>(registration: R, memory: M, _storage: S, uniform: U, payload: P) { + // SAFETY: LoadStoreGroup restricts this dispatch to supported algorithms/implicit + // storage; the owning boundary checks the fixed group interval. + unsafe { cuda_device::library::checked_warp_call_v1::( + registration, memory, uniform, payload) } + } + } + }; +} +warp_group!(Warp, 32); +macro_rules! logical_groups { + ($($width:literal),+ $(,)?) => {$ ( + impl LoadStoreGroup for ThreadsWithinWarp<$width, E> {} + impl LoadStoreGroup for ThreadsWithinWarp<$width, E> {} + impl LoadStoreGroup for ThreadsWithinWarp<$width, E> {} + impl LoadStoreGroup for ThreadsWithinWarp<$width, E> {} + warp_group!(ThreadsWithinWarp<$width, E>, $width, E); + )+}; +} +logical_groups!(1, 2, 4, 8, 16, 32); + +/// 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. +/// +/// Pass an algorithm alone as `Transpose.into()`. To chain controls, start +/// with `LoadOptions::from(Transpose)`. +/// +/// Omitted storage uses compiler-managed reuse synchronization. +/// [`temp_storage`](Self::temp_storage) borrows an explicit descriptor. +/// Controls, including whether a count/default is present, must be uniform +/// across the participating group. They are checked when the operation executes. +#[derive(Clone, Copy)] +pub struct LoadOptions { + storage: S, + algorithm: PhantomData, + valid_items: Option, + oob_default: Option, + offset: usize, +} + +impl Default for LoadOptions { + fn default() -> Self { + Self { + storage: (), + algorithm: PhantomData, + valid_items: None, + oob_default: None, + offset: 0, + } + } +} + +impl From for LoadOptions { + fn from(algorithm: A) -> Self { + LoadOptions::default().algorithm(algorithm) + } +} + +impl LoadOptions { + /// Use this explicit descriptor instead of compiler-managed scratch. + /// + /// The options hold a mutable borrow until consumed by the operation. + /// Construct fresh options with an ordinary reborrow for each call in a loop. + pub fn temp_storage< + const C: usize, + const ALIGN: usize, + const AUTO_SYNC: bool, + const EXCLUSIVE: bool, + >( + self, + storage: &mut TempStorage, + ) -> LoadOptions> { + LoadOptions { + storage, + algorithm: PhantomData, + valid_items: self.valid_items, + oob_default: self.oob_default, + offset: self.offset, + } + } + + /// Select the compile-time algorithm while preserving the other controls. + pub fn algorithm(self, _algorithm: B) -> LoadOptions { + LoadOptions { + storage: self.storage, + 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 group. + 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. A Warp group of width `W` adds + /// its x-fastest group index times `W * ITEMS` automatically. The accessed + /// prefix must fit the source + /// and, for a collective carrier, the group's fixed ownership 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 group. +/// +/// Pass an algorithm alone as `Transpose.into()`. To chain controls, start +/// with `StoreOptions::from(Transpose)`. +/// +/// Omitted storage uses compiler-managed reuse synchronization; +/// [`temp_storage`](Self::temp_storage) selects explicit storage. +#[derive(Clone, Copy)] +pub struct StoreOptions { + storage: S, + algorithm: PhantomData, + valid_items: Option, + offset: usize, +} + +impl Default for StoreOptions { + fn default() -> Self { + Self { + storage: (), + algorithm: PhantomData, + valid_items: None, + offset: 0, + } + } +} + +impl From for StoreOptions { + fn from(algorithm: A) -> Self { + StoreOptions::default().algorithm(algorithm) + } +} + +impl StoreOptions { + /// Use this explicit descriptor instead of compiler-managed scratch. + /// + /// The options hold a mutable borrow until consumed by the operation. + /// Construct fresh options with an ordinary reborrow for each call in a loop. + pub fn temp_storage< + const C: usize, + const ALIGN: usize, + const AUTO_SYNC: bool, + const EXCLUSIVE: bool, + >( + self, + storage: &mut TempStorage, + ) -> StoreOptions> { + StoreOptions { + storage, + algorithm: PhantomData, + valid_items: self.valid_items, + offset: self.offset, + } + } + + /// Select the compile-time algorithm while preserving the other controls. + pub fn algorithm(self, _algorithm: B) -> StoreOptions { + StoreOptions { + storage: self.storage, + 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. A Warp group of width `W` adds + /// its x-fastest group index times `W * ITEMS` automatically. The written + /// prefix must fit the + /// allocation and the group's fixed ownership region. + pub fn offset(mut self, offset: usize) -> Self { + self.offset = offset; + self + } +} + +/// Load a group tile into an existing per-thread payload. +/// +/// `T` is inferred from `source`; `ITEMS` is the positive per-thread extent and +/// `ALIGN` is the payload's minimum alignment. The tile has +/// `group.count() * ITEMS` elements. Striped gives each thread items separated +/// by the group size; the other algorithms give it consecutive items. +/// [`LoadSource`] accepts an immutable slice or a mutable borrow of a collective +/// carrier. [`LoadOptions`] supplies the count, typed default and element offset. +/// The operation returns unit. [`Direct`], [`Striped`] and [`Vectorize`] use no +/// shared scratch. The three transposing algorithms use compiler-planned shared +/// storage and return blocked payloads. +/// +/// Complete one-, two- and three-dimensional blocks support all six algorithms. +/// [`WarpTranspose`] and [`WarpTransposeTimesliced`] require a total thread count +/// divisible by 32; they still require participation by every block thread. +/// Thread rank is x-fastest: `t = x + X * (y + Y * z)` for shape `[X, Y, Z]`. +/// Slot `i` addresses `t + group.count() * i` for Striped Load and `t * ITEMS + i` +/// for the other algorithms, relative to the selected offset. +/// Compile with the supported Oxide backend and an exact block launch contract. +/// Physical [`Warp`] and logical [`ThreadsWithinWarp`] widths 1, 2, 4, 8, 16 and +/// 32 support [`Direct`], [`Striped`], [`Vectorize`] and [`Transpose`] with +/// implicit storage when the exact block size is a multiple of 32. A width-`W` group's origin is the caller's element offset +/// plus its x-fastest group index times `W * ITEMS`. Every participating group +/// thread must reach matching calls with uniform source, controls and +/// participation. Sibling groups may choose different controls. 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 completion covering the producer and +/// consumer groups. A canonical `WarpTile::sync()` completes accesses within +/// its tile; a containing tile can also provide completion. Transitions to or +/// from Block accesses require `cuda_device::thread::sync_threads()`. +/// Warp Transpose receives disjoint shared storage for each group. Reusing its +/// dirty implicit storage inserts matching completion before the next call; +/// first use and scratch-free calls provide no such completion. Canonical +/// containing barriers also complete narrower implicit storage. +/// Ending a borrow is not completion. Raw-mask synchronization is unsupported +/// in checked collective roots. +/// 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< + T: Numeric, + Source: LoadSource, + G: LoadStoreGroup, + const ITEMS: usize, + const ALIGN: usize, + A: Algorithm, + S: TempStorageLike, +>( + group: ThreadGroup, + source: Source, + output: &mut ThreadData, + options: LoadOptions, +) where + Alignment: SupportedAlignment, +{ + source.load(group, output, options); +} + +/// Store a group tile from a per-thread payload. +/// +/// `T` is inferred from the destination and value. [`StoreValue`] accepts a +/// shared or mutable borrow of [`ThreadData`] for [`Direct`], [`Striped`] and [`Vectorize`], a +/// mutable borrow for [`Transpose`], [`WarpTranspose`] or [`WarpTransposeTimesliced`], +/// or a numeric scalar (one item per thread) +/// for any algorithm. +/// `CAPACITY` describes the carrier's ownership grain; it need not equal the +/// payload extent. Direct and Vectorize consume blocked items; Striped consumes striped items. +/// These algorithms preserve every value and initialized bit and use no shared scratch. +/// Transpose uses compiler-planned shared storage and permutes all payload slots +/// into striped order, including a zero-count Store. Destination lane/item +/// `(t, i)` receives the prior blocked slot at linear index `group.count()*i+t`. +/// Both warp-transpose algorithms instead use warp-striped index +/// `(t / 32) * (32 * ITEMS) + 32 * i + t % 32`, where `ITEMS` is the payload extent. +/// They also mutate every slot, including a zero-count Store. All three preserve +/// the initialized bit of the moved value, including unwritten slots and slots +/// outside the write prefix. Scalars use a mutable one-item temporary. +/// +/// The supported groups, compiler/launch and uniformity requirements are those +/// of [`load`]. The destination is the canonical collective carrier: it grants +/// fixed group regions without exposing duplicate mutable lane references. +/// A Warp group of width `W` owns `W * CAPACITY` elements within the block region. +/// When `CAPACITY != ITEMS`, choose an offset whose automatic Warp origin still +/// lies in that fixed region; the carrier is not repartitioned by this call. +/// Every written payload slot must be initialized. Slots outside the written +/// prefix may be unwritten. A following conflicting access to this carrier +/// needs completion covering both accesses, as described by [`load`]; 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 for Block completion. +/// +/// # Panics +/// +/// Invalid counts, offsets, bounds, ownership regions or an unwritten value in +/// the write domain trap on device. Failure does not roll back other lanes' +/// stores. Group/control/resource violations are compilation errors. +/// +/// ```no_run +/// use cuda_coop as coop; +/// fn write_scalar(output: &mut coop::DisjointSlice<'_, i16, coop::CollectiveTiles<1>>) { +/// coop::store(coop::this_block(), output, 7_i16, Default::default()); +/// } +/// ``` +#[inline(always)] +pub fn store< + T: Numeric, + V: StoreValue, + G: LoadStoreGroup, + const CAPACITY: usize, + A: Algorithm, + S: TempStorageLike, +>( + 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< + const ITEMS: usize, + const ALIGN: usize, + A: Algorithm, + S: TempStorageLike, + G: LoadStoreGroup, + >( + self, + group: ThreadGroup, + output: &mut ThreadData, + options: LoadOptions, + ) where + Alignment: SupportedAlignment; +} + +impl sealed::LoadSource for &[T] {} +impl LoadSource for &[T] { + #[inline(always)] + fn load< + const ITEMS: usize, + const ALIGN: usize, + A: Algorithm, + S: TempStorageLike, + G: LoadStoreGroup, + >( + self, + _group: ThreadGroup, + output: &mut ThreadData, + options: LoadOptions, + ) where + Alignment: SupportedAlignment, + { + // SAFETY: the pointer, length and immutable witness describe this slice. + unsafe { + boundary::load::<_, _, ITEMS, ALIGN, A, S, G>( + 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< + const ITEMS: usize, + const ALIGN: usize, + A: Algorithm, + S: TempStorageLike, + G: LoadStoreGroup, + >( + self, + _group: ThreadGroup, + output: &mut ThreadData, + options: LoadOptions, + ) where + Alignment: SupportedAlignment, + { + let valid = match options.valid_items { + Some(count) => count, + None => boundary::group_tile::(), + }; + let offset = boundary::group_offset::(options.offset); + boundary::group_containment::(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::<_, _, ITEMS, ALIGN, A, S, G>(&*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 all algorithms. +/// Shared payload borrows support [`Direct`], [`Striped`] and [`Vectorize`]. A scalar contributes +/// one item per thread. These algorithms preserve every value and bit; +/// the transposing algorithms mutate 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, + ); +} + +macro_rules! shared_store_value { + ($algorithm:ty) => { + impl sealed::StoreValue + for &ThreadData + where + Alignment: SupportedAlignment, + { + } + impl StoreValue + for &ThreadData + where + Alignment: SupportedAlignment, + { + #[inline(always)] + fn store>( + self, + _group: ThreadGroup, + destination: &mut DisjointSlice<'_, T, CollectiveTiles>, + options: StoreOptions<$algorithm, S>, + ) { + boundary::store::<_, I, ALIGN, C, $algorithm, S, G>(destination, self, options); + } + } + }; +} +shared_store_value!(Direct); +shared_store_value!(Striped); +shared_store_value!(Vectorize); +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, + ) { + match A::KIND { + sealed::DIRECT => { + boundary::store::<_, I, ALIGN, C, Direct, S, G>( + destination, + self, + options.algorithm(Direct), + ); + } + sealed::STRIPED => { + boundary::store::<_, I, ALIGN, C, Striped, S, G>( + destination, + self, + options.algorithm(Striped), + ); + } + sealed::VECTORIZE => { + boundary::store::<_, I, ALIGN, C, Vectorize, S, G>( + destination, + self, + options.algorithm(Vectorize), + ); + } + sealed::TRANSPOSE => { + boundary::store_permuting::<_, I, ALIGN, C, _, S, G>( + destination, + self, + options.algorithm(Transpose), + ); + } + sealed::WARP_TRANSPOSE => { + boundary::store_permuting::<_, I, ALIGN, C, _, S, G>( + destination, + self, + options.algorithm(WarpTranspose), + ); + } + sealed::WARP_TRANSPOSE_TIMESLICED => { + boundary::store_permuting::<_, I, ALIGN, C, _, S, G>( + destination, + self, + options.algorithm(WarpTransposeTimesliced), + ); + } + _ => panic!(), + } + } +} +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 Group { + const WARP: bool; + const WIDTH: usize; + fn group_index() -> u32; + unsafe fn call< + T: Copy, + R: Fn() -> &'static str + Copy, + M, + S, + U: Copy, + P: Copy, + const I: usize, + const WRITE: bool, + >( + registration: R, + memory: M, + storage: S, + uniform: U, + payload: P, + ); + } + // Scalar tags let the compiler prune unused native operation arms before + // planning scratch. Only the sealed selector implementations assign them. + pub const DIRECT: u8 = 0; + pub const STRIPED: u8 = 1; + pub const TRANSPOSE: u8 = 2; + pub const VECTORIZE: u8 = 3; + pub const WARP_TRANSPOSE: u8 = 4; + pub const WARP_TRANSPOSE_TIMESLICED: u8 = 5; + pub trait Algorithm { + const KIND: u8; + } +} diff --git a/rust/cuda-coop/src/load_store/boundary.rs b/rust/cuda-coop/src/load_store/boundary.rs new file mode 100644 index 000000000000..b99879d743dd --- /dev/null +++ b/rust/cuda-coop/src/load_store/boundary.rs @@ -0,0 +1,440 @@ +// Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +use super::{ + Algorithm, LoadOptions, StoreOptions, + sealed::{ + DIRECT, Group, STRIPED, TRANSPOSE, VECTORIZE, WARP_TRANSPOSE, WARP_TRANSPOSE_TIMESLICED, + }, +}; +use crate::{Alignment, Numeric, SupportedAlignment, TempStorageLike, 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!(striped_load_full, "COOP_STRIPED_VALIDITY_LOAD_FULL"); +registration!(striped_load_masked, "COOP_STRIPED_VALIDITY_LOAD_MASKED"); +registration!(striped_load_filled, "COOP_STRIPED_VALIDITY_LOAD_FILLED"); +registration!(striped_store_full, "COOP_STRIPED_VALIDITY_STORE_FULL"); +registration!(striped_store_partial, "COOP_STRIPED_VALIDITY_STORE_PARTIAL"); + +registration!(vectorize_load_full, "COOP_VECTORIZE_VALIDITY_LOAD_FULL"); +registration!(vectorize_load_masked, "COOP_VECTORIZE_VALIDITY_LOAD_MASKED"); +registration!(vectorize_load_filled, "COOP_VECTORIZE_VALIDITY_LOAD_FILLED"); +registration!(vectorize_store_full, "COOP_VECTORIZE_VALIDITY_STORE_FULL"); +registration!( + vectorize_store_partial, + "COOP_VECTORIZE_VALIDITY_STORE_PARTIAL" +); + +registration!(transpose_load_full, "COOP_TRANSPOSE_VALIDITY_LOAD_FULL"); +registration!(transpose_load_masked, "COOP_TRANSPOSE_VALIDITY_LOAD_MASKED"); +registration!(transpose_load_filled, "COOP_TRANSPOSE_VALIDITY_LOAD_FILLED"); +registration!(transpose_store_full, "COOP_TRANSPOSE_VALIDITY_STORE_FULL"); +registration!( + transpose_store_partial, + "COOP_TRANSPOSE_VALIDITY_STORE_PARTIAL" +); + +registration!( + warp_transpose_load_full, + "COOP_WARP_TRANSPOSE_VALIDITY_LOAD_FULL" +); +registration!( + warp_transpose_load_masked, + "COOP_WARP_TRANSPOSE_VALIDITY_LOAD_MASKED" +); +registration!( + warp_transpose_load_filled, + "COOP_WARP_TRANSPOSE_VALIDITY_LOAD_FILLED" +); +registration!( + warp_transpose_store_full, + "COOP_WARP_TRANSPOSE_VALIDITY_STORE_FULL" +); +registration!( + warp_transpose_store_partial, + "COOP_WARP_TRANSPOSE_VALIDITY_STORE_PARTIAL" +); +registration!( + warp_transpose_timesliced_load_full, + "COOP_WARP_TRANSPOSE_TIMESLICED_VALIDITY_LOAD_FULL" +); +registration!( + warp_transpose_timesliced_load_masked, + "COOP_WARP_TRANSPOSE_TIMESLICED_VALIDITY_LOAD_MASKED" +); +registration!( + warp_transpose_timesliced_load_filled, + "COOP_WARP_TRANSPOSE_TIMESLICED_VALIDITY_LOAD_FILLED" +); +registration!( + warp_transpose_timesliced_store_full, + "COOP_WARP_TRANSPOSE_TIMESLICED_VALIDITY_STORE_FULL" +); +registration!( + warp_transpose_timesliced_store_partial, + "COOP_WARP_TRANSPOSE_TIMESLICED_VALIDITY_STORE_PARTIAL" +); + +#[inline(always)] +pub(super) fn tile() -> usize { + let threads = thread::blockDim_x() as usize + * 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 <= group_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)); + } +} + +#[inline(always)] +pub(super) fn group_tile() -> usize { + if G::WARP { + assert!(ITEMS > 0 && ITEMS <= usize::MAX / G::WIDTH); + G::WIDTH * ITEMS + } else { + tile::() + } +} + +#[inline(always)] +pub(super) fn group_offset(offset: usize) -> usize { + if G::WARP { + let group = G::group_index() as usize; + let tile = group_tile::(); + assert!(group <= usize::MAX / tile); + let origin = group * tile; + assert!(offset <= usize::MAX - origin); + offset + origin + } else { + offset + } +} + +#[inline(always)] +pub(super) fn group_containment(offset: usize, valid: usize) { + if !G::WARP { + containment::(offset, valid); + return; + } + let block_capacity = tile::(); + let capacity = group_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 / block_capacity); + let period = groups * block_capacity; + if valid != 0 { + let within = if offset < period { + offset + } else { + offset % period + }; + let warp = G::group_index() as usize; + let begin = rank * block_capacity + warp * 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, + S: TempStorageLike, + G: Group, +>( + 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 => group_tile::(), + }; + let offset = group_offset::(options.offset); + extent::(len, 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(offset) }; + let payload = output.block_io_parts(); + let storage = options.storage.into_raw(); + // Each static algorithm keeps its registration function-item identity. The + // storage marker plans the complete provider layout and either synchronizes + // managed reuse or requires the caller's explicit completion. + macro_rules! call { + ($direct:ident, $striped:ident, $vectorize:ident, $transpose:ident, $warp:ident, $timesliced:ident, $controls:expr) => { + match A::KIND { + DIRECT => G::call::( + $direct, memory, storage, $controls, payload, + ), + STRIPED => G::call::( + $striped, memory, storage, $controls, payload, + ), + VECTORIZE => G::call::( + $vectorize, memory, storage, $controls, payload, + ), + TRANSPOSE => G::call::( + $transpose, memory, storage, $controls, payload, + ), + WARP_TRANSPOSE => G::call::( + $warp, memory, storage, $controls, payload, + ), + WARP_TRANSPOSE_TIMESLICED => G::call::( + $timesliced, + memory, + storage, + $controls, + payload, + ), + _ => panic!(), + } + }; + } + // 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 exact scratch and checks the selected completion policy. + unsafe { + match (options.valid_items, options.oob_default) { + (None, None) => call!( + load_full, + striped_load_full, + vectorize_load_full, + transpose_load_full, + warp_transpose_load_full, + warp_transpose_timesliced_load_full, + (pointer,) + ), + (Some(count), None) => { + call!( + load_masked, + striped_load_masked, + vectorize_load_masked, + transpose_load_masked, + warp_transpose_load_masked, + warp_transpose_timesliced_load_masked, + (pointer, count as i32) + ) + } + (Some(count), Some(default)) => call!( + load_filled, + striped_load_filled, + vectorize_load_filled, + transpose_load_filled, + warp_transpose_load_filled, + warp_transpose_timesliced_load_filled, + (pointer, count as i32, default) + ), + (None, Some(_)) => panic!(), + } + } +} + +#[inline(always)] +pub(super) fn store< + T: Numeric, + const ITEMS: usize, + const ALIGN: usize, + const CAPACITY: usize, + A: Algorithm, + S: TempStorageLike, + G: Group, +>( + destination: &mut DisjointSlice<'_, T, CollectiveTiles>, + value: &ThreadData, + options: StoreOptions, +) where + Alignment: SupportedAlignment, +{ + let valid = match options.valid_items { + Some(count) => count, + None => group_tile::(), + }; + let offset = group_offset::(options.offset); + extent::(destination.len(), offset, valid); + group_containment::(offset, valid); + let pointer = destination.as_mut_ptr(); + let payload = value.block_io_read_parts(); + let storage = options.storage.into_raw(); + // 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(offset); + macro_rules! call { + ($direct:ident, $striped:ident, $vectorize:ident, $controls:expr) => { + match A::KIND { + DIRECT => G::call::( + $direct, + destination, + storage, + $controls, + payload, + ), + STRIPED => G::call::( + $striped, + destination, + storage, + $controls, + payload, + ), + VECTORIZE => G::call::( + $vectorize, + destination, + storage, + $controls, + payload, + ), + // Only the shared-preserving Store implementations call this boundary. + _ => panic!(), + } + }; + } + match options.valid_items { + None => call!( + store_full, + striped_store_full, + vectorize_store_full, + (pointer,) + ), + Some(count) => call!( + store_partial, + striped_store_partial, + vectorize_store_partial, + (pointer, count as i32) + ), + } + } +} + +#[inline(always)] +pub(super) fn store_permuting< + T: Numeric, + const ITEMS: usize, + const ALIGN: usize, + const CAPACITY: usize, + A: Algorithm, + S: TempStorageLike, + G: Group, +>( + destination: &mut DisjointSlice<'_, T, CollectiveTiles>, + value: &mut ThreadData, + options: StoreOptions, +) where + Alignment: SupportedAlignment, +{ + let valid = match options.valid_items { + Some(count) => count, + None => group_tile::(), + }; + let offset = group_offset::(options.offset); + extent::(destination.len(), offset, valid); + group_containment::(offset, valid); + let pointer = destination.as_mut_ptr(); + let payload = value.block_io_parts(); + let storage = options.storage.into_raw(); + // 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(offset); + macro_rules! call { + ($transpose:ident, $warp:ident, $timesliced:ident, $controls:expr) => { + match A::KIND { + TRANSPOSE => G::call::( + $transpose, + destination, + storage, + $controls, + payload, + ), + WARP_TRANSPOSE => G::call::( + $warp, + destination, + storage, + $controls, + payload, + ), + WARP_TRANSPOSE_TIMESLICED => G::call::( + $timesliced, + destination, + storage, + $controls, + payload, + ), + _ => panic!(), + } + }; + } + match options.valid_items { + None => call!( + transpose_store_full, + warp_transpose_store_full, + warp_transpose_timesliced_store_full, + (pointer,) + ), + Some(count) => call!( + transpose_store_partial, + warp_transpose_store_partial, + warp_transpose_timesliced_store_partial, + (pointer, count as i32) + ), + } + } +} diff --git a/rust/cuda-coop/src/temp_storage.rs b/rust/cuda-coop/src/temp_storage.rs new file mode 100644 index 000000000000..6e35492539e5 --- /dev/null +++ b/rust/cuda-coop/src/temp_storage.rs @@ -0,0 +1,222 @@ +// Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +use crate::Numeric; +use cuda_device::library; + +/// Compiler-planned shared scratch for complete-block operations. +/// +/// Construct this descriptor inside a supported Oxide kernel and pass a mutable +/// borrow through [`crate::LoadOptions::temp_storage`] or +/// [`crate::StoreOptions::temp_storage`]. The default shares one slice between +/// uses and requires explicit block synchronization before reuse. Omitting the +/// descriptor instead requests compiler-managed reuse synchronization. +/// +/// `CAPACITY` and `ALIGN` default to zero, meaning inferred. Consuming builders +/// select positive explicit values and synchronization/sharing policies before +/// use. Capacity covers the entire plan, including exclusive slices and padding. +/// Each static constructor site denotes one descriptor, including in a loop. +/// Different sites have independent storage. Exclusive mode distinguishes +/// static use sites, not repeated executions of one site in a loop. +/// +/// This is neither a shared-memory reference nor a runtime allocation. All block +/// lanes must agree on the descriptor and reach matching operations; the compiler +/// checks participation, origin and reuse. Dropping it does not synchronize. +#[must_use] +pub struct TempStorage< + const CAPACITY: usize = 0, + const ALIGN: usize = 0, + const AUTO_SYNC: bool = false, + const EXCLUSIVE: bool = false, +> { + raw: library::StorageDescriptor, +} + +impl TempStorage { + /// Request inferred shared storage with automatic synchronization disabled. + #[inline(always)] + pub fn new() -> Self { + // SAFETY: the supported compiler validates uniform construction and + // retains the actual rooted constructor identity. + Self { + raw: unsafe { library::storage_new_v1::<0, 0, false, false>() }, + } + } +} + +impl Default for TempStorage { + fn default() -> Self { + Self::new() + } +} + +impl TempStorage { + const VALID_CAPACITY: () = assert!(C > 0 && C <= u32::MAX as usize); + const VALID_ALIGNMENT: () = assert!(A.is_power_of_two() && A <= u32::MAX as usize); + + /// Set a positive byte capacity for the complete descriptor plan. + /// + /// Undersized capacity, including exclusive-site padding, is a compilation + /// error. For a scratch-using plan, the compiler reserves the requested + /// capacity when it is sufficient. Scratch-free uses allocate nothing. + #[inline(always)] + pub fn with_capacity(self) -> TempStorage { + let () = TempStorage::::VALID_CAPACITY; + // SAFETY: both types are genuine descriptors; the compiler preserves + // this origin and rejects policy changes after use. + TempStorage { + raw: unsafe { library::storage_reconfigure_v1(self.raw) }, + } + } + + /// Set a positive power-of-two minimum alignment for the descriptor base. + /// + /// Each exclusive slice also satisfies its own operation/reservation + /// alignment. This base request does not change every slice's alignment. + #[inline(always)] + pub fn with_alignment(self) -> TempStorage { + let () = TempStorage::::VALID_ALIGNMENT; + // SAFETY: consuming genuine descriptors retains the compiler origin. + TempStorage { + raw: unsafe { library::storage_reconfigure_v1(self.raw) }, + } + } + + /// Insert a trailing block barrier after each scratch-using operation. + /// + /// Scratch-free operations such as Direct Load/Store gain no barrier. + /// Typed reservations are unavailable on an auto-sync descriptor. + #[inline(always)] + pub fn auto_sync(self) -> TempStorage { + // SAFETY: the compiler checks policy changes and actual scratch use. + TempStorage { + raw: unsafe { library::storage_reconfigure_v1(self.raw) }, + } + } + + /// Give distinct static primitive and reservation sites separate slices. + /// + /// Repeating one site still reuses its slice and requires completion. + /// Independent scratch does not complete conflicting global-memory accesses. + #[inline(always)] + pub fn exclusive(self) -> TempStorage { + // SAFETY: the compiler preserves the origin and plans distinct sites. + TempStorage { + raw: unsafe { library::storage_reconfigure_v1(self.raw) }, + } + } +} + +impl TempStorage { + /// Reserve `N` numeric elements with their natural alignment. + /// + /// `N` must be positive. The view is uninitialized and exposes only a raw + /// pointer; see [`StorageReservation`] for access and lifetime obligations. + /// Reservations follow this descriptor's shared/exclusive allocation policy. + #[inline(always)] + pub fn reserve(&mut self) -> StorageReservation { + let () = StorageReservation::::VALID_EXTENT; + // SAFETY: the compiler checks participation, origin, live overlap and + // layout. Returning a raw token performs no access to an element. + StorageReservation { + raw: unsafe { library::storage_reserve_v1::(&mut self.raw) }, + } + } + + /// Reserve `N` numeric elements with a power-of-two minimum alignment. + /// + /// `N` and `ALIGN` must be positive. The requested alignment applies to this + /// reservation, including an exclusive slice at a nonzero offset. + #[inline(always)] + pub fn reserve_aligned( + &mut self, + ) -> StorageReservation { + let () = StorageReservation::::VALID_EXTENT; + let () = TempStorage::::VALID_ALIGNMENT; + // SAFETY: the compiler validates the genuine descriptor and typed + // layout. Actual raw memory accesses retain their unsafe obligations. + StorageReservation { + raw: unsafe { library::storage_reserve_v1::(&mut self.raw) }, + } + } +} + +/// A numeric shared-memory view with explicitly unsafe element access. +/// +/// This non-Copy token grants no initialized value, Rust reference or slice. +/// It does not retain an exclusive Rust borrow of its descriptor, so separate +/// exclusive reservations may remain live together. The supported compiler +/// checks view origins and rejects overlapping planned reuse while a view is live. +/// It does not prove arbitrary unsafe pointer operations. +/// +/// The allocation belongs to the executing block for this kernel execution. +/// Every raw access must prove bounds, initialization before reads, aliasing, +/// non-conflicting accesses and phase synchronization. Complete asynchronous +/// accesses before ending their view. Saved pointers must not be dereferenced +/// after [`release`](Self::release), during overlapping reuse, or after kernel +/// exit. Copies of a pointer do not extend the view's lifetime. +/// +/// Dropping or forgetting the token does not release it or synchronize. +#[must_use = "release the reservation before overlapping storage reuse"] +pub struct StorageReservation { + raw: library::StorageReservation, +} + +impl StorageReservation { + const VALID_EXTENT: () = assert!(N > 0); + + /// Obtain the raw pointer without creating a reference or reading an element. + /// + /// Dereferencing it requires the obligations in [`StorageReservation`]. + #[inline(always)] + pub fn as_mut_ptr(&self) -> *mut T { + self.raw.as_mut_ptr() + } + + /// End the view without completing shared or global-memory accesses. + /// + /// Synchronize the full block **after** release and before overlapping + /// reuse. A barrier before release is insufficient. No saved raw pointer + /// may be accessed afterward; only unsafe code could perform such an access. + #[inline(always)] + pub fn release(self) { + // SAFETY: the compiler validates uniform release of a live view. The + // safe API grants no reference or access; saved raw pointers remain + // subject to the caller's unsafe lifetime and completion obligations. + unsafe { library::storage_release_v1(self.raw) } + } +} + +/// A sealed storage argument carried by Load/Store options. +/// +/// `()` selects compiler-managed scratch; a mutable [`TempStorage`] borrow +/// selects explicit storage. Operations consume options without requiring Copy. +pub trait TempStorageLike: sealed::Storage {} + +impl TempStorageLike for () {} +impl sealed::Storage for () { + type Raw = (); + fn into_raw(self) {} +} + +impl TempStorageLike + for &mut TempStorage +{ +} +impl<'a, const C: usize, const A: usize, const S: bool, const E: bool> sealed::Storage + for &'a mut TempStorage +{ + type Raw = &'a mut library::StorageDescriptor; + + #[inline(always)] + fn into_raw(self) -> Self::Raw { + &mut self.raw + } +} + +pub(crate) mod sealed { + pub trait Storage { + type Raw; + fn into_raw(self) -> Self::Raw; + } +} 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..3e71d7f4bf4a --- /dev/null +++ b/rust/cuda-coop/tests/check_load_store.py @@ -0,0 +1,902 @@ +# 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 barrier covering both accesses", + ), + "write_after_read": ( + reload + + "coop::store(block,&mut output,1_i32,coop::StoreOptions::default().offset(offset));", + "requires a barrier covering both accesses", + ), + "loop_without_completion": ( + load + "let mut iteration=0;while iteration<2 {" + store + "iteration+=1;}", + "requires a barrier covering both accesses", + ), + "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"), + "StoreValue", + ), + transpose_managed_reuse=( + transpose_load + transpose_store + transpose_store, + None, + ), + transpose_no_trailing_completion=( + transpose_load + transpose_store + reload, + "requires a barrier covering both accesses", + ), + 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, + ), +) + +# Public descriptors preserve the existing static-site and completion rules. +new_storage = "let mut scratch=coop::TempStorage::new();" +sync = "thread::sync_threads();" +explicit_load = transpose_load.replace( + ".algorithm(coop::Transpose)", + ".algorithm(coop::Transpose).temp_storage(&mut scratch)", +) +explicit_store = transpose_store.replace( + ".algorithm(coop::Transpose)", + ".algorithm(coop::Transpose).temp_storage(&mut scratch)", +) +reserve = "let view=scratch.reserve::();" +release = "view.release();" +cases.update( + storage_manual=(new_storage + explicit_load + sync + explicit_store, None), + storage_manual_missing_barrier=( + new_storage + explicit_load + explicit_store, + "requires a full-block barrier", + ), + storage_auto_loop=( + new_storage.replace("new()", "new().auto_sync()") + + "let mut iteration=0;while iteration<2 {" + + explicit_load + + explicit_store + + "iteration+=1;}", + None, + ), + storage_direct_auto_no_completion=( + new_storage.replace("new()", "new().auto_sync()") + + load + + store.replace( + ".valid_items(valid)", ".valid_items(valid).temp_storage(&mut scratch)" + ) + + reload, + "requires a barrier covering both accesses", + ), + storage_exclusive_sites=( + new_storage.replace("new()", "new().exclusive()") + + explicit_load + + explicit_load + + explicit_store, + None, + ), + storage_exclusive_loop=( + new_storage.replace("new()", "new().exclusive()") + + "let mut iteration=0;while iteration<2 {" + + explicit_load + + "iteration+=1;}", + "requires a full-block barrier", + ), + storage_reserve_auto=( + new_storage.replace("new()", "new().auto_sync()") + reserve + release, + "no method named `reserve`", + ), + storage_reserve_live=( + new_storage + reserve + sync + explicit_load + release, + "overlaps a live reservation", + ), + storage_release_no_barrier=( + new_storage + reserve + release + explicit_load, + "requires a full-block barrier", + ), + storage_barrier_before_release=( + new_storage + reserve + sync + release + explicit_load, + "requires a full-block barrier", + ), + storage_release_then_barrier=( + new_storage + reserve + release + sync + explicit_load, + None, + ), + storage_reserve_exclusive=( + new_storage.replace("new()", "new().exclusive()") + + reserve + + "let other=scratch.reserve_aligned::();" + + release + + "other.release();", + None, + ), + storage_extent_zero=( + new_storage + reserve.replace("i32,3", "i32,0") + release, + "assertion failed: N > 0", + ), + storage_numeric_only=( + new_storage + reserve.replace("i32,3", "bool,3") + release, + "Numeric", + ), + storage_alignment_invalid=( + "let _scratch=coop::TempStorage::new().with_alignment::<3>();", + "assertion failed", + ), + storage_capacity_zero=( + "let _scratch=coop::TempStorage::new().with_capacity::<0>();", + "assertion failed", + ), + storage_options_consumed=( + new_storage + + "let options=coop::LoadOptions::default().temp_storage(&mut scratch);coop::load(block,input,&mut values,options);coop::load(block,input,&mut values,options);", + "use of moved value: `options`", + ), + storage_policy_after_direct=( + new_storage + + load.replace(".oob_default(0)", ".oob_default(0).temp_storage(&mut scratch)") + + "let _changed=scratch.exclusive();", + "storage policy cannot change after use", + ), + storage_descriptor_mutation=( + "#[inline(never)] fn replace(old:&mut coop::TempStorage,new:coop::TempStorage){*old=new;}" + + new_storage + + "let other=coop::TempStorage::new();replace(&mut scratch,other);", + "mutation through a reference to a resource descriptor", + ), + storage_reservation_mutation=( + "#[inline(never)] fn replace(old:&mut coop::StorageReservation,new:coop::StorageReservation){*old=new;}" + + new_storage + + "let mut view=scratch.reserve::();let mut other=coop::TempStorage::new();let replacement=other.reserve::();replace(&mut view,replacement);view.release();", + "mutation through a reference to a resource descriptor", + ), +) +for capacity in (55, 56): + cases[f"storage_complete_capacity_{capacity}"] = ( + f"let mut scratch=coop::TempStorage::new().with_capacity::<{capacity}>().exclusive();" + + reserve + + "let other=scratch.reserve_aligned::();" + + release + + "other.release();", + "smaller than complete plan" if capacity == 55 else None, + ) + +# Storage-only cases still infer the payload type from a typed producer. +for name, (body, diagnostic) in list(cases.items()): + if name.startswith("storage_"): + cases[name] = ("values.write(0, 0_i32);" + body, diagnostic) + +# Physical Warp operations use the same public calls and canonical carrier. +# Query precision is granted only to the genuine compiler query, not to +# arithmetic that happens to divide the lane rank by 32. +warp_prefix = ( + prefix.replace("block=(128,1,1)", "block=(64,1,1)").replace("*512", "*256") + + "let warp=coop::this_warp();let q=cuda_device::library::warp_group_index_v1::<32>() as usize;\n" +) +warp_load = load.replace("(block,", "(warp,") +warp_store = store.replace("(block,", "(warp,") +warp_reload = reload.replace("(block,", "(warp,") +warp_helper = """#[inline(never)] fn put(group:coop::ThreadGroup, + output:&mut DisjointSlice>,offset:usize) { + coop::store(group,output,7_i32,coop::StoreOptions::default().offset(offset)); } +""" +warp_loop = "let mut i=0;while i() as usize;let mut values=coop::ThreadData::<_,4>::new();let warp=coop::this_warp();let valid=1;let offset=0;" + + warp_loop + + "}cycles(input);" + + load, + None, + ), + warp_partial_geometry=(warp_load, "complete physical warps"), + warp_transpose=( + warp_load.replace( + ".oob_default(0)", ".oob_default(0).algorithm(coop::Transpose)" + ), + None, + ), + warp_explicit_storage=( + new_storage + + warp_load.replace( + ".oob_default(0)", ".oob_default(0).temp_storage(&mut scratch)" + ), + "LoadStoreGroup", + ), +) + +striped_load = load.replace( + ".oob_default(0)", ".oob_default(0).algorithm(coop::Striped)" +) +striped_store = store.replace( + ".valid_items(valid)", ".valid_items(valid).algorithm(coop::Striped)" +) +cases.update( + striped_shared=(striped_load + striped_store, None), + striped_varying_count=( + striped_load.replace( + ".valid_items(valid)", ".valid_items(thread::threadIdx_x() as usize)" + ), + "not proved block-uniform", + ), + striped_missing_completion=( + striped_load + striped_store + reload, + "requires a barrier covering both accesses", + ), +) + + +vectorize_load = load.replace( + ".oob_default(0)", ".oob_default(0).algorithm(coop::Vectorize)" +) +vectorize_store = store.replace( + ".valid_items(valid)", ".valid_items(valid).algorithm(coop::Vectorize)" +) +cases.update( + vectorize_shared=(vectorize_load + vectorize_store, None), + vectorize_varying_offset=( + vectorize_load.replace( + ".offset(offset)", ".offset(offset+thread::threadIdx_x() as usize)" + ), + "not proved block-uniform", + ), + vectorize_missing_completion=( + vectorize_load + vectorize_store + reload, + "requires a barrier covering both accesses", + ), +) + +# Both algorithms participate as complete blocks and mutate the payload. +for name, selector in [ + ("warp_transpose", "WarpTranspose"), + ("warp_timesliced", "WarpTransposeTimesliced"), +]: + warp_load = transpose_load.replace("coop::Transpose", f"coop::{selector}") + warp_store = transpose_store.replace("coop::Transpose", f"coop::{selector}") + cases[name + "_shared"] = ( + load + warp_store.replace("&mut values", "&values"), + f"StoreValue", + ) + cases[name + "_managed_reuse"] = (warp_load + warp_store + warp_store, None) + cases[name + "_scalar"] = ( + warp_load + + f"coop::store(block,&mut output,-17_i32,coop::StoreOptions::default().offset(offset).algorithm(coop::{selector}));", + None, + ) + cases[name + "_partial_block"] = (warp_load, "divisible by 32") + cases[name + "_warp_only_participation"] = ( + "if thread::threadIdx_x()<32 {" + warp_load + "}", + "lane-varying", + ) + +# Logical group precision and completion use canonical public descriptors. +logical_prefix = warp_prefix.replace( + "let warp=coop::this_warp();", "let warp=coop::this_warp().group_by::<8>();" +).replace( + "cuda_device::library::warp_group_index_v1::<32>() as usize", + "warp.rank_at::() as usize", +) +logical_load = load.replace("(block,", "(warp,") +logical_store = store.replace("(block,", "(warp,") +logical_reload = reload.replace("(block,", "(warp,") +logical_sync = "use cuda_device::cooperative_groups::ThreadGroup as _;cuda_device::cooperative_groups::this_thread_block().tiled_partition::<8>().sync();" +wide_sync = logical_sync.replace("::<8>()", "::<32>()") +narrow_sync = logical_sync.replace("::<8>()", "::<4>()") +logical_loop = "let mut i=0;while i();let q=warp.rank_at::() as usize;let valid=1;let offset=0;let mut values=coop::ThreadData::<_,4>::new();" + + logical_loop + + "}cycle(input);" + + load, + None, + ), + logical_lane_control=( + logical_load.replace( + ".valid_items(valid)", ".valid_items(thread::threadIdx_x() as usize)" + ), + "not proved Warp-uniform", + ), + logical_public_rank_control=( + logical_load.replace(".valid_items(valid)", ".valid_items(q+1)"), + None, + ), + logical_group6=( + logical_load.replace( + "(warp,", "(coop::this_warp().group_by_non_exhaustive::<6>()," + ), + "LoadStoreGroup", + ), + logical_width1_origin=( + "let chosen=if thread::threadIdx_x()&1==0 {&mut output}else{&mut other};" + + logical_store.replace( + "(warp,&mut output", "(coop::this_warp().group_by::<1>(),chosen" + ), + "not proved Warp-uniform", + ), + logical_structural_width6=( + 'values.write(0,0_i32);let mut local=[0_i32;4];let flags=[1_u8;4];fn registration()->&\'static str {"unused"} unsafe {cuda_device::library::checked_warp_call_v1::(registration,&mut output,(core::ptr::null_mut::(),),(local.as_mut_ptr(),flags.as_ptr()));}', + "power-of-two width in 1..=32", + ), +) +# Finite sibling cycles reconverge structurally; value precision and memory +# completion remain independent requirements for a subsequent wider call. +cycle = logical_loop.replace("i()," + ), + None, + ) + + +# Replicated implicit Warp pools use actual provider extents and scoped completion. +logical_transpose_load = logical_load.replace( + ".oob_default(0)", ".oob_default(0).algorithm(coop::Transpose)" +) +logical_transpose_store = logical_store.replace("&values", "&mut values").replace( + ".valid_items(valid)", ".valid_items(valid).algorithm(coop::Transpose)" +) +logical_transpose_reload = logical_reload.replace( + ".valid_items(valid)", ".valid_items(valid).algorithm(coop::Transpose)" +) +warp16 = "coop::this_warp().group_by::<16>()" +logical_fixed_view = "let prefix=unsafe {cuda_device::DynamicSharedArray::<[u64;8],32>::get_fixed()};let _=prefix;" +cases.update( + logical_transpose_managed=( + logical_transpose_load + + logical_transpose_store + + logical_transpose_store + + logical_transpose_reload, + None, + ), + logical_transpose_shared=( + logical_transpose_load + + logical_transpose_store.replace("&mut values", "&values"), + "StoreValue", + ), + logical_transpose_sibling_count=( + logical_transpose_load.replace(".valid_items(valid)", ".valid_items(q)"), + None, + ), + logical_transpose_lane_count=( + logical_transpose_load.replace( + ".valid_items(valid)", ".valid_items(thread::threadIdx_x() as usize)" + ), + "not proved Warp-uniform", + ), + logical_transpose_block_pool=( + transpose_load + + logical_transpose_load + + logical_transpose_load + + transpose_load, + None, + ), + logical_transpose_containing_sync=( + logical_transpose_load + + logical_transpose_load.replace("(warp,", f"({warp16},") + + logical_sync.replace("::<8>()", "::<16>()") + + logical_transpose_load, + None, + ), + logical_transpose_wider_clean=( + logical_transpose_load + + logical_transpose_store + + logical_transpose_reload.replace("(warp,", f"({warp16},"), + "requires a barrier covering both accesses", + ), + logical_transpose_wider_completed=( + logical_transpose_load + + logical_transpose_store + + logical_sync.replace("::<8>()", "::<16>()") + + logical_transpose_reload.replace("(warp,", f"({warp16},"), + None, + ), + logical_transpose_block_completion=( + logical_transpose_load + logical_transpose_store + reload, + "requires a barrier covering both accesses", + ), + logical_transpose_live_reservation=( + new_storage + + reserve + + logical_transpose_load + + logical_sync + + explicit_load + + release, + "overlaps a live reservation", + ), + logical_transpose_explicit_dirty=( + new_storage + + explicit_load + + logical_transpose_load + + logical_transpose_load + + explicit_load, + "requires a full-block barrier", + ), + logical_transpose_loop=( + "let mut i=0;while i<2 {" + + logical_transpose_load + + logical_transpose_store + + "i+=1;}" + + logical_transpose_reload, + None, + ), + logical_shared_getter_lookalike=( + "struct DynamicSharedArray;impl DynamicSharedArray {#[inline(never)] fn get_fixed()->*mut u8 {if thread::threadIdx_x()&1==0 {thread::sync_threads();}core::ptr::null_mut()}}let _=DynamicSharedArray::get_fixed();" + + logical_transpose_load, + "block barrier is reached through lane-varying control flow", + ), + logical_transpose_fixed_prefix=(logical_fixed_view + logical_transpose_load, None), + logical_transpose_raw_view=( + logical_fixed_view.replace("get_fixed()", "get_raw()") + logical_transpose_load, + "planned storage cannot include an unconstrained native dynamic shared view", + ), + logical_zero_raw_view=( + logical_fixed_view.replace("get_fixed()", "get_raw()") + logical_load, + "planned storage cannot include an unconstrained native dynamic shared view", + ), + block_zero_raw_numeric=( + "let raw=unsafe {cuda_device::DynamicSharedArray::<[i32;128],16>::get_raw()};unsafe {raw.add(thread::threadIdx_x() as usize).write(1);}" + + load, + "planned storage cannot include an unconstrained native dynamic shared view", + ), + logical_zero_no_completion=( + logical_load + logical_store + logical_reload, + "requires a barrier covering both accesses", + ), +) + + +# Scratch-free Warp choices preserve shared payloads without completion credit. +for selector in ("Striped", "Vectorize"): + selected_load = logical_load.replace( + ".oob_default(0)", f".oob_default(0).algorithm(coop::{selector})" + ) + selected_store = logical_store.replace( + ".valid_items(valid)", f".valid_items(valid).algorithm(coop::{selector})" + ) + key = "logical_" + selector.lower() + cases[key + "_shared"] = (selected_load + selected_store, None) + cases[key + "_non_exhaustive"] = ( + selected_load.replace( + "(warp,", "(coop::this_warp().group_by_non_exhaustive::<8>()," + ) + + selected_store.replace( + "(warp,", "(coop::this_warp().group_by_non_exhaustive::<8>()," + ), + None, + ) + cases[key + "_physical"] = ( + selected_load.replace("(warp,", "(coop::this_warp(),") + + selected_store.replace("(warp,", "(coop::this_warp(),"), + None, + ) + cases[key + "_missing_completion"] = ( + selected_load + selected_store + logical_reload, + "requires a barrier covering both accesses", + ) + cases[key + "_explicit_storage"] = ( + new_storage + + selected_load.replace( + ".oob_default(0)", ".oob_default(0).temp_storage(&mut scratch)" + ), + f"LoadStoreGroup<{selector}, &mut TempStorage>", + ) + + +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(): + block_warp = name.startswith(("warp_transpose_", "warp_timesliced_")) + case_prefix = warp_prefix if name.startswith("warp_") and not block_warp else prefix + if name.startswith("logical_"): + case_prefix = logical_prefix + if name == "logical_width1_origin": + case_prefix = case_prefix.replace( + "valid:usize)", + "valid:usize,mut other:DisjointSlice>)", + ) + case_prefix = case_prefix.replace("pub fn probe(", "pub fn probe<'a>(").replace( + "DisjointSlice 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..514e255e8636 --- /dev/null +++ b/rust/cuda-coop/tests/load_store/Cargo.lock @@ -0,0 +1,1011 @@ +# 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=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "libnvvm-sys", + "nvjitlink-sys", + "serde", + "sha2", + "thiserror 2.0.21", +] + +[[package]] +name = "cuda-async" +version = "0.5.0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +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.5.0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +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.5.0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "anyhow", + "cuda-bindings", + "cuda-core-derive", + "half", + "oxide-artifacts", +] + +[[package]] +name = "cuda-core-derive" +version = "0.5.0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +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=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "cuda-macros", +] + +[[package]] +name = "cuda-host" +version = "0.2.1" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "cuda-artifact-finalizer", + "cuda-async", + "cuda-core", + "cuda-device", + "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=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +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=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" + +[[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=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +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=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +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=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "object", +] + +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +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=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" + +[[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=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" + +[[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..626172bef415 --- /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 = "96caed46c776eb49226137eac849e795b83b7ce0" } +cuda-host = { git = "https://github.com/tpn/cuda-rust.git", rev = "96caed46c776eb49226137eac849e795b83b7ce0", features = ["async"] } +cuda-core = { git = "https://github.com/tpn/cuda-rust.git", rev = "96caed46c776eb49226137eac849e795b83b7ce0" } +cuda-async = { git = "https://github.com/tpn/cuda-rust.git", rev = "96caed46c776eb49226137eac849e795b83b7ce0" } diff --git a/rust/cuda-coop/tests/load_store/src/block_warp.rs b/rust/cuda-coop/tests/load_store/src/block_warp.rs new file mode 100644 index 000000000000..42e058a840b5 --- /dev/null +++ b/rust/cuda-coop/tests/load_store/src/block_warp.rs @@ -0,0 +1,366 @@ +// Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +//! Full-block warp exchange observes every payload slot before and after Store. +use cuda_coop as coop; +use cuda_device::{CollectiveTiles, DisjointSlice, kernel, launch_contract, thread}; +use cuda_host::cuda_module; + +#[cuda_module] +mod kernels { + use super::*; + + fn expected( + input: &[T], + offset: usize, + index: usize, + items: usize, + valid: usize, + default: T, + mode: u32, + ) -> Option { + if mode == 1 || mode == 5 || mode == 6 { + if index < valid { + Some(input[offset + index]) + } else { + None + } + } else if mode == 2 || mode == 4 { + if index < valid { + Some(input[offset + index]) + } else { + Some(default) + } + } else if mode == 3 { + let lane = index / items; + let item = index % items; + if index < valid || (lane & 1 == 0 && item & 1 == 1) || (lane & 1 == 1 && item & 1 == 0) + { + Some(input[offset + index]) + } else { + None + } + } else { + Some(input[offset + index]) + } + } + + #[inline(never)] + fn exercise< + T: coop::Numeric + PartialEq, + const ITEMS: usize, + const CAPACITY: usize, + A: coop::Algorithm, + >( + input: &[T], + output: &mut DisjointSlice>, + valid: usize, + default: T, + mode: u32, + algorithm: A, + ) { + let block = coop::this_block(); + let x = thread::blockDim_x() as usize; + let y = thread::blockDim_y() as usize; + let threads = x * y * thread::blockDim_z() as usize; + let lane = thread::threadIdx_x() as usize + + x * (thread::threadIdx_y() as usize + y * thread::threadIdx_z() as usize); + let group = thread::blockIdx_x() as usize + + thread::gridDim_x() as usize + * (thread::blockIdx_y() as usize + + thread::gridDim_y() as usize * thread::blockIdx_z() as usize); + let offset = group * threads * CAPACITY; + let mut values = coop::ThreadData::<_, ITEMS>::new(); + if mode != 9 { + coop::load( + block, + input, + &mut values, + coop::LoadOptions::default() + .offset(offset) + .algorithm(algorithm), + ); + } + if mode == 1 || mode == 5 || mode == 6 { + coop::load( + block, + input, + &mut values, + coop::LoadOptions::default() + .offset(offset) + .valid_items(valid) + .algorithm(algorithm), + ); + } else if mode == 2 || mode == 4 { + coop::load( + block, + input, + &mut values, + coop::LoadOptions::default() + .offset(offset) + .valid_items(valid) + .oob_default(default) + .algorithm(algorithm), + ); + } else if mode == 3 { + let incoming = values; + values = coop::ThreadData::new(); + let mut item = 0; + while item < ITEMS { + let index = lane * ITEMS + item; + let initialize = index < valid + || (lane & 1 == 0 && item & 1 == 1) + || (lane & 1 == 1 && item & 1 == 0); + if let (true, Some(value)) = (initialize, incoming.get(item)) { + values.write(item, value); + } + item += 1; + } + } + let mut correct = true; + let mut item = 0; + while item < ITEMS { + correct &= values.get(item) + == expected( + input, + offset, + lane * ITEMS + item, + ITEMS, + valid, + default, + mode, + ); + item += 1; + } + if mode == 0 || mode == 2 { + coop::store( + block, + &mut *output, + &mut values, + coop::StoreOptions::default() + .offset(offset) + .algorithm(algorithm), + ); + } else { + let written = if mode == 5 { + valid + 1 + } else if mode == 6 { + valid / 2 + } else { + valid + }; + coop::store( + block, + &mut *output, + &mut values, + coop::StoreOptions::default() + .offset(offset) + .valid_items(written) + .algorithm(algorithm), + ); + } + item = 0; + while item < ITEMS { + let source = (lane / 32) * (32 * ITEMS) + 32 * item + lane % 32; + correct &= + values.get(item) == expected(input, offset, source, ITEMS, valid, default, mode); + item += 1; + } + // Managed scratch reuse inserts completion before this second operation. + // A zero count still permutes all values and bits, with no global writes. + coop::store( + block, + &mut *output, + &mut values, + coop::StoreOptions::default() + .offset(offset) + .valid_items(0) + .algorithm(algorithm), + ); + item = 0; + while item < ITEMS { + let first = (lane / 32) * (32 * ITEMS) + 32 * item + lane % 32; + let source = + ((first / ITEMS) / 32) * (32 * ITEMS) + 32 * (first % ITEMS) + (first / ITEMS) % 32; + correct &= + values.get(item) == expected(input, offset, source, ITEMS, valid, default, mode); + item += 1; + } + if mode == 4 { + thread::sync_threads(); + coop::load( + block, + &mut *output, + &mut values, + coop::LoadOptions::default() + .offset(offset) + .valid_items(valid) + .oob_default(default) + .algorithm(algorithm), + ); + coop::store( + block, + &mut *output, + &mut values, + coop::StoreOptions::default() + .offset(offset) + .valid_items(valid) + .algorithm(algorithm), + ); + item = 0; + while item < ITEMS { + correct &= values.get(item) + == expected( + input, + offset, + (lane / 32) * (32 * ITEMS) + 32 * item + lane % 32, + ITEMS, + valid, + default, + mode, + ); + item += 1; + } + } + assert!(correct); + } + #[kernel] + #[launch_contract(domain=3,block=(8,4,2),dynamic_shared=inferred)] + pub fn block_warp( + input: &[i32], + mut output: DisjointSlice>, + valid: usize, + mode: u32, + ) { + exercise::<_, ITEMS_PER_THREAD, 4, _>( + input, + &mut output, + valid, + -17, + mode, + coop::WarpTranspose, + ); + } + #[kernel] + #[launch_contract(domain=3,block=(8,4,2),dynamic_shared=inferred)] + pub fn block_warp_timesliced( + input: &[i32], + mut output: DisjointSlice>, + valid: usize, + mode: u32, + ) { + exercise::<_, ITEMS_PER_THREAD, 4, _>( + input, + &mut output, + valid, + -17, + mode, + coop::WarpTransposeTimesliced, + ); + } +} + +pub fn run(trap: Option<&str>) -> Result<(), Box> { + use cuda_async::simt::device_context::{init_device_contexts, with_cuda_context}; + use cuda_async::simt::device_operation::DeviceOperation; + use cuda_core::{DeviceBuffer, LaunchConfig3D}; + init_device_contexts(0, 1)?; + let context = with_cuda_context(0, |context| context.clone())?; + let stream = context.default_stream(); + // SAFETY: the typed declarations and device kernels are compiled together. + let module = unsafe { kernels::load_async(0)? }; + if let Some(trap) = trap { + let plan = module.prepare_block_warp_timesliced::<4>(LaunchConfig3D::auto_shared( + (1, 1, 1), + (8, 4, 2), + ))?; + let source = DeviceBuffer::from_host(&stream, &vec![11_i32; 256])?; + let target = DeviceBuffer::from_host(&stream, &vec![0_i32; 256])?; + let (valid, mode) = match trap { + "block-warp-count" => (257, 1), + "block-warp-unwritten" => (256, 9), + "block-warp-different-count" => (31, 5), + _ => panic!("unknown block warp trap"), + }; + let _ = module + .block_warp_timesliced_async_owned::<4, _, _>(&plan, source, target, valid, mode) + .sync()?; + panic!("expected block warp trap returned successfully"); + } + let mut cases = 0; + macro_rules! run { + ($prepare:ident,$launch:ident,$items:literal,$scratch:literal) => {{ + let plan = + module.$prepare::<$items>(LaunchConfig3D::auto_shared((2, 1, 1), (8, 4, 2)))?; + assert_eq!(plan.__raw_config().shared_mem_bytes, $scratch); + assert!( + module + .$prepare::<$items>(LaunchConfig3D::auto_shared((2, 1, 1), (64, 1, 1))) + .is_err() + ); + for mode in [0, 1, 2, 3, 4, 6] { + for valid in [0, 1, 32 * $items + 1, 64 * $items - 1, 64 * $items] { + if mode == 0 && valid != 64 * $items { + continue; + } + let input: Vec = (0..512).map(|i| i + 19).collect(); + let source = DeviceBuffer::from_host(&stream, &input)?; + let target = DeviceBuffer::from_host(&stream, &vec![0_i32; 512])?; + let (_, target) = module + .$launch::<$items, _, _>(&plan, source, target, valid, mode) + .sync()?; + let mut result = vec![0_i32; 512]; + target.copy_to_host(&stream, &mut result)?; + stream.synchronize()?; + for (index, &value) in result.iter().enumerate() { + let within = index % 256; + let written = if mode == 0 || mode == 2 { + 64 * $items + } else if mode == 6 { + valid / 2 + } else { + valid + }; + let expected = if within < written { + if mode == 2 && within >= valid { + -17 + } else { + input[index] + } + } else { + 0 + }; + assert_eq!( + value, + expected, + "{} items={} mode={} valid={} index={}", + stringify!($launch), + $items, + mode, + valid, + index + ); + } + cases += 1; + } + } + }}; + } + // CUB maximum plus the independent counted-Store bitmap plane, alignment16. + run!(prepare_block_warp, block_warp_async_owned, 1, 272); + run!(prepare_block_warp, block_warp_async_owned, 4, 1072); + run!( + prepare_block_warp_timesliced, + block_warp_timesliced_async_owned, + 1, + 144 + ); + run!( + prepare_block_warp_timesliced, + block_warp_timesliced_async_owned, + 4, + 560 + ); + println!("PASS {cases} public block warp Load/Store cases"); + Ok(()) +} diff --git a/rust/cuda-coop/tests/load_store/src/logical_plain.rs b/rust/cuda-coop/tests/load_store/src/logical_plain.rs new file mode 100644 index 000000000000..cf3d30497db4 --- /dev/null +++ b/rust/cuda-coop/tests/load_store/src/logical_plain.rs @@ -0,0 +1,228 @@ +// Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +//! Scratch-free Warp algorithms, adjusted pointers, and shared Store preservation. +use cuda_coop as coop; +use cuda_device::{cooperative_groups as cg, kernel, launch_contract, thread}; +use cuda_host::cuda_module; + +#[cuda_module] +mod kernels { + use super::*; + use cg::ThreadGroup as _; + + #[inline(never)] + #[allow(clippy::too_many_arguments)] + fn exercise< + T: coop::Numeric + PartialEq + From + core::ops::Add, + A: coop::Algorithm + Default, + const W: usize, + const I: usize, + const N: u32, + >( + input: &[T], + output: &mut coop::DisjointSlice>, + valid: usize, + shift: usize, + mode: u32, + striped: bool, + ) where + coop::ThreadsWithinWarp: coop::LoadStoreGroup, + for<'a> &'a coop::ThreadData: coop::StoreValue, + { + let group = coop::this_warp().group_by::(); + let tile = cg::this_thread_block().tiled_partition::(); + let lane = group.rank() as usize; + let q = coop::this_block().rank_at::() as usize * (32 / W) + + group.rank_at::() as usize; + let offset = thread::blockIdx_x() as usize * 512 + q * W * (8 - I) + shift; + let count = if mode == 0 { W * I } else { valid }; + let mut values = coop::ThreadData::<_, I>::new(); + if mode == 4 { + let mut i = 0; + while i < I { + let position = if striped { lane + W * i } else { lane * I + i }; + if position < count || (lane + i) % 3 == 1 { + values.write(i, T::from(position as i8) + T::from(17)); + } + i += 1; + } + } else { + coop::load( + group, + input, + &mut values, + coop::LoadOptions::default() + .algorithm(A::default()) + .offset(offset), + ); + if mode != 0 { + let options = coop::LoadOptions::default() + .algorithm(A::default()) + .offset(offset) + .valid_items(count); + if mode == 2 { + coop::load(group, input, &mut values, options.oob_default(T::from(-7))); + } else { + coop::load(group, input, &mut values, options); + } + } + } + let before = values; + let options = coop::StoreOptions::default() + .algorithm(A::default()) + .offset(offset); + if mode == 0 { + coop::store(group, &mut *output, &values, options); + } else { + coop::store(group, &mut *output, &values, options.valid_items(count)); + } + let after = values; + tile.sync(); + coop::store( + group, + &mut *output, + &values, + coop::StoreOptions::default() + .algorithm(A::default()) + .offset(offset) + .valid_items(0), + ); + let twice = values; + tile.sync(); + coop::load( + group, + output, + &mut values, + coop::LoadOptions::default() + .algorithm(A::default()) + .offset(offset) + .valid_items(count), + ); + let mut matches = true; + let mut i = 0; + while i < I { + let position = if striped { lane + W * i } else { lane * I + i }; + let initialized = if mode == 4 { + position < count || (lane + i) % 3 == 1 + } else { + mode == 0 || mode == 2 || position < count + }; + let expected = if initialized { + Some(if mode == 4 { + T::from(position as i8) + T::from(17) + } else if mode == 2 && position >= count { + T::from(-7) + } else { + input[offset + q * W * I + position] + }) + } else { + None + }; + matches &= + before.get(i) == expected && after.get(i) == expected && twice.get(i) == expected; + matches &= values.get(i) == if position < count { expected } else { None }; + i += 1; + } + assert!(matches); + } + + #[kernel] + #[launch_contract(domain=1,block=(64,1,1),dynamic_shared=inferred)] + pub fn logical_plain< + T: coop::Numeric + PartialEq + From + core::ops::Add, + A: coop::Algorithm + Default, + const W: usize, + const I: usize, + const N: u32, + >( + input: &[T], + mut output: coop::DisjointSlice>, + valid: usize, + shift: usize, + mode: u32, + striped: bool, + ) where + coop::ThreadsWithinWarp: coop::LoadStoreGroup, + for<'a> &'a coop::ThreadData: coop::StoreValue, + { + exercise::(input, &mut output, valid, shift, mode, striped); + } +} + +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}; + init_device_contexts(0, 1)?; + let stream = with_cuda_context(0, |context| context.default_stream())?; + // SAFETY: the typed module and launchers compile together. + let module = unsafe { kernels::load_async(0)? }; + let mut cases = 0; + macro_rules! exercise { + ($ty:ty,$algorithm:ty,$striped:literal,$w:literal,$i:literal,$controls:expr) => {{ + let input: Vec<$ty> = (0..1536).map(|i| (i + 31) as $ty).collect(); + let plan = module.prepare_logical_plain::<$ty, $algorithm, $w, $i, $w>( + LaunchConfig1D::auto_shared(3, 64), + )?; + assert_eq!(plan.__raw_config().shared_mem_bytes, 0); + for shift in [0, 1] { + for (mode, count) in $controls { + let source = DeviceBuffer::from_host(&stream, &input)?; + let target = DeviceBuffer::from_host(&stream, &vec![-99 as $ty; input.len()])?; + let (_, target) = module + .logical_plain_async_owned::<$ty, $algorithm, $w, $i, $w, _, _>( + &plan, source, target, count, shift, mode, $striped, + ) + .sync()?; + let mut result = vec![0 as $ty; input.len()]; + target.copy_to_host(&stream, &mut result)?; + stream.synchronize()?; + let mut expected = vec![-99 as $ty; input.len()]; + for block in 0..3 { + for q in 0..64 / $w { + let valid = if mode == 0 { $w * $i } else { count }; + for i in 0..valid { + let index = block * 512 + q * $w * 8 + shift + i; + expected[index] = if mode == 4 { + i as $ty + 17 as $ty + } else { + input[index] + }; + } + } + } + assert_eq!(result, expected); + cases += 1; + } + } + }}; + } + let controls = |width: usize, items: usize| { + let mut counts = vec![0, 1, width * items - 1, width * items]; + counts.sort_unstable(); + counts.dedup(); + [0, 1, 2, 4] + .into_iter() + .flat_map(|mode| { + counts.iter().copied().filter_map(move |count| { + (mode != 0 || count == width * items).then_some((mode, count)) + }) + }) + .collect::>() + }; + exercise!(i32, coop::Striped, true, 1, 1, controls(1, 1)); + exercise!(i32, coop::Striped, true, 8, 4, controls(8, 4)); + exercise!(i32, coop::Striped, true, 32, 4, controls(32, 4)); + exercise!(i32, coop::Vectorize, false, 1, 1, controls(1, 1)); + exercise!(i32, coop::Vectorize, false, 8, 4, controls(8, 4)); + exercise!(i32, coop::Vectorize, false, 32, 4, controls(32, 4)); + exercise!(i8, coop::Striped, true, 8, 4, [(0, 32), (1, 31)]); + exercise!(i8, coop::Vectorize, false, 8, 4, [(0, 32), (1, 31)]); + exercise!(f64, coop::Striped, true, 8, 4, [(0, 32), (1, 31)]); + exercise!(f64, coop::Vectorize, false, 8, 4, [(0, 32), (1, 31)]); + println!( + "PASS {cases} Warp Striped/Vectorize cases: shared Store preservation, offsets0/1, zero scratch" + ); + Ok(()) +} diff --git a/rust/cuda-coop/tests/load_store/src/logical_transpose.rs b/rust/cuda-coop/tests/load_store/src/logical_transpose.rs new file mode 100644 index 000000000000..037324ec2d03 --- /dev/null +++ b/rust/cuda-coop/tests/load_store/src/logical_transpose.rs @@ -0,0 +1,513 @@ +// Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +//! Replicated implicit Warp scratch and observed Transpose payload mutation. +use cuda_coop as coop; +use cuda_device::{kernel, launch_contract, thread}; +use cuda_host::cuda_module; + +#[cuda_module] +mod kernels { + use super::*; + + #[inline(never)] + #[allow(clippy::too_many_arguments)] + fn exercise, const I: usize>( + group: coop::ThreadGroup, + input: &[i32], + output: &mut coop::DisjointSlice>, + width: usize, + q: usize, + lane: usize, + offset: usize, + valid: usize, + mode: u32, + ) { + let count = if mode == 0 { + width * I + } else if mode == 3 { + (valid + q).min(width * I) + } else { + valid + }; + let mut values = coop::ThreadData::<_, I>::new(); + if mode == 4 { + let mut i = 0; + while i < I { + if lane * I + i < count || (lane + i) % 3 == 1 { + values.write(i, (lane * I + i) as i32 + 17); + } + i += 1; + } + } else { + coop::load( + group, + input, + &mut values, + coop::LoadOptions::default() + .algorithm(coop::Transpose) + .offset(offset), + ); + if mode != 0 { + let options = coop::LoadOptions::default() + .algorithm(coop::Transpose) + .offset(offset) + .valid_items(count); + if mode == 2 { + coop::load(group, input, &mut values, options.oob_default(-7)); + } else { + coop::load(group, input, &mut values, options); + } + } + } + let before = values; + coop::store( + group, + &mut *output, + &mut values, + coop::StoreOptions::default() + .algorithm(coop::Transpose) + .offset(offset) + .valid_items(count), + ); + let after = values; + coop::store( + group, + &mut *output, + &mut values, + coop::StoreOptions::default() + .algorithm(coop::Transpose) + .offset(offset) + .valid_items(0), + ); + let twice = values; + // Managed scratch reuse supplies matching completion before this read. + coop::load( + group, + output, + &mut values, + coop::LoadOptions::default() + .algorithm(coop::Transpose) + .offset(offset) + .valid_items(count), + ); + let mut matches = true; + for i in 0..I { + let first = width * i + lane; + let second = width * (first % I) + first / I; + let mut observation = 0; + while observation < 3 { + let (snapshot, index) = match observation { + 0 => (before.get(i), lane * I + i), + 1 => (after.get(i), first), + _ => (twice.get(i), second), + }; + let initialized = if mode == 4 { + index < count || (index / I + index % I) % 3 == 1 + } else { + mode == 0 || mode == 2 || index < count + }; + let expected = if initialized { + Some(if mode == 4 { + index as i32 + 17 + } else if mode == 2 && index >= count { + -7 + } else { + input[offset + q * width * I + index] + }) + } else { + None + }; + matches &= snapshot == expected; + observation += 1; + } + let index = lane * I + i; + let expected = if index < count { + Some(if mode == 4 { + index as i32 + 17 + } else { + input[offset + q * width * I + index] + }) + } else { + None + }; + matches &= values.get(i) == expected; + } + assert!(matches); + } + + #[inline(never)] + fn selected( + input: &[i32], + output: &mut coop::DisjointSlice>, + valid: usize, + mode: u32, + ) where + coop::ThreadsWithinWarp: coop::LoadStoreGroup, + { + let group = coop::this_warp().group_by::(); + let q = coop::this_block().rank_at::() as usize * (32 / W) + + group.rank_at::() as usize; + let block = 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); + exercise( + group, + input, + output, + W, + q, + group.rank() as usize, + block * 64 * I, + valid, + mode, + ); + } + #[kernel] + #[launch_contract(domain=1,block=(64,1,1),dynamic_shared=inferred)] + pub fn logical_transpose( + input: &[i32], + mut output: coop::DisjointSlice>, + valid: usize, + mode: u32, + ) where + coop::ThreadsWithinWarp: coop::LoadStoreGroup, + { + selected::(input, &mut output, valid, mode); + } + #[kernel] + #[launch_contract(domain=1,block=(128,1,1),dynamic_shared=inferred)] + pub fn logical_transpose_b128( + input: &[i32], + mut output: coop::DisjointSlice>, + ) { + let group = coop::this_warp().group_by::<8>(); + let q = coop::this_block().rank_at::() as usize * 4 + + group.rank_at::() as usize; + exercise( + group, + input, + &mut output, + 8, + q, + group.rank() as usize, + thread::blockIdx_x() as usize * 512, + 17, + 4, + ); + } + #[kernel] + #[launch_contract(domain=3,block=(8,4,2),dynamic_shared=inferred)] + pub fn logical_transpose_xyz( + input: &[i32], + mut output: coop::DisjointSlice>, + valid: usize, + mode: u32, + ) { + selected::<8, 4>(input, &mut output, valid, mode); + } + #[kernel] + #[launch_contract(domain=1,block=(64,1,1),dynamic_shared=inferred)] + pub fn logical_transpose_prefix( + input: &[i32], + mut output: coop::DisjointSlice>, + valid: usize, + mode: u32, + ) { + // SAFETY: one whole fixed prefix, disjoint lane writes and Block completion. + let lane = thread::threadIdx_x() as usize; + let prefix = + unsafe { cuda_device::DynamicSharedArray::<[i32; 64], 32>::get_fixed() as *mut i32 }; + unsafe { + prefix.add(lane).write(lane as i32); + } + thread::sync_threads(); + selected::<8, 4>(input, &mut output, valid, mode); + let neighbor = (lane + 1) % 64; + assert!(unsafe { prefix.add(neighbor).read() } == neighbor as i32); + } + #[kernel] + #[launch_contract(domain=1,block=(64,1,1),dynamic_shared=inferred)] + pub fn logical_transpose_containing( + input: &[i32], + mut output: coop::DisjointSlice>, + ) { + use cuda_device::cooperative_groups::{ThreadGroup, this_thread_block}; + let warp8 = coop::this_warp().group_by::<8>(); + let warp16 = coop::this_warp().group_by::<16>(); + let tile16 = this_thread_block().tiled_partition::<16>(); + let offset = thread::blockIdx_x() as usize * 256; + let lane = thread::threadIdx_x() as usize; + let mut values = coop::ThreadData::<_, 4>::new(); + coop::load( + warp8, + input, + &mut values, + coop::LoadOptions::default() + .algorithm(coop::Transpose) + .offset(offset), + ); + coop::store( + warp8, + &mut output, + &mut values, + coop::StoreOptions::default() + .algorithm(coop::Transpose) + .offset(offset) + .valid_items(32), + ); + tile16.sync(); + coop::load( + warp16, + &mut output, + &mut values, + coop::LoadOptions::default() + .algorithm(coop::Transpose) + .offset(offset), + ); + let first = values; + coop::load( + warp16, + input, + &mut values, + coop::LoadOptions::default() + .algorithm(coop::Transpose) + .offset(offset), + ); + let second = values; + tile16.sync(); + coop::load( + warp8, + &mut output, + &mut values, + coop::LoadOptions::default() + .algorithm(coop::Transpose) + .offset(offset), + ); + let mut matches = true; + for i in 0..4 { + let expected = Some(input[offset + lane * 4 + i]); + matches &= + first.get(i) == expected && second.get(i) == expected && values.get(i) == expected; + } + assert!(matches); + } + #[kernel] + #[launch_contract(domain=1,block=(64,1,1),dynamic_shared=inferred)] + pub fn logical_transpose_mixed( + input: &[i32], + mut output: coop::DisjointSlice>, + ) { + let block = coop::this_block(); + let warp = coop::this_warp().group_by::<8>(); + let offset = thread::blockIdx_x() as usize * 256; + let lane = thread::threadIdx_x() as usize; + let mut values = coop::ThreadData::<_, 4>::new(); + coop::load( + block, + input, + &mut values, + coop::LoadOptions::default() + .algorithm(coop::Transpose) + .offset(offset), + ); + let first = values; + coop::load( + warp, + input, + &mut values, + coop::LoadOptions::default() + .algorithm(coop::Transpose) + .offset(offset), + ); + let second = values; + coop::load( + warp, + input, + &mut values, + coop::LoadOptions::default() + .algorithm(coop::Transpose) + .offset(offset), + ); + let third = values; + // W8 reuse must leave Block scratch dirty, requiring a Block barrier here. + coop::load( + block, + input, + &mut values, + coop::LoadOptions::default() + .algorithm(coop::Transpose) + .offset(offset), + ); + coop::store( + block, + &mut output, + &values, + coop::StoreOptions::default().offset(offset), + ); + let mut matches = true; + for i in 0..4 { + let expected = Some(input[offset + lane * 4 + i]); + matches &= first.get(i) == expected + && second.get(i) == expected + && third.get(i) == expected + && values.get(i) == expected; + } + assert!(matches); + } +} + +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}; + init_device_contexts(0, 1)?; + let stream = with_cuda_context(0, |context| context.default_stream())?; + // SAFETY: the typed module and its launchers compile together. + let module = unsafe { kernels::load_async(0)? }; + let mut cases = 0; + macro_rules! exercise { + ($w:literal,$i:literal,$scratch:literal,$prepare:expr,$launch:expr,$config:expr) => {{ + let plan = ($prepare)(($config)(3))?; + assert_eq!(plan.__raw_config().shared_mem_bytes, $scratch); + let input: Vec = (0..3 * 64 * $i).map(|i| i as i32 + 31).collect(); + for mode in 0..5 { + for valid in [0, 1, $w * $i - 1, $w * $i] { + let source = DeviceBuffer::from_host(&stream, &input)?; + let target = DeviceBuffer::from_host(&stream, &vec![-99_i32; input.len()])?; + let (_, target) = ($launch)(&plan, source, target, valid, mode).sync()?; + let mut result = vec![0; input.len()]; + target.copy_to_host(&stream, &mut result)?; + stream.synchronize()?; + let mut expected = vec![-99; input.len()]; + for block in 0..3 { + for q in 0..64 / $w { + let count = if mode == 0 { + $w * $i + } else if mode == 3 { + (valid + q).min($w * $i) + } else { + valid + }; + for i in 0..count { + let index = block * 64 * $i + q * $w * $i + i; + expected[index] = if mode == 4 { + i as i32 + 17 + } else { + input[index] + }; + } + } + } + assert_eq!(result, expected); + cases += 1; + } + } + }}; + } + exercise!( + 1, + 1, + 512, + |c| module.prepare_logical_transpose::<1, 1>(c), + |p, s, t, v, m| module.logical_transpose_async_owned::<1, 1, _, _>(p, s, t, v, m), + |g| LaunchConfig1D::auto_shared(g, 64) + ); + exercise!( + 8, + 4, + 1184, + |c| module.prepare_logical_transpose::<8, 4>(c), + |p, s, t, v, m| module.logical_transpose_async_owned::<8, 4, _, _>(p, s, t, v, m), + |g| LaunchConfig1D::auto_shared(g, 64) + ); + exercise!( + 32, + 4, + 1064, + |c| module.prepare_logical_transpose::<32, 4>(c), + |p, s, t, v, m| module.logical_transpose_async_owned::<32, 4, _, _>(p, s, t, v, m), + |g| LaunchConfig1D::auto_shared(g, 64) + ); + exercise!( + 8, + 4, + 1184, + |c| module.prepare_logical_transpose_xyz(c), + |p, s, t, v, m| module.logical_transpose_xyz_async_owned(p, s, t, v, m), + |g| cuda_core::LaunchConfig3D::auto_shared((1, g, 1), (8, 4, 2)) + ); + exercise!( + 8, + 4, + 1440, + |c| module.prepare_logical_transpose_prefix(c), + |p, s, t, v, m| module.logical_transpose_prefix_async_owned(p, s, t, v, m), + |g| LaunchConfig1D::auto_shared(g, 64) + ); + assert!( + module + .prepare_logical_transpose::<8, 4>(LaunchConfig1D::new(3, 64, 1183)) + .is_err() + ); + let plan = module.prepare_logical_transpose_mixed(LaunchConfig1D::auto_shared(3, 64))?; + assert_eq!(plan.__raw_config().shared_mem_bytes, 2080); + let input: Vec = (0..768).map(|i| i + 53).collect(); + let source = DeviceBuffer::from_host(&stream, &input)?; + let target = DeviceBuffer::from_host(&stream, &vec![-99_i32; 768])?; + let (_, target) = module + .logical_transpose_mixed_async_owned(&plan, source, target) + .sync()?; + let mut result = vec![0; 768]; + target.copy_to_host(&stream, &mut result)?; + stream.synchronize()?; + assert_eq!(result, input); + let plan = module.prepare_logical_transpose_containing(LaunchConfig1D::auto_shared(3, 64))?; + assert_eq!(plan.__raw_config().shared_mem_bytes, 2224); + let source = DeviceBuffer::from_host(&stream, &input)?; + let target = DeviceBuffer::from_host(&stream, &vec![-99_i32; 768])?; + let (_, target) = module + .logical_transpose_containing_async_owned(&plan, source, target) + .sync()?; + target.copy_to_host(&stream, &mut result)?; + stream.synchronize()?; + assert_eq!(result, input); + println!( + "PASS {cases} logical Transpose cases plus mixed Block/W8 and W8/W16 pool reuse, containing synchronization and undersized launch rejection" + ); + run_b128()?; + Ok(()) +} + +/// One changed-replica-count contrast using the same value/validity helper. +fn run_b128() -> Result<(), Box> { + use cuda_async::simt::device_context::with_cuda_context; + use cuda_async::simt::device_operation::DeviceOperation; + use cuda_core::{DeviceBuffer, LaunchConfig1D}; + let stream = with_cuda_context(0, |context| context.default_stream())?; + // SAFETY: the typed module and its launchers compile together. + let module = unsafe { kernels::load_async(0)? }; + let plan = module.prepare_logical_transpose_b128(LaunchConfig1D::auto_shared(3, 128))?; + assert_eq!(plan.__raw_config().shared_mem_bytes, 148 * 16); + let input: Vec = (0..1536).map(|i| i + 31).collect(); + let source = DeviceBuffer::from_host(&stream, &input)?; + let target = DeviceBuffer::from_host(&stream, &vec![-99_i32; input.len()])?; + let (_, target) = module + .logical_transpose_b128_async_owned(&plan, source, target) + .sync()?; + let mut result = vec![0; input.len()]; + target.copy_to_host(&stream, &mut result)?; + stream.synchronize()?; + let mut expected = vec![-99; input.len()]; + for block in 0..3 { + for q in 0..16 { + for i in 0..17 { + expected[block * 512 + q * 32 + i] = i as i32 + 17; + } + } + } + assert_eq!(result, expected); + println!("PASS B128/W8 Transpose: 16 replicas, 2368 bytes, observed values/validity"); + Ok(()) +} diff --git a/rust/cuda-coop/tests/load_store/src/logical_warp.rs b/rust/cuda-coop/tests/load_store/src/logical_warp.rs new file mode 100644 index 000000000000..7d9b3af76c45 --- /dev/null +++ b/rust/cuda-coop/tests/load_store/src/logical_warp.rs @@ -0,0 +1,504 @@ +// Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +//! Logical Warp Direct through inferred public calls and a generic helper. +use cuda_coop as coop; +use cuda_device::{cooperative_groups as cg, kernel, launch_contract, thread}; +use cuda_host::cuda_module; + +#[cuda_module] +mod kernels { + use super::*; + use cg::ThreadGroup as _; + + #[inline(never)] + #[allow(clippy::too_many_arguments)] + fn exercise, const I: usize, const N: u32>( + group: coop::ThreadGroup, + tile: cg::WarpTile, + input: &[i32], + output: &mut coop::DisjointSlice>, + width: usize, + q: usize, + lane: usize, + offset: usize, + valid: usize, + mode: u32, + ) { + let count = if mode == 0 { + width * I + } else if mode == 3 { + (valid + q).min(width * I) + } else { + valid + }; + let mut values = coop::ThreadData::<_, I>::new(); + if mode == 4 { + let mut other = coop::ThreadData::<_, I>::new(); + let mut i = 0; + while i < I { + if lane * I + i < count || i % 2 == 0 { + values.write(i, -7); + } + if lane * I + i < count || i % 2 != 0 { + other.write(i, -7); + } + i += 1; + } + values = if lane & 1 == 0 { values } else { other }; + } else if mode == 0 { + coop::load( + group, + input, + &mut values, + coop::LoadOptions::default().offset(offset), + ); + } else if mode == 2 { + coop::load( + group, + input, + &mut values, + coop::LoadOptions::default() + .offset(offset) + .valid_items(count) + .oob_default(-7), + ); + } else { + // A preceding full Load makes the exact masked reset observable. + coop::load( + group, + input, + &mut values, + coop::LoadOptions::default().offset(offset), + ); + coop::load( + group, + input, + &mut values, + coop::LoadOptions::default() + .offset(offset) + .valid_items(count), + ); + } + let before = values; + coop::store( + group, + &mut *output, + &values, + coop::StoreOptions::default() + .offset(offset) + .valid_items(count), + ); + let after = values; + tile.sync(); + coop::load( + group, + output, + &mut values, + coop::LoadOptions::default() + .offset(offset) + .valid_items(count), + ); + let mut matches = true; + let mut i = 0; + while i < I { + matches &= before.get(i) == after.get(i); + if lane * I + i < count { + let expected = if mode == 4 { + -7 + } else { + input[offset + q * width * I + lane * I + i] + }; + matches &= values.get(i) == Some(expected); + } else { + matches &= values.get(i).is_none(); + if mode == 1 || mode == 3 { + matches &= before.get(i).is_none(); + } + if mode == 2 { + matches &= before.get(i) == Some(-7); + } + } + i += 1; + } + assert!(matches); + } + + #[inline(never)] + fn selected( + input: &[i32], + output: &mut coop::DisjointSlice>, + valid: usize, + mode: u32, + ) where + coop::ThreadsWithinWarp: coop::LoadStoreGroup, + { + let group = coop::this_warp().group_by::(); + let q = coop::this_block().rank_at::() as usize * (32 / W) + + group.rank_at::() as usize; + let block = 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); + exercise( + group, + cg::this_thread_block().tiled_partition::(), + input, + output, + W, + q, + group.rank() as usize, + block * 64 * I, + valid, + mode, + ); + } + #[kernel] + #[launch_contract(domain=1,block=(64,1,1),dynamic_shared=inferred)] + pub fn logical_direct( + input: &[i32], + mut output: coop::DisjointSlice>, + valid: usize, + mode: u32, + ) where + coop::ThreadsWithinWarp: coop::LoadStoreGroup, + { + selected::(input, &mut output, valid, mode); + } + #[kernel] + #[launch_contract(domain=3,block=(8,4,2),dynamic_shared=inferred)] + pub fn logical_xyz( + input: &[i32], + mut output: coop::DisjointSlice>, + valid: usize, + mode: u32, + ) { + selected::<8, 4, 8>(input, &mut output, valid, mode); + } + #[kernel] + #[launch_contract(domain=1,block=(64,1,1),dynamic_shared=inferred)] + pub fn logical_nonexhaustive( + input: &[i32], + mut output: coop::DisjointSlice>, + valid: usize, + mode: u32, + ) { + let group = coop::this_warp().group_by_non_exhaustive::<8>(); + let q = coop::this_block().rank_at::() as usize * 4 + + group.rank_at::() as usize; + exercise( + group, + cg::this_thread_block().tiled_partition::<8>(), + input, + &mut output, + 8, + q, + group.rank() as usize, + thread::blockIdx_x() as usize * 256, + valid, + mode, + ); + } + #[kernel] + #[launch_contract(domain=1,block=(64,1,1),dynamic_shared=inferred)] + pub fn logical_grain( + mut output: coop::DisjointSlice>, + adapt: bool, + ) { + let group = coop::this_warp().group_by::<8>(); + let q = coop::this_block().rank_at::() as usize * 4 + + group.rank_at::() as usize; + let offset = thread::blockIdx_x() as usize * 256 + if adapt { q * 24 } else { 0 }; + coop::store( + group, + &mut output, + 7_i32, + coop::StoreOptions::default().offset(offset), + ); + } + + #[kernel] + #[launch_contract(domain=1,block=(64,1,1),dynamic_shared=inferred)] + pub fn logical_cycle( + input: &[i32], + mut output: coop::DisjointSlice>, + ) { + let group = coop::this_warp().group_by::<8>(); + let q = coop::this_block().rank_at::() as usize * 4 + + group.rank_at::() as usize; + let block = thread::blockIdx_x() as usize; + let lane = thread::threadIdx_x() as usize; + let base = block * 9 * 256; + let mut values = coop::ThreadData::<_, 4>::new(); + let mut sums = [0_i32; 4]; + let mut matches = true; + let mut iteration = 0; + // Siblings execute zero through seven iterations over immutable input. + while iteration < q { + let offset = base + iteration * 256; + coop::load( + group, + input, + &mut values, + coop::LoadOptions::default().offset(offset), + ); + let mut i = 0; + while i < 4 { + matches &= values.get(i) == Some(input[offset + lane * 4 + i]); + sums[i] += values.get(i).unwrap_or(0); + i += 1; + } + iteration += 1; + } + if COMPLETE { + cg::this_thread_block().tiled_partition::<32>().sync(); + } + coop::load( + coop::this_warp(), + input, + &mut values, + coop::LoadOptions::default().offset(base + 7 * 256), + ); + let mut i = 0; + while i < 4 { + matches &= values.get(i) == Some(input[base + 7 * 256 + lane * 4 + i]); + sums[i] += 3 * values.get(i).unwrap_or(0); + i += 1; + } + if COMPLETE { + thread::sync_threads(); + } + coop::load( + coop::this_block(), + input, + &mut values, + coop::LoadOptions::default().offset(base + 8 * 256), + ); + let mut i = 0; + while i < 4 { + matches &= values.get(i) == Some(input[base + 8 * 256 + lane * 4 + i]); + values.write(i, sums[i] + 5 * values.get(i).unwrap_or(0)); + i += 1; + } + coop::store( + coop::this_block(), + &mut output, + &values, + coop::StoreOptions::default().offset(block * 256), + ); + assert!(matches); + } + + #[kernel] + #[launch_contract(domain=1,block=(64,1,1),dynamic_shared=inferred)] + pub fn logical_mixed_block( + input: &[i32], + mut output: coop::DisjointSlice>, + ) { + let block = coop::this_block(); + let group = coop::this_warp().group_by::<8>(); + let offset = thread::blockIdx_x() as usize * 256; + let lane = thread::threadIdx_x() as usize; + let mut values = coop::ThreadData::<_, 4>::new(); + coop::load( + block, + input, + &mut values, + coop::LoadOptions::default() + .offset(offset) + .algorithm(coop::Transpose), + ); + coop::store( + group, + &mut output, + &values, + coop::StoreOptions::default().offset(offset), + ); + cg::this_thread_block().tiled_partition::<8>().sync(); + coop::load( + group, + &mut output, + &mut values, + coop::LoadOptions::default().offset(offset), + ); + let mut matches = (values.get(0) == Some(input[offset + lane * 4])) + & (values.get(1) == Some(input[offset + lane * 4 + 1])) + & (values.get(2) == Some(input[offset + lane * 4 + 2])) + & (values.get(3) == Some(input[offset + lane * 4 + 3])); + // Tile completion leaves the Block scratch dirty. Managed reuse must + // insert a full-block barrier before this second exchange. + coop::load( + block, + input, + &mut values, + coop::LoadOptions::default() + .offset(offset) + .algorithm(coop::Transpose), + ); + for i in 0..4 { + matches &= values.get(i) == Some(input[offset + lane * 4 + i]); + } + assert!(matches); + } +} + +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}; + init_device_contexts(0, 1)?; + let stream = with_cuda_context(0, |context| context.default_stream())?; + // SAFETY: this module and its typed launchers compile together. + let module = unsafe { kernels::load_async(0)? }; + let mut cases = 0; + macro_rules! exercise { + ($width:literal,$items:literal,$prepare:expr,$launch:expr,$config:expr) => {{ + let groups = 3; + let config = ($config)(groups); + let plan = ($prepare)(config)?; + assert_eq!(plan.__raw_config().shared_mem_bytes, 0); + let input: Vec = (0..groups as usize * 64 * $items) + .map(|i| i as i32 + 31) + .collect(); + for mode in 0..5 { + for valid in [0, 1, $width * $items - 1, $width * $items] { + let source = DeviceBuffer::from_host(&stream, &input)?; + let target = DeviceBuffer::from_host(&stream, &vec![-99_i32; input.len()])?; + let (_, target) = ($launch)(&plan, source, target, valid, mode).sync()?; + let mut result = vec![0; input.len()]; + target.copy_to_host(&stream, &mut result)?; + stream.synchronize()?; + let mut expected = vec![-99; input.len()]; + for block in 0..groups as usize { + for q in 0..64 / $width { + let count = if mode == 0 { + $width * $items + } else if mode == 3 { + (valid + q).min($width * $items) + } else { + valid + }; + for i in 0..count { + let index = block * 64 * $items + q * $width * $items + i; + expected[index] = if mode == 4 { -7 } else { input[index] }; + } + } + } + assert_eq!(result, expected); + cases += 1; + } + } + }}; + } + exercise!( + 1, + 1, + |c| module.prepare_logical_direct::<1, 1, 1>(c), + |p, s, t, v, m| module.logical_direct_async_owned::<1, 1, 1, _, _>(p, s, t, v, m), + |g| LaunchConfig1D::auto_shared(g, 64) + ); + exercise!( + 8, + 4, + |c| module.prepare_logical_direct::<8, 4, 8>(c), + |p, s, t, v, m| module.logical_direct_async_owned::<8, 4, 8, _, _>(p, s, t, v, m), + |g| LaunchConfig1D::auto_shared(g, 64) + ); + exercise!( + 16, + 4, + |c| module.prepare_logical_direct::<16, 4, 16>(c), + |p, s, t, v, m| module.logical_direct_async_owned::<16, 4, 16, _, _>(p, s, t, v, m), + |g| LaunchConfig1D::auto_shared(g, 64) + ); + exercise!( + 8, + 4, + |c| module.prepare_logical_xyz(c), + |p, s, t, v, m| module.logical_xyz_async_owned(p, s, t, v, m), + |g| cuda_core::LaunchConfig3D::auto_shared((1, g, 1), (8, 4, 2)) + ); + exercise!( + 8, + 4, + |c| module.prepare_logical_nonexhaustive(c), + |p, s, t, v, m| module.logical_nonexhaustive_async_owned(p, s, t, v, m), + |g| LaunchConfig1D::auto_shared(g, 64) + ); + let grain_plan = module.prepare_logical_grain(LaunchConfig1D::auto_shared(2, 64))?; + let target = DeviceBuffer::from_host(&stream, &vec![-99_i32; 512])?; + let adapt = !std::env::args().any(|arg| arg == "logical-grain-trap"); + let target = module + .logical_grain_async_owned(&grain_plan, target, adapt) + .sync()?; + assert!(adapt, "expected fixed-grain trap completed successfully"); + let mut result = vec![0; 512]; + target.copy_to_host(&stream, &mut result)?; + stream.synchronize()?; + let mut expected = vec![-99; 512]; + for block in 0..2 { + for q in 0..8 { + expected[block * 256 + q * 32..block * 256 + q * 32 + 8].fill(7); + } + } + assert_eq!(result, expected); + println!("PASS {cases} logical Warp mask/default/scattered cases and fixed-grain offset24"); + + let plan = module.prepare_logical_mixed_block(LaunchConfig1D::auto_shared(3, 64))?; + assert_eq!(plan.__raw_config().shared_mem_bytes, 1024); + let input: Vec = (0..768).map(|i| i + 53).collect(); + let source = DeviceBuffer::from_host(&stream, &input)?; + let target = DeviceBuffer::from_host(&stream, &vec![-99_i32; input.len()])?; + let (_, target) = module + .logical_mixed_block_async_owned(&plan, source, target) + .sync()?; + let mut result = vec![0; input.len()]; + target.copy_to_host(&stream, &mut result)?; + stream.synchronize()?; + assert_eq!(result, input); + println!("PASS mixed Block Transpose/logical Warp Direct: 768 values"); + run_cycles()?; + Ok(()) +} + +fn run_cycles() -> Result<(), Box> { + use cuda_async::simt::device_context::with_cuda_context; + use cuda_async::simt::device_operation::DeviceOperation; + use cuda_core::{DeviceBuffer, LaunchConfig1D}; + let stream = with_cuda_context(0, |context| context.default_stream())?; + // SAFETY: this module and its typed launchers compile together. + let module = unsafe { kernels::load_async(0)? }; + let input: Vec = (0..3 * 9 * 256).map(|i| i + 17).collect(); + let expected: Vec = (0..3 * 256) + .map(|index| { + let base = index / 256 * 9 * 256; + let element = index % 256; + let q = element / 4 / 8; + (0..q) + .map(|iteration| input[base + iteration * 256 + element]) + .sum::() + + 3 * input[base + 7 * 256 + element] + + 5 * input[base + 8 * 256 + element] + }) + .collect(); + macro_rules! exercise { + ($complete:literal) => {{ + let plan = + module.prepare_logical_cycle::<$complete>(LaunchConfig1D::auto_shared(3, 64))?; + assert_eq!(plan.__raw_config().shared_mem_bytes, 0); + let source = DeviceBuffer::from_host(&stream, &input)?; + let target = DeviceBuffer::from_host(&stream, &vec![-99_i32; expected.len()])?; + let (_, target) = module + .logical_cycle_async_owned::<$complete, _, _>(&plan, source, target) + .sync()?; + let mut result = vec![0; expected.len()]; + target.copy_to_host(&stream, &mut result)?; + stream.synchronize()?; + assert_eq!(result, expected); + }}; + } + exercise!(false); + exercise!(true); + println!("PASS finite sibling cycles then W32/Block: 1536 values, without/with barriers"); + Ok(()) +} 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..2c27de8e667f --- /dev/null +++ b/rust/cuda-coop/tests/load_store/src/main.rs @@ -0,0 +1,766 @@ +// 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 block_warp; +mod logical_plain; +mod logical_transpose; +mod logical_warp; +mod storage; +mod striped; +mod transpose; +mod vectorize; +mod warp; + +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 std::env::args().any(|arg| arg == "logical-plain") { + return logical_plain::run(); + } + if std::env::args().any(|arg| arg == "logical-transpose") { + return logical_transpose::run(); + } + if matches!( + std::env::args().nth(1).as_deref(), + Some("--warp" | "--warp-grain-trap") + ) { + return warp::run(); + } + if let Some(argument) = std::env::args().nth(1) { + if matches!(argument.as_str(), "logical-warp" | "logical-grain-trap") { + return logical_warp::run(); + } + if argument == "storage" { + return storage::run(); + } + if argument == "block-warp" { + return block_warp::run(None); + } + if argument.starts_with("block-warp-") { + return block_warp::run(Some(&argument)); + } + if argument == "vectorize" { + return vectorize::run(None); + } + if argument.starts_with("vectorize-") { + return vectorize::run(Some(&argument)); + } + if argument == "striped" { + return striped::run(None); + } + if argument.starts_with("striped-") { + return striped::run(Some(&argument)); + } + 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/storage.rs b/rust/cuda-coop/tests/load_store/src/storage.rs new file mode 100644 index 000000000000..36435cd6d454 --- /dev/null +++ b/rust/cuda-coop/tests/load_store/src/storage.rs @@ -0,0 +1,357 @@ +// Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +//! Public storage policies and observable cross-lane reservation access. +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 load_tile( + input: &[i32], + values: &mut coop::ThreadData, + offset: usize, + valid: usize, + scratch: &mut coop::TempStorage, + ) { + coop::load( + coop::this_block(), + input, + values, + coop::LoadOptions::default() + .offset(offset) + .valid_items(valid) + .oob_default(-7) + .temp_storage(scratch) + .algorithm(coop::Transpose), + ); + } + + #[inline(never)] + fn store_tile( + output: &mut DisjointSlice>, + values: &mut coop::ThreadData, + offset: usize, + valid: usize, + scratch: &mut coop::TempStorage, + ) { + coop::store( + coop::this_block(), + output, + values, + coop::StoreOptions::default() + .offset(offset) + .valid_items(valid) + .algorithm(coop::Transpose) + .temp_storage(scratch), + ); + } + + #[inline(never)] + fn descriptor() -> coop::TempStorage { + coop::TempStorage::new() + } + + #[kernel] + #[launch_contract(domain = 1, block = (33, 1, 1), dynamic_shared = inferred)] + pub fn exercise( + input: &[i32], + mut output: DisjointSlice>, + valid: usize, + ) { + let block = coop::this_block(); + let offset = thread::blockIdx_x() as usize * 33 * 4; + let lane = thread::threadIdx_x() as usize; + let mut values = coop::ThreadData::<_, ITEMS>::new(); + match MODE { + 0 => { + coop::load( + block, + input, + &mut values, + coop::LoadOptions::default() + .offset(offset) + .valid_items(valid) + .oob_default(-7) + .algorithm(coop::Transpose), + ); + coop::store( + block, + &mut output, + &mut values, + coop::StoreOptions::default() + .offset(offset) + .valid_items(valid) + .algorithm(coop::Transpose), + ); + } + 1 => { + let mut scratch = descriptor(); + load_tile(input, &mut values, offset, valid, &mut scratch); + thread::sync_threads(); + store_tile(&mut output, &mut values, offset, valid, &mut scratch); + } + 2 => { + let mut scratch = coop::TempStorage::new().auto_sync(); + let mut iteration = 0; + while iteration < 2 + (valid & 1) { + load_tile(input, &mut values, offset, valid, &mut scratch); + store_tile(&mut output, &mut values, offset, valid, &mut scratch); + iteration += 1; + } + } + 3 => { + let mut scratch = coop::TempStorage::new().with_alignment::<64>().exclusive(); + // The same helper at two sites receives distinct exclusive slices. + load_tile(input, &mut values, offset, valid, &mut scratch); + load_tile(input, &mut values, offset, valid, &mut scratch); + store_tile(&mut output, &mut values, offset, valid, &mut scratch); + } + 4 => { + let mut first = descriptor(); + let mut second = descriptor(); + load_tile(input, &mut values, offset, valid, &mut first); + load_tile(input, &mut values, offset, valid, &mut second); + thread::sync_threads(); + store_tile(&mut output, &mut values, offset, valid, &mut first); + } + 5 => { + let mut scratch = descriptor(); + let view = scratch.reserve::(); + // SAFETY: each lane initializes its own in-bounds element, then + // all writers complete before a different lane reads it. + unsafe { core::ptr::write(view.as_mut_ptr().add(lane), lane as i32) }; + thread::sync_threads(); + let saved = unsafe { view.as_mut_ptr().add((lane + 1) % 33).read() }; + view.release(); + thread::sync_threads(); + load_tile(input, &mut values, offset, valid, &mut scratch); + let mut item = 0; + while item < ITEMS { + // Filled Load initialized every slot; the saved cross-lane + // value remains observable after the shared region is reused. + let value = values.get(item).unwrap_or(-100_000); + values.write(item, value + saved); + item += 1; + } + thread::sync_threads(); + store_tile(&mut output, &mut values, offset, valid, &mut scratch); + } + 6 => { + let mut scratch = coop::TempStorage::new().exclusive(); + let first = scratch.reserve::(); + let second = scratch.reserve_aligned::(); + // SAFETY: distinct exclusive sites do not overlap. Each lane + // writes its own slots before the full-block completion barrier. + unsafe { + core::ptr::write(first.as_mut_ptr().add(lane), lane as i32); + core::ptr::write(second.as_mut_ptr().add(lane), (2 * lane + 1) as u64); + } + thread::sync_threads(); + let neighbor = (lane + 1) % 33; + let saved = unsafe { + first.as_mut_ptr().add(neighbor).read() + + second.as_mut_ptr().add(neighbor).read() as i32 + }; + first.release(); + second.release(); + thread::sync_threads(); + load_tile(input, &mut values, offset, valid, &mut scratch); + let mut item = 0; + while item < ITEMS { + let value = values.get(item).unwrap_or(-100_000); + values.write(item, value + saved); + item += 1; + } + store_tile(&mut output, &mut values, offset, valid, &mut scratch); + } + 7 => { + let mut scratch = coop::TempStorage::new().auto_sync(); + coop::load( + block, + input, + &mut values, + coop::LoadOptions::default() + .offset(offset) + .valid_items(valid) + .oob_default(-7) + .temp_storage(&mut scratch), + ); + coop::store( + block, + &mut output, + &values, + coop::StoreOptions::default() + .offset(offset) + .valid_items(valid) + .temp_storage(&mut scratch), + ); + } + _ => { + let mut scratch = coop::TempStorage::new() + .with_capacity::<4096>() + .with_alignment::<64>() + .auto_sync(); + load_tile(input, &mut values, offset, valid, &mut scratch); + store_tile(&mut output, &mut values, offset, valid, &mut scratch); + } + } + } + + #[kernel] + #[launch_contract(domain = 1, block = (64, 1, 1), dynamic_shared = inferred)] + pub fn timesliced_reuse( + input: &[i32], + mut output: DisjointSlice>, + valid: usize, + ) { + let block = coop::this_block(); + let offset = thread::blockIdx_x() as usize * 256; + let lane = thread::threadIdx_x() as usize; + let mut values = coop::ThreadData::<_, 4>::new(); + let mut scratch = coop::TempStorage::new().auto_sync(); + let mut correct = true; + let mut iteration = 0; + while iteration < 2 { + coop::load( + block, + input, + &mut values, + coop::LoadOptions::default() + .offset(offset) + .valid_items(valid) + .algorithm(coop::WarpTransposeTimesliced) + .temp_storage(&mut scratch), + ); + coop::store( + block, + &mut output, + &mut values, + coop::StoreOptions::default() + .offset(offset) + .valid_items(valid) + .algorithm(coop::WarpTransposeTimesliced) + .temp_storage(&mut scratch), + ); + let mut item = 0; + while item < 4 { + let index = (lane / 32) * 128 + 32 * item + lane % 32; + let expected = if index < valid { + Some((offset + index) as i32 * 7 + 3) + } else { + None + }; + correct &= values.get(item) == expected; + item += 1; + } + iteration += 1; + } + assert!(correct); + } +} + +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}; + + init_device_contexts(0, 1)?; + let context = with_cuda_context(0, |context| context.clone())?; + let stream = context.default_stream(); + // SAFETY: the embedded module and typed declarations are built together. + let module = unsafe { kernels::load_async(0)? }; + let input: Vec = (0..3 * 33 * 4).map(|i| i * 7 + 3).collect(); + let mut cases = 0; + macro_rules! run { + ($items:literal, $mode:literal) => {{ + let plan = + module.prepare_exercise::<$items, $mode>(LaunchConfig1D::auto_shared(3, 33))?; + let scratch = plan.__raw_config().shared_mem_bytes; + if $mode == 7 { + assert_eq!(scratch, 0); + } else if $mode == 8 { + assert_eq!(scratch, 4096); + } else { + assert!(scratch > 0); + } + println!("storage mode={} items={} shared={scratch}", $mode, $items); + for valid in [0, 1, 33 * $items - 1, 33 * $items] { + let source = DeviceBuffer::from_host(&stream, &input)?; + let target = DeviceBuffer::from_host(&stream, &vec![-19_i32; input.len()])?; + let (_, target) = module + .exercise_async_owned::<$items, $mode, _, _>(&plan, source, target, valid) + .sync()?; + let mut output = vec![0; input.len()]; + target.copy_to_host(&stream, &mut output)?; + stream.synchronize()?; + for (index, actual) in output.into_iter().enumerate() { + let local = index % (33 * 4); + let expected = if local < valid { + let neighbor = (local / $items + 1) % 33; + input[index] + + if $mode == 5 { + neighbor as i32 + } else if $mode == 6 { + (3 * neighbor + 1) as i32 + } else { + 0 + } + } else { + -19 + }; + assert_eq!( + actual, expected, + "mode={} items={} valid={valid} index={index}", + $mode, $items + ); + } + cases += 1; + } + }}; + } + macro_rules! modes { + ($items:literal) => { + run!($items, 0); + run!($items, 1); + run!($items, 2); + run!($items, 3); + run!($items, 4); + run!($items, 5); + run!($items, 6); + run!($items, 7); + run!($items, 8); + }; + } + modes!(1); + modes!(4); + // Counted timesliced Store needs its full, non-timesliced validity plane: + // 528 CUB bytes plus 32 mask bytes for two warps at four items per thread. + let plan = module.prepare_timesliced_reuse(LaunchConfig1D::auto_shared(3, 64))?; + assert_eq!(plan.__raw_config().shared_mem_bytes, 560); + let input: Vec = (0..3 * 256).map(|i| i * 7 + 3).collect(); + for valid in [0, 1, 255, 256] { + let source = DeviceBuffer::from_host(&stream, &input)?; + let target = DeviceBuffer::from_host(&stream, &vec![-19_i32; input.len()])?; + let (_, target) = module + .timesliced_reuse_async_owned(&plan, source, target, valid) + .sync()?; + let mut output = vec![0; input.len()]; + target.copy_to_host(&stream, &mut output)?; + stream.synchronize()?; + for (index, actual) in output.into_iter().enumerate() { + let expected = if index % 256 < valid { + input[index] + } else { + -19 + }; + assert_eq!(actual, expected, "timesliced valid={valid} index={index}"); + } + cases += 1; + } + println!("PASS: {cases} public storage cases"); + Ok(()) +} diff --git a/rust/cuda-coop/tests/load_store/src/striped.rs b/rust/cuda-coop/tests/load_store/src/striped.rs new file mode 100644 index 000000000000..003de60cc877 --- /dev/null +++ b/rust/cuda-coop/tests/load_store/src/striped.rs @@ -0,0 +1,415 @@ +// Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +//! Striped payload mapping and preservation, observed independently of roundtrip I/O. +use cuda_coop as coop; +use cuda_device::{CollectiveTiles, DisjointSlice, kernel, launch_contract, thread}; +use cuda_host::cuda_module; + +#[cuda_module] +mod kernels { + use super::*; + + #[kernel] + #[launch_contract(domain = 1, block = (33, 1, 1), dynamic_shared = inferred)] + pub fn copy_33( + input: &[i32], + mut output: DisjointSlice>, + ) { + let block = coop::this_block(); + let lane = thread::threadIdx_x() as usize; + let offset = thread::blockIdx_x() as usize * 33 * ITEMS_PER_THREAD; + let mut values = coop::ThreadData::<_, ITEMS_PER_THREAD>::new(); + coop::load( + block, + input, + &mut values, + coop::LoadOptions::default() + .offset(offset) + .algorithm(coop::Striped), + ); + let before = values; + coop::store( + block, + &mut output, + &values, + coop::StoreOptions::default() + .algorithm(coop::Striped) + .offset(offset), + ); + let mut observations_ok = true; + let mut item = 0; + while item < ITEMS_PER_THREAD { + let expected = input[offset + lane + 33 * item]; + observations_ok &= before.get(item) == Some(expected); + observations_ok &= values.get(item) == before.get(item); + item += 1; + } + assert!(observations_ok); + } + + fn exercise( + input: &[T], + output: &mut DisjointSlice>, + valid: usize, + shift: usize, + default: T, + mode: u32, + ) { + let block = coop::this_block(); + let x = thread::blockDim_x() as usize; + let y = thread::blockDim_y() as usize; + let threads = x * y * thread::blockDim_z() as usize; + let lane = thread::threadIdx_x() as usize + + x * (thread::threadIdx_y() as usize + y * thread::threadIdx_z() as usize); + let group = thread::blockIdx_x() as usize + + thread::gridDim_x() as usize + * (thread::blockIdx_y() as usize + + thread::gridDim_y() as usize * thread::blockIdx_z() as usize); + let offset = group * threads * 4 + shift; + let mut values = coop::ThreadData::<_, ITEMS>::new(); + match mode { + 0 => coop::load( + block, + input, + &mut values, + coop::LoadOptions::default() + .algorithm(coop::Striped) + .offset(offset), + ), + 1 | 5 | 6 => { + // A second masked Load must clear validity left by a full Load. + coop::load( + block, + input, + &mut values, + coop::LoadOptions::default() + .algorithm(coop::Striped) + .offset(offset), + ); + coop::load( + block, + input, + &mut values, + coop::LoadOptions::default() + .algorithm(coop::Striped) + .offset(offset) + .valid_items(valid), + ); + } + 2 | 4 => coop::load( + block, + input, + &mut values, + coop::LoadOptions::default() + .algorithm(coop::Striped) + .offset(offset) + .valid_items(valid) + .oob_default(default), + ), + 7 => {} + _ => { + let mut item = 0; + while item < ITEMS { + let selected = lane + threads * item < valid; + if (selected || (lane + item) & 1 == 0) + && !(mode == 8 && lane + threads * item == valid - 1) + { + values.write(item, default); + } + item += 1; + } + } + } + let before = values; + if mode == 0 || mode == 2 { + coop::store( + block, + &mut *output, + &values, + coop::StoreOptions::default() + .algorithm(coop::Striped) + .offset(offset), + ); + } else { + let written = if mode == 5 { + valid / 2 + } else if mode == 6 { + valid + 1 + } else { + valid + }; + coop::store( + block, + &mut *output, + &values, + coop::StoreOptions::default() + .algorithm(coop::Striped) + .offset(offset) + .valid_items(written), + ); + } + let mut observations_ok = true; + let mut item = 0; + while item < ITEMS { + let position = lane + threads * item; + let expected = match mode { + 3 | 8 => { + if position < valid || (lane + item) & 1 == 0 { + Some(default) + } else { + None + } + } + 7 => None, + _ => { + if position < valid || mode == 0 { + Some(input[offset + position]) + } else if mode == 2 || mode == 4 { + Some(default) + } else { + None + } + } + }; + observations_ok &= before.get(item) == expected; + observations_ok &= values.get(item) == before.get(item); + item += 1; + } + if mode == 3 || mode == 4 { + thread::sync_threads(); + // A repeated zero Store preserves arbitrary initialized/invalid slots. + coop::store( + block, + &mut *output, + &values, + coop::StoreOptions::default() + .algorithm(coop::Striped) + .offset(offset) + .valid_items(0), + ); + let mut item = 0; + while item < ITEMS { + observations_ok &= values.get(item) == before.get(item); + item += 1; + } + } + assert!(observations_ok); + } + + #[kernel] + #[launch_contract(domain = 1, block = (33, 1, 1), dynamic_shared = inferred)] + pub fn partial_33( + input: &[T], + mut output: DisjointSlice>, + valid: usize, + shift: usize, + default: T, + mode: u32, + ) { + exercise::<_, ITEMS>(input, &mut output, valid, shift, default, mode); + } + + #[kernel] + #[launch_contract(domain = 3, block = (8, 4, 4), dynamic_shared = inferred)] + pub fn partial_3d( + input: &[T], + mut output: DisjointSlice>, + valid: usize, + shift: usize, + default: T, + mode: u32, + ) { + exercise::<_, ITEMS>(input, &mut output, valid, shift, default, mode); + } + + #[kernel] + #[launch_contract(domain = 1, block = (33, 1, 1), dynamic_shared = inferred)] + pub fn striped_scalar(mut output: DisjointSlice>, valid: usize) { + let offset = thread::blockIdx_x() as usize * 33; + coop::store( + coop::this_block(), + &mut output, + 17_i32, + coop::StoreOptions::default() + .algorithm(coop::Striped) + .offset(offset) + .valid_items(valid), + ); + } +} + +pub fn run(trap: Option<&str>) -> Result<(), Box> { + use cuda_async::simt::device_context::{init_device_contexts, with_cuda_context}; + use cuda_async::simt::device_operation::DeviceOperation; + use cuda_core::{DeviceBuffer, LaunchConfig1D, LaunchConfig3D}; + init_device_contexts(0, 1)?; + let context = with_cuda_context(0, |context| context.clone())?; + let stream = context.default_stream(); + // SAFETY: embedded kernels and typed declarations compile together. + let module = unsafe { kernels::load_async(0)? }; + if let Some(trap) = trap { + let mode = if trap == "striped-larger" { 6 } else { 8 }; + let plan = module.prepare_partial_33::(LaunchConfig1D::auto_shared(1, 33))?; + let source = DeviceBuffer::from_host(&stream, &vec![7_i32; 132])?; + let target = DeviceBuffer::from_host(&stream, &vec![0_i32; 132])?; + module + .partial_33_async_owned::(&plan, source, target, 34, 0, -7, mode) + .sync()?; + panic!("invalid Striped Store unexpectedly completed"); + } + let input: Vec = (0..2 * 33 * 4).map(|i| i * 7 + 3).collect(); + let source = DeviceBuffer::from_host(&stream, &input)?; + let target = DeviceBuffer::from_host(&stream, &vec![-1_i32; input.len()])?; + let plan = module.prepare_copy_33::<4>(LaunchConfig1D::auto_shared(2, 33))?; + assert_eq!(plan.__raw_config().shared_mem_bytes, 0); + let (_, target) = module + .copy_33_async_owned::<4, _, _>(&plan, source, target) + .sync()?; + let mut output = vec![0_i32; input.len()]; + target.copy_to_host(&stream, &mut output)?; + stream.synchronize()?; + assert_eq!(output, input); + let mut cases = 1; + macro_rules! run { + ($t:ty, $items:literal, $threads:literal, $prepare:ident, $launch:ident, $config:expr, $extra:expr) => {{ + let tile = $threads * $items; + let plan = module.$prepare::<$t, $items>($config)?; + assert_eq!(plan.__raw_config().shared_mem_bytes, 0); + let prefix = ($threads + 1).min(tile); + let mut selected = vec![ + (0, tile, 0), + (1, prefix, 0), + (2, 1, 0), + (3, prefix, 0), + (5, prefix, 0), + ]; + if $extra { + for valid in [0, 1, $threads - 1, $threads, prefix, tile] { + for mode in [1, 2, 3, 4, 5] { + selected.push((mode, valid, 0)); + } + } + for mode in [1, 3, 4, 5] { + selected.push((mode, tile - 7, 7)); + } + selected.push((7, 0, 0)); + } + selected.sort_unstable(); + selected.dedup(); + for (mode, valid, shift) in selected { + let written = if mode == 0 || mode == 2 { + tile + } else if mode == 5 { + valid / 2 + } else { + valid + }; + // Keep the final allocation shorter than the full tile when + // counted Store writes a prefix. Striped per-lane iterator + // formation must not create an out-of-allocation pointer. + let output_len = $threads * 4 + shift + written; + let input: Vec<$t> = (0..2 * $threads * 4 + shift) + .map(|i| (i % 101) as $t) + .collect(); + let default = 111 as $t; + let sentinel = 123 as $t; + let source = DeviceBuffer::from_host(&stream, &input)?; + let target = DeviceBuffer::from_host(&stream, &vec![sentinel; output_len])?; + let (_, target) = module + .$launch::<$t, $items, _, _>(&plan, source, target, valid, shift, default, mode) + .sync()?; + let mut output = vec![sentinel; output_len]; + target.copy_to_host(&stream, &mut output)?; + stream.synchronize()?; + for (index, actual) in output.into_iter().enumerate() { + let local = index % ($threads * 4); + let expected = if local >= shift && local - shift < written { + if mode == 3 || (mode == 2 && local - shift >= valid) { + default + } else { + input[index] + } + } else { + sentinel + }; + assert_eq!( + actual, + expected, + "type={} items={} mode={mode} valid={valid} shift={shift} index={index}", + stringify!($t), + $items + ); + } + cases += 1; + } + }}; + } + macro_rules! width { + ($t:ty) => { + run!( + $t, + 1, + 33, + prepare_partial_33, + partial_33_async_owned, + LaunchConfig1D::auto_shared(2, 33), + false + ); + run!( + $t, + 4, + 33, + prepare_partial_33, + partial_33_async_owned, + LaunchConfig1D::auto_shared(2, 33), + false + ); + }; + } + width!(i8); + width!(u8); + width!(i16); + width!(u16); + width!(i32); + width!(u32); + width!(i64); + width!(u64); + width!(f32); + width!(f64); + run!( + i32, + 4, + 33, + prepare_partial_33, + partial_33_async_owned, + LaunchConfig1D::auto_shared(2, 33), + true + ); + run!( + i32, + 4, + 128, + prepare_partial_3d, + partial_3d_async_owned, + LaunchConfig3D::auto_shared((1, 2, 1), (8, 4, 4)), + true + ); + for valid in [0, 1, 32, 33] { + let plan = module.prepare_striped_scalar(LaunchConfig1D::auto_shared(2, 33))?; + assert_eq!(plan.__raw_config().shared_mem_bytes, 0); + let target = DeviceBuffer::from_host(&stream, &vec![-1_i32; 33 + valid])?; + let target = module + .striped_scalar_async_owned(&plan, target, valid) + .sync()?; + let mut output = vec![0_i32; 33 + valid]; + target.copy_to_host(&stream, &mut output)?; + stream.synchronize()?; + for (i, value) in output.into_iter().enumerate() { + assert_eq!(value, if i % 33 < valid { 17 } else { -1 }); + } + cases += 1; + } + println!("PASS: {cases} Striped cases with observed payload mapping and validity"); + Ok(()) +} 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/load_store/src/vectorize.rs b/rust/cuda-coop/tests/load_store/src/vectorize.rs new file mode 100644 index 000000000000..c4708797870d --- /dev/null +++ b/rust/cuda-coop/tests/load_store/src/vectorize.rs @@ -0,0 +1,380 @@ +// Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +//! Vectorize full/guarded access and post-Store validity, including misaligned offsets. +use cuda_coop as coop; +use cuda_device::{CollectiveTiles, DisjointSlice, kernel, launch_contract, thread}; +use cuda_host::cuda_module; + +#[cuda_module] +mod kernels { + use super::*; + + fn exercise( + input: &[T], + output: &mut DisjointSlice>, + valid: usize, + shift: usize, + default: T, + mode: u32, + ) { + let block = coop::this_block(); + let x = thread::blockDim_x() as usize; + let y = thread::blockDim_y() as usize; + let threads = x * y * thread::blockDim_z() as usize; + let lane = thread::threadIdx_x() as usize + + x * (thread::threadIdx_y() as usize + y * thread::threadIdx_z() as usize); + let group = thread::blockIdx_x() as usize + + thread::gridDim_x() as usize + * (thread::blockIdx_y() as usize + + thread::gridDim_y() as usize * thread::blockIdx_z() as usize); + let offset = group * threads * 8 + shift; + let mut values = coop::ThreadData::<_, ITEMS>::new(); + match mode { + 0 => coop::load( + block, + input, + &mut values, + coop::LoadOptions::default() + .algorithm(coop::Vectorize) + .offset(offset), + ), + 1 | 5 | 6 | 9 => { + // A second masked Load must clear validity left by a full Load. + coop::load( + block, + input, + &mut values, + coop::LoadOptions::default() + .algorithm(coop::Vectorize) + .offset(offset), + ); + coop::load( + block, + input, + &mut values, + coop::LoadOptions::default() + .algorithm(coop::Vectorize) + .offset(offset) + .valid_items(valid), + ); + } + 2 | 4 => coop::load( + block, + input, + &mut values, + coop::LoadOptions::default() + .algorithm(coop::Vectorize) + .offset(offset) + .valid_items(valid) + .oob_default(default), + ), + 7 => {} + _ => { + let mut item = 0; + while item < ITEMS { + let selected = lane * ITEMS + item < valid; + if (selected || (lane + item) & 1 == 0) + && !(mode == 8 && lane * ITEMS + item == valid - 1) + { + values.write(item, default); + } + item += 1; + } + } + } + let before = values; + if mode == 0 || mode == 2 || mode == 9 { + coop::store( + block, + &mut *output, + &values, + coop::StoreOptions::default() + .algorithm(coop::Vectorize) + .offset(offset), + ); + } else { + let written = if mode == 5 { + valid / 2 + } else if mode == 6 { + valid + 1 + } else { + valid + }; + coop::store( + block, + &mut *output, + &values, + coop::StoreOptions::default() + .algorithm(coop::Vectorize) + .offset(offset) + .valid_items(written), + ); + } + let mut observations_ok = true; + let mut item = 0; + while item < ITEMS { + let position = lane * ITEMS + item; + let expected = match mode { + 3 | 8 => { + if position < valid || (lane + item) & 1 == 0 { + Some(default) + } else { + None + } + } + 7 => None, + _ => { + if position < valid || mode == 0 { + Some(input[offset + position]) + } else if mode == 2 || mode == 4 { + Some(default) + } else { + None + } + } + }; + observations_ok &= before.get(item) == expected; + observations_ok &= values.get(item) == before.get(item); + item += 1; + } + if mode == 3 || mode == 4 { + thread::sync_threads(); + // A repeated zero Store preserves arbitrary initialized/invalid slots. + coop::store( + block, + &mut *output, + &values, + coop::StoreOptions::default() + .algorithm(coop::Vectorize) + .offset(offset) + .valid_items(0), + ); + let mut item = 0; + while item < ITEMS { + observations_ok &= values.get(item) == before.get(item); + item += 1; + } + } + assert!(observations_ok); + } + + #[kernel] + #[launch_contract(domain = 1, block = (33, 1, 1), dynamic_shared = inferred)] + pub fn partial_33( + input: &[T], + mut output: DisjointSlice>, + valid: usize, + shift: usize, + default: T, + mode: u32, + ) { + exercise::<_, ITEMS>(input, &mut output, valid, shift, default, mode); + } + + #[kernel] + #[launch_contract(domain = 3, block = (8, 4, 4), dynamic_shared = inferred)] + pub fn partial_3d( + input: &[T], + mut output: DisjointSlice>, + valid: usize, + shift: usize, + default: T, + mode: u32, + ) { + exercise::<_, ITEMS>(input, &mut output, valid, shift, default, mode); + } + + #[kernel] + #[launch_contract(domain = 1, block = (33, 1, 1), dynamic_shared = inferred)] + pub fn vectorize_scalar(mut output: DisjointSlice>, valid: usize) { + let offset = thread::blockIdx_x() as usize * 33; + coop::store( + coop::this_block(), + &mut output, + 17_i32, + coop::StoreOptions::default() + .algorithm(coop::Vectorize) + .offset(offset) + .valid_items(valid), + ); + } +} + +pub fn run(trap: Option<&str>) -> Result<(), Box> { + use cuda_async::simt::device_context::{init_device_contexts, with_cuda_context}; + use cuda_async::simt::device_operation::DeviceOperation; + use cuda_core::{DeviceBuffer, LaunchConfig1D, LaunchConfig3D}; + init_device_contexts(0, 1)?; + let context = with_cuda_context(0, |context| context.clone())?; + let stream = context.default_stream(); + // SAFETY: embedded kernels and typed declarations compile together. + let module = unsafe { kernels::load_async(0)? }; + if let Some(trap) = trap { + let mode = match trap { + "vectorize-larger" => 6, + "vectorize-full-hole" => 9, + _ => 8, + }; + let plan = module.prepare_partial_33::(LaunchConfig1D::auto_shared(1, 33))?; + let source = DeviceBuffer::from_host(&stream, &vec![7_i32; 132])?; + let target = DeviceBuffer::from_host(&stream, &vec![0_i32; 132])?; + module + .partial_33_async_owned::(&plan, source, target, 34, 0, -7, mode) + .sync()?; + panic!("invalid Vectorize Store unexpectedly completed"); + } + let mut cases = 0; + macro_rules! run { + ($t:ty, $items:literal, $threads:literal, $prepare:ident, $launch:ident, $config:expr, $extra:expr) => {{ + let tile = $threads * $items; + let plan = module.$prepare::<$t, $items>($config)?; + assert_eq!(plan.__raw_config().shared_mem_bytes, 0); + let prefix = ($threads + 1).min(tile); + let mut selected = vec![ + (0, tile, 0), + (0, tile, 1), + (0, tile, 4), + (1, prefix, 0), + (2, 1, 0), + (3, prefix, 0), + (5, prefix, 0), + ]; + if $extra { + for valid in [0, 1, $threads - 1, $threads, prefix, tile] { + for mode in [1, 2, 3, 4, 5] { + selected.push((mode, valid, 0)); + } + } + for mode in [1, 3, 4, 5] { + selected.push((mode, tile - 7, 7)); + } + selected.push((7, 0, 0)); + } + selected.sort_unstable(); + selected.dedup(); + for (mode, valid, shift) in selected { + let written = if mode == 0 || mode == 2 || mode == 9 { + tile + } else if mode == 5 { + valid / 2 + } else { + valid + }; + // Keep the final allocation shorter than the full tile when + // counted Store writes a prefix. Its guarded iterator must not + // form an out-of-allocation pointer. + let output_len = $threads * 8 + shift + written; + let input: Vec<$t> = (0..2 * $threads * 8 + shift) + .map(|i| (i % 101) as $t) + .collect(); + let default = 111 as $t; + let sentinel = 123 as $t; + let source = DeviceBuffer::from_host(&stream, &input)?; + let target = DeviceBuffer::from_host(&stream, &vec![sentinel; output_len])?; + let (_, target) = module + .$launch::<$t, $items, _, _>(&plan, source, target, valid, shift, default, mode) + .sync()?; + let mut output = vec![sentinel; output_len]; + target.copy_to_host(&stream, &mut output)?; + stream.synchronize()?; + for (index, actual) in output.into_iter().enumerate() { + let local = index % ($threads * 8); + let expected = if local >= shift && local - shift < written { + if mode == 3 || (mode == 2 && local - shift >= valid) { + default + } else { + input[index] + } + } else { + sentinel + }; + assert_eq!( + actual, + expected, + "type={} items={} mode={mode} valid={valid} shift={shift} index={index}", + stringify!($t), + $items + ); + } + cases += 1; + } + }}; + } + macro_rules! width { + ($t:ty) => { + run!( + $t, + 1, + 33, + prepare_partial_33, + partial_33_async_owned, + LaunchConfig1D::auto_shared(2, 33), + false + ); + run!( + $t, + 4, + 33, + prepare_partial_33, + partial_33_async_owned, + LaunchConfig1D::auto_shared(2, 33), + false + ); + }; + } + width!(i8); + width!(u8); + width!(i16); + width!(u16); + width!(i32); + width!(u32); + width!(i64); + width!(u64); + width!(f32); + width!(f64); + run!( + i32, + 4, + 33, + prepare_partial_33, + partial_33_async_owned, + LaunchConfig1D::auto_shared(2, 33), + true + ); + run!( + i32, + 4, + 128, + prepare_partial_3d, + partial_3d_async_owned, + LaunchConfig3D::auto_shared((1, 2, 1), (8, 4, 4)), + true + ); + run!( + i32, + 3, + 33, + prepare_partial_33, + partial_33_async_owned, + LaunchConfig1D::auto_shared(2, 33), + false + ); + for valid in [0, 1, 32, 33] { + let plan = module.prepare_vectorize_scalar(LaunchConfig1D::auto_shared(2, 33))?; + assert_eq!(plan.__raw_config().shared_mem_bytes, 0); + let target = DeviceBuffer::from_host(&stream, &vec![-1_i32; 33 + valid])?; + let target = module + .vectorize_scalar_async_owned(&plan, target, valid) + .sync()?; + let mut output = vec![0_i32; 33 + valid]; + target.copy_to_host(&stream, &mut output)?; + stream.synchronize()?; + for (i, value) in output.into_iter().enumerate() { + assert_eq!(value, if i % 33 < valid { 17 } else { -1 }); + } + cases += 1; + } + println!("PASS: {cases} Vectorize cases with observed payload mapping and validity"); + Ok(()) +} diff --git a/rust/cuda-coop/tests/load_store/src/warp.rs b/rust/cuda-coop/tests/load_store/src/warp.rs new file mode 100644 index 000000000000..eec19289f87b --- /dev/null +++ b/rust/cuda-coop/tests/load_store/src/warp.rs @@ -0,0 +1,334 @@ +// Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +//! Physical Warp Direct uses the ordinary public API and owned async launch. +use cuda_coop as coop; +use cuda_device::{kernel, launch_contract, library, thread}; +use cuda_host::cuda_module; + +#[cuda_module] +mod kernels { + use super::*; + + #[inline(never)] + fn exercise( + input: &[i32], + output: &mut coop::DisjointSlice>, + valid: usize, + mode: u32, + ) { + let warp = coop::this_warp(); + let q = library::warp_group_index_v1::<32>() as usize; + let block = thread::blockIdx_x() as usize; + let offset = block * 64 * I; + let lane = (thread::threadIdx_x() as usize + + thread::blockDim_x() as usize + * (thread::threadIdx_y() as usize + + thread::blockDim_y() as usize * thread::threadIdx_z() as usize)) + % 32; + let count = if mode == 0 { + 32 * I + } else if mode == 3 { + (valid + q).min(32 * I) + } else { + valid + }; + let mut values = coop::ThreadData::<_, I>::new(); + if mode == 4 { + let mut other = coop::ThreadData::<_, I>::new(); + let mut i = 0; + while i < I { + if lane * I + i < count || i % 2 == 0 { + values.write(i, -7); + } + if lane * I + i < count || i % 2 != 0 { + other.write(i, -7); + } + i += 1; + } + values = if lane & 1 == 0 { values } else { other }; + } else if mode == 0 { + coop::load( + warp, + input, + &mut values, + coop::LoadOptions::default().offset(offset), + ); + } else if mode == 2 { + coop::load( + warp, + input, + &mut values, + coop::LoadOptions::default() + .offset(offset) + .valid_items(count) + .oob_default(-7), + ); + } else { + // A preceding full Load makes the exact masked reset observable. + coop::load( + warp, + input, + &mut values, + coop::LoadOptions::default().offset(offset), + ); + coop::load( + warp, + input, + &mut values, + coop::LoadOptions::default() + .offset(offset) + .valid_items(count), + ); + } + let before = values; + coop::store( + warp, + &mut *output, + &values, + coop::StoreOptions::default() + .offset(offset) + .valid_items(count), + ); + let after = values; + thread::sync_threads(); + coop::load( + warp, + output, + &mut values, + coop::LoadOptions::default() + .offset(offset) + .valid_items(count), + ); + let mut matches = true; + let mut i = 0; + while i < I { + matches &= before.get(i) == after.get(i); + if lane * I + i < count { + let expected = if mode == 4 { + -7 + } else { + input[offset + q * 32 * I + lane * I + i] + }; + matches &= values.get(i) == Some(expected); + } else { + matches &= values.get(i).is_none(); + if mode == 1 || mode == 3 { + matches &= before.get(i).is_none(); + } + if mode == 2 { + matches &= before.get(i) == Some(-7); + } + } + i += 1; + } + assert!(matches); + } + + #[kernel] + #[launch_contract(domain=1,block=(64,1,1),dynamic_shared=inferred)] + pub fn warp_direct_64( + input: &[i32], + mut output: coop::DisjointSlice>, + valid: usize, + mode: u32, + ) { + exercise::(input, &mut output, valid, mode); + } + #[kernel] + #[launch_contract(domain=3,block=(8,4,2),dynamic_shared=inferred)] + pub fn warp_direct_xyz( + input: &[i32], + mut output: coop::DisjointSlice>, + valid: usize, + mode: u32, + ) { + exercise::(input, &mut output, valid, mode); + } + + #[kernel] + #[launch_contract(domain=1,block=(64,1,1),dynamic_shared=inferred)] + pub fn warp_mixed_block_transpose( + input: &[i32], + mut output: coop::DisjointSlice>, + ) { + let offset = thread::blockIdx_x() as usize * 256; + let mut values = coop::ThreadData::<_, 4>::new(); + coop::load( + coop::this_block(), + input, + &mut values, + coop::LoadOptions::default() + .offset(offset) + .algorithm(coop::Transpose), + ); + coop::store( + coop::this_warp(), + &mut output, + &values, + coop::StoreOptions::default().offset(offset), + ); + thread::sync_threads(); + coop::load( + coop::this_warp(), + &mut output, + &mut values, + coop::LoadOptions::default().offset(offset).valid_items(127), + ); + let lane = thread::threadIdx_x() as usize; + let mut matches = true; + let mut i = 0; + while i < 4 { + let expected = if (lane % 32) * 4 + i < 127 { + Some(input[offset + lane * 4 + i]) + } else { + None + }; + matches &= values.get(i) == expected; + i += 1; + } + assert!(matches); + } + + #[kernel] + #[launch_contract(domain=1,block=(64,1,1),dynamic_shared=inferred)] + pub fn warp_fixed_grain( + mut output: coop::DisjointSlice>, + adapt: bool, + ) { + let q = library::warp_group_index_v1::<32>() as usize; + let offset = thread::blockIdx_x() as usize * 256 + if adapt { q * 96 } else { 0 }; + // Automatic origin is q*32 for this one-item scalar. Carrier capacity + // stays four, so warp one needs caller offset96 to reach its region128. + coop::store( + coop::this_warp(), + &mut output, + 7_i32, + coop::StoreOptions::default().offset(offset), + ); + } +} + +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, LaunchConfig3D}; + init_device_contexts(0, 1)?; + let stream = with_cuda_context(0, |context| context.default_stream())?; + // SAFETY: the typed launcher and embedded native kernels compile together. + let module = unsafe { kernels::load_async(0)? }; + if std::env::args().any(|argument| argument == "--warp-grain-trap") { + let plan = module.prepare_warp_fixed_grain(LaunchConfig1D::auto_shared(1, 64))?; + let target = DeviceBuffer::from_host(&stream, &vec![-99_i32; 256])?; + module + .warp_fixed_grain_async_owned(&plan, target, false) + .sync()?; + panic!("expected ownership trap returned successfully"); + } + for groups in [1_u32, 3] { + let plan = module.prepare_warp_fixed_grain(LaunchConfig1D::auto_shared(groups, 64))?; + let target = DeviceBuffer::from_host(&stream, &vec![-99_i32; groups as usize * 256])?; + let target = module + .warp_fixed_grain_async_owned(&plan, target, true) + .sync()?; + let mut result = vec![0; groups as usize * 256]; + target.copy_to_host(&stream, &mut result)?; + stream.synchronize()?; + let mut expected = vec![-99_i32; result.len()]; + for block in 0..groups as usize { + for warp in 0..2 { + expected[block * 256 + warp * 128..block * 256 + warp * 128 + 32].fill(7); + } + } + assert_eq!(result, expected); + } + // A storage-bearing Block operation and scratch-free Warp operations + // share one root plan and the same collective allocation. + let plan = module.prepare_warp_mixed_block_transpose(LaunchConfig1D::auto_shared(2, 64))?; + assert!(plan.__raw_config().shared_mem_bytes > 0); + let input: Vec = (0..512).map(|i| i + 17).collect(); + let source = DeviceBuffer::from_host(&stream, &input)?; + let target = DeviceBuffer::from_host(&stream, &vec![-99_i32; input.len()])?; + let (_, target) = module + .warp_mixed_block_transpose_async_owned(&plan, source, target) + .sync()?; + let mut result = vec![0; input.len()]; + target.copy_to_host(&stream, &mut result)?; + stream.synchronize()?; + assert_eq!(result, input); + let mut cases = 0; + macro_rules! run { + ($i:literal,$prepare:ident,$launch:ident,$config:expr,$wrong:expr) => {{ + assert!(module.$prepare::<$i>($wrong).is_err()); + for groups in [1_u32, 3] { + let input: Vec = (0..groups as usize * 64 * $i) + .map(|i| i as i32 + 11) + .collect(); + let plan = module.$prepare::<$i>($config(groups))?; + assert_eq!(plan.__raw_config().shared_mem_bytes, 0); + for mode in 0..5 { + for valid in [0, 1, 32 * $i - 1, 32 * $i] { + let source = DeviceBuffer::from_host(&stream, &input)?; + let target = DeviceBuffer::from_host(&stream, &vec![-99_i32; input.len()])?; + let (_, target) = module + .$launch::<$i, _, _>(&plan, source, target, valid, mode) + .sync()?; + let mut result = vec![0; input.len()]; + target.copy_to_host(&stream, &mut result)?; + stream.synchronize()?; + let mut expected = vec![-99; input.len()]; + for block in 0..groups as usize { + for q in 0..2 { + let count = if mode == 0 { + 32 * $i + } else if mode == 3 { + (valid + q).min(32 * $i) + } else { + valid + }; + for i in 0..count { + let index = block * 64 * $i + q * 32 * $i + i; + expected[index] = if mode == 4 { -7 } else { input[index] }; + } + } + } + assert_eq!(result, expected); + cases += 1; + } + } + } + }}; + } + run!( + 1, + prepare_warp_direct_64, + warp_direct_64_async_owned, + |g| LaunchConfig1D::auto_shared(g, 64), + LaunchConfig1D::auto_shared(1, 32) + ); + run!( + 4, + prepare_warp_direct_64, + warp_direct_64_async_owned, + |g| LaunchConfig1D::auto_shared(g, 64), + LaunchConfig1D::auto_shared(1, 32) + ); + run!( + 1, + prepare_warp_direct_xyz, + warp_direct_xyz_async_owned, + |g| LaunchConfig3D::auto_shared((g, 1, 1), (8, 4, 2)), + LaunchConfig3D::auto_shared((1, 1, 1), (4, 4, 4)) + ); + run!( + 4, + prepare_warp_direct_xyz, + warp_direct_xyz_async_owned, + |g| LaunchConfig3D::auto_shared((g, 1, 1), (8, 4, 2)), + LaunchConfig3D::auto_shared((1, 1, 1), (4, 4, 4)) + ); + println!( + "PASS {cases} physical Warp Direct cases and 2 fixed-grain cases; mixed Block Transpose/Warp Direct passed" + ); + 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..fcb8eb44bbf6 --- /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 = "1.99.1" +components = ["clippy", "rustfmt"] +profile = "minimal" From 2fa4ba8f78a0aaea9b39b4b4a2186752c2694b1f Mon Sep 17 00:00:00 2001 From: Trent Nelson Date: Sat, 24 Oct 2026 17:38:00 -0700 Subject: [PATCH 2/3] Add Rust cooperative Scan and Sum variants Expose Block and Warp scans with explicit payload validity, prefix registration and native provider integration. Signed-off-by: Trent Nelson --- docs/rust/index.rst | 43 +- docs/rust/scan.rst | 136 +++ rust/cuda-coop-codegen/src/lib.rs | 87 +- rust/cuda-coop-codegen/src/oxide.rs | 222 +++- rust/cuda-coop-codegen/src/scan.rs | 351 ++++++ rust/cuda-coop-codegen/tests/native.rs | 133 ++- rust/cuda-coop/build.rs | 67 +- rust/cuda-coop/examples/scan/.gitignore | 7 + rust/cuda-coop/examples/scan/Cargo.lock | 1019 +++++++++++++++++ rust/cuda-coop/examples/scan/Cargo.toml | 19 + .../cuda-coop/examples/scan/helper/Cargo.toml | 11 + .../cuda-coop/examples/scan/helper/src/lib.rs | 39 + rust/cuda-coop/examples/scan/src/main.rs | 75 ++ rust/cuda-coop/src/group.rs | 3 + rust/cuda-coop/src/lib.rs | 12 + rust/cuda-coop/src/scan.rs | 622 ++++++++++ rust/cuda-coop/src/scan/boundary.rs | 334 ++++++ rust/cuda-coop/tests/check_load_store.py | 110 +- rust/cuda-coop/tests/scan/.gitignore | 7 + rust/cuda-coop/tests/scan/Cargo.lock | 1011 ++++++++++++++++ rust/cuda-coop/tests/scan/Cargo.toml | 17 + rust/cuda-coop/tests/scan/src/main.rs | 540 +++++++++ 22 files changed, 4823 insertions(+), 42 deletions(-) create mode 100644 docs/rust/scan.rst create mode 100644 rust/cuda-coop-codegen/src/scan.rs create mode 100644 rust/cuda-coop/examples/scan/.gitignore create mode 100644 rust/cuda-coop/examples/scan/Cargo.lock create mode 100644 rust/cuda-coop/examples/scan/Cargo.toml create mode 100644 rust/cuda-coop/examples/scan/helper/Cargo.toml create mode 100644 rust/cuda-coop/examples/scan/helper/src/lib.rs create mode 100644 rust/cuda-coop/examples/scan/src/main.rs create mode 100644 rust/cuda-coop/src/scan.rs create mode 100644 rust/cuda-coop/src/scan/boundary.rs create mode 100644 rust/cuda-coop/tests/scan/.gitignore create mode 100644 rust/cuda-coop/tests/scan/Cargo.lock create mode 100644 rust/cuda-coop/tests/scan/Cargo.toml create mode 100644 rust/cuda-coop/tests/scan/src/main.rs diff --git a/docs/rust/index.rst b/docs/rust/index.rst index d037aa8adbf5..ad383f5f66e5 100644 --- a/docs/rust/index.rst +++ b/docs/rust/index.rst @@ -5,28 +5,28 @@ ===================== Write group-first cooperative operations in an ordinary generic Rust kernel. -This kernel copies partial block tiles through a reusable helper in another -crate. The item count stays a const generic from the host launch through the -kernel and helper: +This kernel computes exclusive prefixes for partial block tiles through a +reusable helper in another crate. The item count stays a const generic from the +host launch through the kernel and helper: -.. literalinclude:: ../../rust/cuda-coop/examples/load_store/src/main.rs +.. literalinclude:: ../../rust/cuda-coop/examples/scan/src/main.rs :language: rust - :lines: 12-35 + :lines: 12-34 :dedent: 4 -The helper uses the following calls for a partial tile. ``ThreadData`` infers its -element type from ``input``. Load fills the payload in place; Store writes the -valid prefix and leaves the rest of the destination unchanged: +The helper infers the payload element type from ``input``. Load initializes +the payload, filling the missing tail with zero. Scan returns fresh prefixes +without changing its input, and Store writes the valid output prefix: -.. literalinclude:: ../../rust/cuda-coop/examples/load_store/helper/src/lib.rs +.. literalinclude:: ../../rust/cuda-coop/examples/scan/helper/src/lib.rs :language: rust - :lines: 55-73 - :dedent: 8 + :lines: 20-38 + :dedent: 4 -The complete :doc:`load_store` example includes the helper, imports, host -allocation and launch. It runs Direct, Striped, Vectorize, Transpose, -WarpTranspose and WarpTransposeTimesliced with one and four items per thread. -These listings come directly from that executable source. +The complete :doc:`scan` example includes the helper, imports, host allocation +and launch for one and four items per thread. These listings come directly +from that executable source. The :doc:`load_store` tutorial develops the +underlying tile-copy and temporary-storage operations. Start with a kernel ------------------- @@ -43,10 +43,10 @@ as the remaining input length capped at the tile size. The kernel above shows that idiom, including a short final tile. Counts and offsets describe the whole group's tile, and must be uniform across that group. -The current public operations are Load and Store, with the supported groups, -algorithms and numeric types described in :doc:`load_store`. Additional -primitive families are under development. The crate is unpublished and its -interface can change. +The public operations include :doc:`load_store` and :doc:`scan`. Scan returns +fresh per-group prefixes while preserving its input. Additional primitive +families are under development. The crate is unpublished and its interface can +change. Build and run ------------- @@ -57,7 +57,7 @@ NVIDIA GPU. From a complete CCCL checkout with the matching Oxide installed: .. code-block:: shell - cd rust/cuda-coop/examples/load_store + cd rust/cuda-coop/examples/scan CUDA_HOME=/usr/local/cuda-13.3 CUDA_PATH=/usr/local/cuda-13.3 \ cargo oxide run --arch sm_120 --materialize-cubin @@ -80,7 +80,7 @@ provider build script, manual linker inputs or scratch-byte calculation. How compilation works --------------------- -``cargo oxide`` discovers the concrete Load/Store uses, including those inside +``cargo oxide`` discovers the concrete cooperative calls, including those inside generic helpers from another crate. The dependency's host provider specializes CUB for their types, item counts, algorithms and launch geometry. NVRTC returns device LTO-IR and the actual C++ temporary-storage layouts. Oxide checks @@ -113,6 +113,7 @@ Read next payload groups load_store + scan provider Generate the API reference with ``cargo doc --workspace --no-deps`` from diff --git a/docs/rust/scan.rst b/docs/rust/scan.rst new file mode 100644 index 000000000000..238619aa0b5c --- /dev/null +++ b/docs/rust/scan.rst @@ -0,0 +1,136 @@ +.. Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. +.. SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +Block and Warp Scan +=================== + +``scan`` computes a prefix within each participating group and returns a fresh +result. Its defaults select an exclusive sum: the first output is zero, and +each following output sums the preceding inputs. ``exclusive_sum``, +``inclusive_sum``, ``exclusive_scan`` and ``inclusive_scan`` provide the same +operation with their mode or operator fixed by the function name. All five +functions take the group first and preserve the input. + +A Block accepts a scalar or a shared or mutable borrow of ``ThreadData``. +The result has the same element type and, for a payload, the same item count +and alignment. Block payloads use blocked order: each thread contributes its +consecutive items, followed by the next thread in x-fastest rank order. +Complete one-, two- and three-dimensional blocks are supported. Each block's +prefix starts independently; this is not a device-wide scan. + +Physical Warp and logical lane groups of widths 1, 2, 4, 8, 16 and 32 accept +scalar inputs and return scalars. Prefixes restart at each group's first lane. +The enclosing block must contain complete physical warps. For example: + +.. code-block:: rust + + let group = coop::this_warp().group_by::<8>(); + let prefix = coop::exclusive_sum(group, value, Default::default()); + +A complete program +------------------ + +This independent Cargo workspace uses a generic helper in another crate. The +kernel and helper share the const-generic ``ITEMS_PER_THREAD``; the payload +element type and fresh result type are inferred. The program checks one and +four items per thread across three Block tiles, including a short final tile. + +.. literalinclude:: ../../rust/cuda-coop/examples/scan/src/main.rs + :language: rust + :start-after: // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +The helper initializes the missing tail with zero before Scan, then stores only +the valid prefix. Load returns unit; Scan returns a new payload; Store reads +that result through a shared borrow. + +.. literalinclude:: ../../rust/cuda-coop/examples/scan/helper/src/lib.rs + :language: rust + :start-after: // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +Use the Oxide revision selected by the manifests and ``1.99.1``. +With that toolchain installed, run from a complete CCCL checkout: + +.. code-block:: shell + + cd rust/cuda-coop/examples/scan + 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. The dependency supplies its host provider; +the consumer needs no build script, C++ shim, manual linker inputs or scratch +size. ``auto_shared`` uses the compiler's complete shared-memory plan. + +Modes, operators, and seeds +--------------------------- + +``ScanOptions::default()`` selects exclusive mode, ``Sum``, the default Block +algorithm, an implicit zero seed and compiler-managed storage. Consuming +setters select compile-time options and the runtime seed: + +.. code-block:: rust + + let inclusive = coop::scan(block, &values, + coop::ScanOptions::default().mode(coop::Inclusive)); + let products = coop::exclusive_scan(block, &values, + coop::ScanOptions::default() + .scan_op(coop::Multiplies).initial_value(1)); + let maxima = coop::inclusive_scan(block, &values, + coop::ScanOptions::default() + .scan_op(coop::Max).algorithm(coop::RakingMemoize)); + +``Exclusive`` emits the prefix before the current input; ``Inclusive`` includes +the current input. An exclusive seed participates in every prefix and has +exactly the input element type. Exclusive ``Sum`` may omit its seed; every +other exclusive operator requires one. Inclusive operations reject a seed. +The sum functions fix ``Sum`` and do not accept a replacement operator or seed; +the named exclusive and inclusive functions do not accept a mode override. + +The supported elements are ``i8``, ``u8``, ``i16``, ``u16``, ``i32``, ``u32``, +``i64``, ``u64``, ``f32`` and ``f64``. Integers support ``Sum``, ``Multiplies``, +``Min``, ``Max``, ``BitAnd``, ``BitOr`` and ``BitXor``. Floating types support the +first four. These operators use the input dtype; the API does not widen narrow +integers. CUB may regroup associative operations, so floating-point results +need not match a sequential host sum exactly. Arbitrary callables are not yet +supported. + +Validity and participation +-------------------------- + +Every input slot must be initialized. Scan does not accept a valid count or +skip unwritten items. A masked Load without a default can leave such slots; +Scan traps before reading an unwritten input. For a partial tile, use a filled +Load or initialize each missing slot explicitly. Choose a fill appropriate for +the operator and the output you need. Every returned slot is initialized, +including outputs corresponding to the filled tail. + +All group members must reach the same collective. Seeds and storage selections +must be uniform within the group, though separate logical groups may use +different seeds. Input values remain lane-local. Rust borrowing preserves local +ownership; the supported compiler also checks collective participation and +shared-storage use. + +Algorithms and temporary storage +-------------------------------- + +Block Scan supports ``Raking`` (the default), ``RakingMemoize`` and +``WarpScans``. Select one with ``.algorithm(...)``. ``WarpScans`` requires a +total block thread count divisible by 32. Warp Scan uses its implicit native +algorithm and storage; it rejects an explicit Block algorithm or +``TempStorage`` argument. + +Block calls can use the same public temporary-storage descriptor as +:doc:`Load/Store `: + +.. code-block:: rust + + let mut scratch = coop::TempStorage::new().auto_sync(); + let prefixes = coop::exclusive_scan(block, &values, + coop::ScanOptions::default() + .initial_value(seed).temp_storage(&mut scratch)); + +Omitting a descriptor requests managed synchronization before scratch reuse. +An explicit descriptor defaults to manual completion; ``auto_sync()`` adds a +trailing barrier after a scratch-using operation. The provider queries the +actual CUB layout, and Warp storage is replicated for each group. Scan itself +does not promise completion of unrelated global Load/Store accesses. Retain +the required group barriers when reusing a destination or explicit scratch. diff --git a/rust/cuda-coop-codegen/src/lib.rs b/rust/cuda-coop-codegen/src/lib.rs index 006776db9414..ede8ed4450b6 100644 --- a/rust/cuda-coop-codegen/src/lib.rs +++ b/rust/cuda-coop-codegen/src/lib.rs @@ -3,7 +3,7 @@ //! Native, host-only CUB specialization for Rust cooperative algorithms. //! -//! This private development crate compiles block and Warp Load/Store providers with +//! This private development crate compiles block and Warp Load/Store and Scan 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. @@ -28,6 +28,14 @@ //! Transpose at power-of-two widths from 1 through 32, within complete physical //! warps. Their scratch layout describes one group, not the complete block. //! +//! [`ScanRequest`] selects a separate numeric Scan family: all ten dtypes, seven +//! built-in operators (bitwise for integers only), inclusive or exclusive output, +//! and Raking/RakingMemoize/WarpScans for Blocks. Warp Scan accepts one scalar +//! per lane with implicit algorithm selection. Scan checks every input validity +//! bit, preserves the input, and initializes every fresh output slot. Seeded +//! exclusive forms carry one runtime value of exactly the input dtype. +//! [`Provider::compile_native`] batches both families in one native compilation. +//! //! [`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 @@ -56,6 +64,8 @@ #![deny(unsafe_code)] mod headers; +mod scan; +pub use scan::{ScanAlgorithm, ScanOperation, ScanOperator, ScanRecipe, ScanRequest}; #[allow(unsafe_code)] mod nvrtc; pub mod oxide; @@ -431,6 +441,48 @@ impl Request { } } +/// One member of an ordered native batch, with family-specific validated inputs. +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub enum NativeRequest { + /// Existing Load/Store operation and layout algorithm. + LoadStore(Request), + /// Numeric Scan with preserved input and fresh initialized output. + Scan(ScanRequest), +} + +impl NativeRequest { + fn source(self, symbol: &str) -> String { + match self { + Self::LoadStore(request) => source(request, symbol), + Self::Scan(request) => request.source(symbol), + } + } + + fn token(self) -> String { + match self { + Self::LoadStore(request) => request.operation.token().into(), + Self::Scan(request) => request.recipe.operation().token(), + } + } + + fn scratch_free(self) -> bool { + match self { + Self::LoadStore(request) => request.algorithm.scratch_free(), + Self::Scan(_) => false, + } + } + + fn mutates_payload(self) -> bool { + match self { + Self::LoadStore(request) => { + !request.operation.is_load() && !request.algorithm.scratch_free() + } + // Scan writes a fresh output, never rearranges its input payload. + Self::Scan(_) => false, + } + } +} + /// Byte size and minimum byte alignment of a storage region. #[derive(Clone, Copy, Debug, Eq, PartialEq)] pub struct Layout { @@ -444,6 +496,10 @@ pub struct Layout { /// /// # Device ABI obligations /// +/// The memory-tile and payload-layout requirements below describe Load/Store. +/// Scan uses lane-private input and fresh output; its ABI and validity contract +/// are documented by [`ScanRequest`]. +/// /// Every thread in the selected group executes the provider in converged control /// flow. The group is the exact block, or one aligned Warp of the selected width. /// Memory is device-accessible and properly aligned; input/output spans @@ -585,6 +641,24 @@ impl Provider { /// 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 { + let requests: Vec<_> = requests + .iter() + .copied() + .map(NativeRequest::LoadStore) + .collect(); + self.compile_native(&requests, target) + } + + /// Compile one ordered mixed Load/Store and Scan batch in a single NVRTC program. + /// + /// Per-specialization identity is independent of unrelated batch members. + /// Scan scratch describes one group and includes the actual CUB TempStorage, + /// even when its C++ empty type occupies one byte. The caller plans replication. + /// + /// # Errors + /// Rejects empty/duplicate requests and malformed targets as [`Self::compile`] + /// does; native compilation and layout decoding errors retain their diagnostics. + pub fn compile_native(&self, requests: &[NativeRequest], target: &str) -> Result { if !target.starts_with("compute_") || !target[8..].chars().all(|c| c.is_ascii_alphanumeric()) || target.len() <= 8 @@ -617,6 +691,7 @@ impl Provider { include_bytes!("headers.rs"), include_bytes!("nvrtc.rs"), include_bytes!("block_io.cuh"), + include_bytes!("scan.rs"), ] { hash_field(&mut inputs, bytes); } @@ -632,13 +707,13 @@ impl Provider { let mut hash = inputs.clone(); hash_field( &mut hash, - self::source(request, "cuda_coop_specialization").as_bytes(), + request.source("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)); + let symbol = format!("cuda_coop_{}_{}", request.token(), suffix); + source.push_str(&request.source(&symbol)); expressions.push(format!("&layout_probe")); expressions.push(format!("&layout_probe")); specializations.push(Specialization { @@ -652,8 +727,8 @@ impl Provider { size: 0, alignment: 1, }, - scratch_free: request.algorithm.scratch_free(), - mutates_payload: !request.operation.is_load() && !request.algorithm.scratch_free(), + scratch_free: request.scratch_free(), + mutates_payload: request.mutates_payload(), }); } let output = self diff --git a/rust/cuda-coop-codegen/src/oxide.rs b/rust/cuda-coop-codegen/src/oxide.rs index 367375864e51..12807ec40b78 100644 --- a/rust/cuda-coop-codegen/src/oxide.rs +++ b/rust/cuda-coop-codegen/src/oxide.rs @@ -10,12 +10,15 @@ //! 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 raw i32 and typed validity recipes for all six block algorithms -//! documented by [`Operation`]. +//! The protocol supports raw i32 and typed validity Load/Store recipes documented +//! by [`Operation`], and numeric Block/scalar Warp Scan through [`ScanRecipe`]. //! 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 crate::{ + Algorithm, Artifact, Dtype, Error, Layout, NativeRequest, Operation, Provider, Request, Result, + ScanRecipe, ScanRequest, +}; use serde::{Deserialize, Serialize}; use std::io::{Read, Write}; use std::path::{Path, PathBuf}; @@ -116,7 +119,29 @@ struct SpecializationRequest { } impl SpecializationRequest { - fn native(&self) -> Result { + fn native(&self) -> Result { + if self.operation.starts_with("scan_") { + let recipe = ScanRecipe::parse(&self.operation, &self.algorithm)?; + let dtype = self.dtype.dtype()?; + let request = match (recipe, self.warp_width) { + (ScanRecipe::Block(operation, algorithm), None) => ScanRequest::block( + dtype, + operation, + algorithm, + self.block, + self.items_per_thread, + )?, + (ScanRecipe::Warp(operation), Some(width)) if self.items_per_thread == 1 => { + ScanRequest::warp(dtype, operation, self.block, width)? + } + _ => { + return Err(Error::message( + "Scan recipe disagrees with Warp scope or scalar extent", + )); + } + }; + return Ok(NativeRequest::Scan(request)); + } let operation = match self.operation.as_str() { "load_full" => Operation::LoadFull, "load_filled" => Operation::LoadFilled, @@ -156,8 +181,8 @@ impl SpecializationRequest { self.items_per_thread, )?; match self.warp_width { - Some(width) => request.with_warp_width(width), - None => Ok(request), + Some(width) => request.with_warp_width(width).map(NativeRequest::LoadStore), + None => Ok(NativeRequest::LoadStore(request)), } } } @@ -199,7 +224,23 @@ struct BatchResponse { specializations: Vec, } -fn arguments(request: Request) -> Vec { +fn arguments(request: NativeRequest) -> Vec { + let request = match request { + NativeRequest::Scan(request) => { + let dtype = Scalar::from(request.dtype); + let mut arguments = vec![ + Argument::ConstPointer(dtype), + Argument::MutPointer(dtype), + Argument::ConstPointer(Scalar::Unsigned(8)), + Argument::MutPointer(Scalar::Unsigned(8)), + ]; + if request.recipe.operation().seeded() { + arguments.push(Argument::Value(dtype)); + } + return arguments; + } + NativeRequest::LoadStore(request) => request, + }; let operation = request.operation(); let dtype = Scalar::from(request.dtype()); let shared_payload = @@ -309,6 +350,29 @@ pub fn registration( .map_err(|error| Error::message(format!("cannot encode provider registration: {error}"))) } +/// Encode one Scan recipe for the existing dependency provider executable. +/// +/// The Block recipe selects a Block Scan algorithm; the Warp recipe has no +/// explicit algorithm. Dtype, geometry and extent arrive with specialization. +/// Seeded exclusive operations take one runtime value of exactly that dtype; +/// inclusive and zero-seeded sum take no scalar controls. +/// +/// # Errors +/// Rejects a nonabsolute executable or a path that cannot be encoded as JSON. +pub fn scan_registration(executable: &Path, recipe: ScanRecipe) -> 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: &recipe.operation().token(), + algorithm: recipe.algorithm_token(), + }) + .map_err(|error| Error::message(format!("cannot encode provider registration: {error}"))) +} + #[derive(Deserialize)] #[serde(deny_unknown_fields)] struct InputRequest { @@ -420,11 +484,15 @@ pub fn serve(checkout: &Path, input: impl Read, mut output: impl Write) -> Resul .map(SpecializationRequest::native) .collect::>()?; let artifact = Provider::from_checkout(checkout, &request.cuda_root)? - .compile(&requests, &request.target)?; + .compile_native(&requests, &request.target)?; write_response(&requests, artifact, &mut output) } -fn write_response(requests: &[Request], artifact: Artifact, output: &mut impl Write) -> Result<()> { +fn write_response( + requests: &[NativeRequest], + 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", @@ -469,6 +537,136 @@ fn write_response(requests: &[Request], artifact: Artifact, output: &mut impl Wr #[cfg(test)] mod tests { use super::*; + use crate::{ScanAlgorithm, ScanOperation, ScanOperator}; + + #[test] + fn scan_registration_preserves_typed_seed_and_empty_controls() { + for (dtype, recipe) in [ + ( + Dtype::I8, + ScanRecipe::Block( + ScanOperation::Exclusive(ScanOperator::Min), + ScanAlgorithm::Raking, + ), + ), + ( + Dtype::U16, + ScanRecipe::Warp(ScanOperation::Exclusive(ScanOperator::Sum)), + ), + ( + Dtype::F64, + ScanRecipe::Warp(ScanOperation::Inclusive(ScanOperator::Max)), + ), + ( + Dtype::I32, + ScanRecipe::Block(ScanOperation::ExclusiveSum, ScanAlgorithm::WarpScans), + ), + ] { + let registration: serde_json::Value = + serde_json::from_str(&scan_registration(Path::new("/provider"), recipe).unwrap()) + .unwrap(); + let scalar = Scalar::from(dtype); + let request = SpecializationRequest { + operation: registration["operation"].as_str().unwrap().into(), + algorithm: registration["algorithm"].as_str().unwrap().into(), + dtype: scalar, + block: [4, 4, 4], + items_per_thread: 1, + warp_width: matches!(recipe, ScanRecipe::Warp(_)).then_some(8), + } + .native() + .unwrap(); + let mut expected = vec![ + Argument::ConstPointer(scalar), + Argument::MutPointer(scalar), + Argument::ConstPointer(Scalar::Unsigned(8)), + Argument::MutPointer(Scalar::Unsigned(8)), + ]; + if recipe.operation().seeded() { + expected.push(Argument::Value(scalar)); + } + assert_eq!(arguments(request), expected); + assert_eq!(registration["provider"], "cuda_coop"); + } + } + + #[test] + fn rejects_invalid_scan_before_opening_toolkit() { + let base = serde_json::json!({"version":1,"target":"compute_120","cuda_root":"/absent", "requests":[{ + "operation":"scan_inclusive_bit_xor_validity_bits_v1","algorithm":"raking", + "dtype":{"kind":"signed","bits":32},"block":[8,4,2],"items_per_thread":9,"warp_width":null + }]}); + for (path, value, diagnostic) in [ + ( + "/requests/0/dtype", + serde_json::json!({"kind":"float","bits":32}), + "integer", + ), + ( + "/requests/0/algorithm", + serde_json::json!("transpose"), + "Scan algorithm", + ), + ( + "/requests/0/operation", + serde_json::json!("scan_exclusive_min_validity_bits_v1"), + "seed form", + ), + ( + "/requests/0/operation", + serde_json::json!("scan_inclusive_sum_seeded_validity_bits_v1"), + "Scan operator", + ), + ( + "/requests/0/items_per_thread", + serde_json::json!(0), + "positive items", + ), + ( + "/requests/0/warp_width", + serde_json::json!(8), + "scope or scalar extent", + ), + ] { + let mut request = base.clone(); + *request.pointer_mut(path).unwrap() = value; + let error = serve( + Path::new("/absent"), + request.to_string().as_bytes(), + Vec::new(), + ) + .unwrap_err(); + assert!(error.to_string().contains(diagnostic), "{error}"); + } + for (block, items, width, diagnostic) in [ + ([48, 1, 1], 1, 8, "complete physical warps"), + ([64, 1, 1], 4, 8, "scalar extent"), + ([64, 1, 1], 1, 3, "power of two"), + ] { + let mut request = base.clone(); + request["requests"][0]["algorithm"] = "warp".into(); + request["requests"][0]["block"] = serde_json::json!(block); + request["requests"][0]["items_per_thread"] = items.into(); + request["requests"][0]["warp_width"] = width.into(); + let error = serve( + Path::new("/absent"), + request.to_string().as_bytes(), + Vec::new(), + ) + .unwrap_err(); + assert!(error.to_string().contains(diagnostic), "{error}"); + } + assert!( + ScanRequest::block( + Dtype::I32, + ScanOperation::ExclusiveSum, + ScanAlgorithm::WarpScans, + [8, 3, 2], + 1 + ) + .is_err() + ); + } #[test] fn typed_validity_abi_preserves_width_and_separate_flags() { @@ -558,7 +756,7 @@ mod tests { if operation.counted() { expected.push(Argument::Value(Scalar::Signed(32))); } - assert_eq!(arguments(request), expected); + assert_eq!(arguments(NativeRequest::LoadStore(request)), expected); } } for scalar in [Scalar::Signed(1), Scalar::Unsigned(128), Scalar::Float(16)] { @@ -632,7 +830,9 @@ mod tests { items_per_thread: 4, warp_width: None, }; - let native = request.native().unwrap(); + let NativeRequest::LoadStore(native) = request.native().unwrap() else { + panic!("expected Load/Store") + }; assert_eq!(native.block_shape(), [8, 4, 4]); assert_eq!(native.block_threads(), 128); assert_eq!(native.dtype(), request.dtype.dtype().unwrap()); diff --git a/rust/cuda-coop-codegen/src/scan.rs b/rust/cuda-coop-codegen/src/scan.rs new file mode 100644 index 000000000000..047187db6169 --- /dev/null +++ b/rust/cuda-coop-codegen/src/scan.rs @@ -0,0 +1,351 @@ +// Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +//! Validated Scan recipes, separate from Load/Store algorithms. + +use crate::{Dtype, Error, Result}; + +/// Static associative operator used by a native Scan specialization. +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub enum ScanOperator { + /// Addition in the input dtype. + Sum, + /// Multiplication in the input dtype. + Multiplies, + /// Minimum. + Min, + /// Maximum. + Max, + /// Integer bitwise AND. + BitAnd, + /// Integer bitwise OR. + BitOr, + /// Integer bitwise XOR. + BitXor, +} + +impl ScanOperator { + fn token(self) -> &'static str { + match self { + Self::Sum => "sum", + Self::Multiplies => "multiplies", + Self::Min => "min", + Self::Max => "max", + Self::BitAnd => "bit_and", + Self::BitOr => "bit_or", + Self::BitXor => "bit_xor", + } + } + + fn parse(token: &str) -> Result { + match token { + "sum" => Ok(Self::Sum), + "multiplies" => Ok(Self::Multiplies), + "min" => Ok(Self::Min), + "max" => Ok(Self::Max), + "bit_and" => Ok(Self::BitAnd), + "bit_or" => Ok(Self::BitOr), + "bit_xor" => Ok(Self::BitXor), + _ => Err(Error::message("unsupported Scan operator")), + } + } + + fn functor(self) -> &'static str { + match self { + Self::Sum => "cuda::std::plus", + Self::Multiplies => "cuda::std::multiplies", + Self::Min => "cuda::minimum", + Self::Max => "cuda::maximum", + Self::BitAnd => "cuda::std::bit_and", + Self::BitOr => "cuda::std::bit_or", + Self::BitXor => "cuda::std::bit_xor", + } + } +} + +/// Scan mode and runtime seed presence. +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub enum ScanOperation { + /// Exclusive sum with typed zero, without a runtime seed argument. + ExclusiveSum, + /// Exclusive scan with one runtime seed of exactly the input dtype. + Exclusive(ScanOperator), + /// Inclusive scan without a seed argument. + Inclusive(ScanOperator), +} + +impl ScanOperation { + pub(crate) fn seeded(self) -> bool { + matches!(self, Self::Exclusive(_)) + } + + pub(crate) fn token(self) -> String { + match self { + Self::ExclusiveSum => "scan_exclusive_sum_validity_bits_v1".into(), + Self::Exclusive(op) => { + format!("scan_exclusive_{}_seeded_validity_bits_v1", op.token()) + } + Self::Inclusive(op) => format!("scan_inclusive_{}_validity_bits_v1", op.token()), + } + } + + fn parse(token: &str) -> Result { + if token == "scan_exclusive_sum_validity_bits_v1" { + return Ok(Self::ExclusiveSum); + } + if let Some(op) = token + .strip_prefix("scan_exclusive_") + .and_then(|s| s.strip_suffix("_seeded_validity_bits_v1")) + { + return Ok(Self::Exclusive(ScanOperator::parse(op)?)); + } + if let Some(op) = token + .strip_prefix("scan_inclusive_") + .and_then(|s| s.strip_suffix("_validity_bits_v1")) + { + return Ok(Self::Inclusive(ScanOperator::parse(op)?)); + } + Err(Error::message("unsupported Scan operation or seed form")) + } +} + +/// CUB algorithm selected for a complete Block Scan. +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub enum ScanAlgorithm { + /// Shared-memory raking, the default Block algorithm. + Raking, + /// Raking with register-cached partial reductions. + RakingMemoize, + /// Combine physical-warp scans; requires a block size divisible by 32. + WarpScans, +} + +/// Registration identity independent of dtype, geometry and payload extent. +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub enum ScanRecipe { + /// Complete Block operation with its selected Block algorithm. + Block(ScanOperation, ScanAlgorithm), + /// Scalar physical/logical Warp operation with implicit algorithm selection. + Warp(ScanOperation), +} + +impl ScanRecipe { + /// Operation and seed presence selected by this recipe. + pub fn operation(self) -> ScanOperation { + match self { + Self::Block(operation, _) | Self::Warp(operation) => operation, + } + } + + pub(crate) fn algorithm_token(self) -> &'static str { + match self { + Self::Block(_, ScanAlgorithm::Raking) => "raking", + Self::Block(_, ScanAlgorithm::RakingMemoize) => "raking_memoize", + Self::Block(_, ScanAlgorithm::WarpScans) => "warp_scans", + Self::Warp(_) => "warp", + } + } + + pub(crate) fn parse(operation: &str, algorithm: &str) -> Result { + let operation = ScanOperation::parse(operation)?; + match algorithm { + "raking" => Ok(Self::Block(operation, ScanAlgorithm::Raking)), + "raking_memoize" => Ok(Self::Block(operation, ScanAlgorithm::RakingMemoize)), + "warp_scans" => Ok(Self::Block(operation, ScanAlgorithm::WarpScans)), + "warp" => Ok(Self::Warp(operation)), + _ => Err(Error::message("unsupported Scan algorithm")), + } + } +} + +/// Validated numeric Scan specialization with exact root geometry. +/// +/// Native calls preserve all input values and validity bits and initialize every +/// fresh output slot. Their ABI is `(const T* input, T* output, u8* scratch, +/// const u8* input_bits, u8* output_bits[, T seed])`. Input and output must be +/// live, disjoint lane-local storage of the declared extent; validity is packed +/// low-bit-first. Every input bit is checked before numeric input is read. +/// All group members must participate, seeds must be group-uniform, and scratch +/// ownership/completion belongs to the checked compiler integration. +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub struct ScanRequest { + pub(crate) recipe: ScanRecipe, + pub(crate) dtype: Dtype, + block: [u32; 3], + items: u32, + width: Option, +} + +impl ScanRequest { + /// Validate a complete Block with positive fixed payload extent. + /// + /// # Errors + /// Rejects zero/overflowing axes, more than 1024 threads, zero items or a + /// total extent exceeding i32. WarpScans requires complete physical warps; + /// bitwise operators require an integer dtype. + pub fn block( + dtype: Dtype, + operation: ScanOperation, + algorithm: ScanAlgorithm, + block: [u32; 3], + items: u32, + ) -> Result { + Self::validate( + dtype, + ScanRecipe::Block(operation, algorithm), + block, + items, + None, + ) + } + + /// Validate one scalar per lane in a complete physical/logical Warp. + /// + /// # Errors + /// Rejects invalid Block geometry, incomplete physical warps, widths other + /// than 1/2/4/8/16/32, or a bitwise operator with a floating dtype. + pub fn warp( + dtype: Dtype, + operation: ScanOperation, + block: [u32; 3], + width: u32, + ) -> Result { + Self::validate(dtype, ScanRecipe::Warp(operation), block, 1, Some(width)) + } + + fn validate( + dtype: Dtype, + recipe: ScanRecipe, + block: [u32; 3], + items: u32, + width: Option, + ) -> Result { + let threads = block + .into_iter() + .try_fold(1_u32, u32::checked_mul) + .filter(|n| (1..=1024).contains(n)); + let Some(threads) = threads else { + return Err(Error::message( + "Scan requires 1..=1024 block threads with positive axes", + )); + }; + if items == 0 || u64::from(threads) * u64::from(items) > i32::MAX as u64 { + return Err(Error::message( + "Scan requires positive items and tile extent <= i32::MAX", + )); + } + if matches!( + recipe, + ScanRecipe::Block(_, ScanAlgorithm::WarpScans) | ScanRecipe::Warp(_) + ) && threads % 32 != 0 + { + return Err(Error::message( + "Warp Scan and Block WarpScans require complete physical warps", + )); + } + if width.is_some_and(|w| w == 0 || w > 32 || !w.is_power_of_two()) { + return Err(Error::message( + "Warp Scan width must be a power of two in 1..=32", + )); + } + if matches!(dtype, Dtype::F32 | Dtype::F64) + && matches!( + recipe.operation(), + ScanOperation::Exclusive( + ScanOperator::BitAnd | ScanOperator::BitOr | ScanOperator::BitXor + ) | ScanOperation::Inclusive( + ScanOperator::BitAnd | ScanOperator::BitOr | ScanOperator::BitXor + ) + ) + { + return Err(Error::message( + "bitwise Scan operators require integer elements", + )); + } + Ok(Self { + recipe, + dtype, + block, + items, + width, + }) + } + + pub(crate) fn source(self, symbol: &str) -> String { + let (dtype, _) = self.dtype.native(); + let [x, y, z] = self.block; + let items = self.items; + let primitive = match self.recipe { + ScanRecipe::Block(_, algorithm) => { + let algorithm = match algorithm { + ScanAlgorithm::Raking => "BLOCK_SCAN_RAKING", + ScanAlgorithm::RakingMemoize => "BLOCK_SCAN_RAKING_MEMOIZE", + ScanAlgorithm::WarpScans => "BLOCK_SCAN_WARP_SCANS", + }; + format!("cub::BlockScan") + } + ScanRecipe::Warp(_) => format!("cub::WarpScan", self.width.unwrap()), + }; + let controls = if self.recipe.operation().seeded() { + format!(", {dtype} seed") + } else { + String::new() + }; + let operands = if items == 1 { + "local_input[0], local_output[0]" + } else { + "local_input, local_output" + }; + let call = match self.recipe.operation() { + ScanOperation::ExclusiveSum => format!("operation.ExclusiveSum({operands});"), + ScanOperation::Exclusive(op) => format!( + "operation.ExclusiveScan({operands}, seed, {}{{}});", + op.functor() + ), + ScanOperation::Inclusive(op) => { + format!("operation.InclusiveScan({operands}, {}{{}});", op.functor()) + } + }; + let tail = if items % 8 == 0 { + String::new() + } else { + format!( + "output_bits[{}] = {};", + items / 8, + (1_u32 << (items % 8)) - 1 + ) + }; + format!( + r#" +#include +#include +#include +#include +namespace {symbol}_detail {{ +using value_type = {dtype}; +using primitive = {primitive}; +using scratch_storage = primitive::TempStorage; +}} +extern "C" __device__ void {symbol}(const {dtype}* input, {dtype}* output, + unsigned char* scratch, const unsigned char* input_bits, unsigned char* output_bits{controls}) {{ + using namespace {symbol}_detail; +#pragma unroll + for (int i = 0; i < {items}; ++i) + if ((input_bits[i / 8] & (1u << (i % 8))) == 0) __trap(); + value_type local_input[{items}]; + value_type local_output[{items}]; +#pragma unroll + for (int i = 0; i < {items}; ++i) local_input[i] = input[i]; + primitive operation{{*reinterpret_cast(scratch)}}; + {call} +#pragma unroll + for (int i = 0; i < {items}; ++i) output[i] = local_output[i]; +#pragma unroll + for (int i = 0; i < {full_bytes}; ++i) output_bits[i] = 255; + {tail} +}} +"#, + full_bytes = items / 8 + ) + } +} diff --git a/rust/cuda-coop-codegen/tests/native.rs b/rust/cuda-coop-codegen/tests/native.rs index d29a5dd6ad7c..59acafc02b5f 100644 --- a/rust/cuda-coop-codegen/tests/native.rs +++ b/rust/cuda-coop-codegen/tests/native.rs @@ -3,9 +3,140 @@ //! Real native compiler coverage; ordinary host tests need no CUDA installation. -use cuda_coop_codegen::{Algorithm, Dtype, Layout, Operation, Provider, Request}; +use cuda_coop_codegen::{ + Algorithm, Dtype, Layout, NativeRequest, Operation, Provider, Request, ScanAlgorithm, + ScanOperation, ScanOperator, ScanRequest, +}; use std::{env, path::Path, time::Instant}; +#[test] +#[ignore = "requires CUDA_HOME and a Linux CUDA toolkit with NVRTC LTO support"] +fn native_mixed_scan_recipes() -> 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)?; + // Each operator is compiled in both modes; types/algorithms are spread + // across those cases rather than multiplying independent dimensions. + let types = [ + Dtype::I8, + Dtype::U8, + Dtype::I16, + Dtype::U16, + Dtype::I32, + Dtype::U32, + Dtype::I64, + Dtype::U64, + Dtype::F32, + Dtype::F64, + Dtype::I8, + Dtype::I16, + Dtype::U16, + Dtype::U64, + ]; + let operators = [ + ScanOperator::Sum, + ScanOperator::Multiplies, + ScanOperator::Min, + ScanOperator::Max, + ScanOperator::BitAnd, + ScanOperator::BitOr, + ScanOperator::BitXor, + ]; + let algorithms = [ + ScanAlgorithm::Raking, + ScanAlgorithm::RakingMemoize, + ScanAlgorithm::WarpScans, + ]; + let mut requests = Vec::new(); + for (index, dtype) in types.into_iter().enumerate() { + let op = operators[index % operators.len()]; + let operation = if index < operators.len() { + ScanOperation::Exclusive(op) + } else { + ScanOperation::Inclusive(op) + }; + requests.push(NativeRequest::Scan(ScanRequest::block( + dtype, + operation, + algorithms[index % algorithms.len()], + [4, 4, 4], + if index % 2 == 0 { 1 } else { 4 }, + )?)); + } + requests.push(NativeRequest::Scan(ScanRequest::block( + Dtype::I8, + ScanOperation::ExclusiveSum, + ScanAlgorithm::Raking, + [8, 4, 1], + 9, + )?)); + for (index, width) in [1, 2, 4, 8, 16, 32].into_iter().enumerate() { + requests.push(NativeRequest::Scan(ScanRequest::warp( + if index % 2 == 0 { + Dtype::I8 + } else { + Dtype::F64 + }, + match index % 3 { + 0 => ScanOperation::ExclusiveSum, + 1 => ScanOperation::Exclusive(ScanOperator::Min), + _ => ScanOperation::Inclusive(ScanOperator::Max), + }, + [4, 4, 4], + width, + )?)); + } + for operation in [Operation::LoadFullValidity, Operation::StoreFullValidity] { + requests.push(NativeRequest::LoadStore(Request::new_typed( + Dtype::I32, + operation, + Algorithm::Direct, + 64, + 4, + )?)); + } + let artifact = provider.compile_native(&requests, &target)?; + assert_eq!(artifact.specializations.len(), requests.len()); + for (request, entry) in requests.iter().zip(&artifact.specializations) { + if matches!(request, NativeRequest::Scan(_)) { + assert!(!entry.scratch_free); + assert_eq!(entry.scratch, entry.cub_layout); + assert!(entry.scratch.size > 0 && entry.scratch.alignment.is_power_of_two()); + assert!(!entry.mutates_payload, "Scan must preserve its input"); + } else { + assert!(entry.scratch_free); + assert_eq!(entry.scratch.size, 0); + } + eprintln!( + "{request:?}: CUB {:?}, scratch {:?}", + entry.cub_layout, entry.scratch + ); + } + let single = provider.compile_native(&requests[..1], &target)?; + assert_eq!( + single.specializations[0].identity, + artifact.specializations[0].identity + ); + assert_eq!( + single.specializations[0].symbol, + artifact.specializations[0].symbol + ); + assert!( + provider + .compile_native(&[requests[0]; 2], &target) + .unwrap_err() + .to_string() + .contains("duplicate") + ); + eprintln!( + "{} mixed Scan/Load/Store specializations, {} LTO bytes", + requests.len(), + artifact.ltoir.len() + ); + Ok(()) +} + #[test] fn rejects_invalid_shapes() { for (threads, items) in [(0, 1), (1025, 1), (128, 0), (1024, u32::MAX)] { diff --git a/rust/cuda-coop/build.rs b/rust/cuda-coop/build.rs index 1052620aed55..6d85c64cc097 100644 --- a/rust/cuda-coop/build.rs +++ b/rust/cuda-coop/build.rs @@ -1,8 +1,10 @@ // 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}; +//! Register the owning host provider for native Load, Store and Scan. +use cuda_coop_codegen::{ + Algorithm, Operation, ScanAlgorithm, ScanOperation, ScanOperator, ScanRecipe, oxide, +}; fn main() { oxide::build_script().expect("cuda.coop native provider registration"); @@ -13,6 +15,67 @@ fn main() { return; } let executable = std::env::current_exe().expect("cuda.coop provider executable"); + for (operation, name) in [ + (ScanOperation::ExclusiveSum, "EXCLUSIVE_ZERO"), + (ScanOperation::Exclusive(ScanOperator::Sum), "EXCLUSIVE_SUM"), + ( + ScanOperation::Exclusive(ScanOperator::Multiplies), + "EXCLUSIVE_MULTIPLIES", + ), + (ScanOperation::Exclusive(ScanOperator::Min), "EXCLUSIVE_MIN"), + (ScanOperation::Exclusive(ScanOperator::Max), "EXCLUSIVE_MAX"), + ( + ScanOperation::Exclusive(ScanOperator::BitAnd), + "EXCLUSIVE_BIT_AND", + ), + ( + ScanOperation::Exclusive(ScanOperator::BitOr), + "EXCLUSIVE_BIT_OR", + ), + ( + ScanOperation::Exclusive(ScanOperator::BitXor), + "EXCLUSIVE_BIT_XOR", + ), + (ScanOperation::Inclusive(ScanOperator::Sum), "INCLUSIVE_SUM"), + ( + ScanOperation::Inclusive(ScanOperator::Multiplies), + "INCLUSIVE_MULTIPLIES", + ), + (ScanOperation::Inclusive(ScanOperator::Min), "INCLUSIVE_MIN"), + (ScanOperation::Inclusive(ScanOperator::Max), "INCLUSIVE_MAX"), + ( + ScanOperation::Inclusive(ScanOperator::BitAnd), + "INCLUSIVE_BIT_AND", + ), + ( + ScanOperation::Inclusive(ScanOperator::BitOr), + "INCLUSIVE_BIT_OR", + ), + ( + ScanOperation::Inclusive(ScanOperator::BitXor), + "INCLUSIVE_BIT_XOR", + ), + ] { + for (recipe, prefix) in [ + ( + ScanRecipe::Block(operation, ScanAlgorithm::Raking), + "RAKING", + ), + ( + ScanRecipe::Block(operation, ScanAlgorithm::RakingMemoize), + "RAKING_MEMOIZE", + ), + ( + ScanRecipe::Block(operation, ScanAlgorithm::WarpScans), + "WARP_SCANS", + ), + (ScanRecipe::Warp(operation), "WARP"), + ] { + let value = + oxide::scan_registration(&executable, recipe).expect("cuda.coop scan registration"); + println!("cargo::rustc-env=COOP_SCAN_{prefix}_{name}={value}"); + } + } for (operation, name) in [ (Operation::LoadFullValidity, "LOAD_FULL"), (Operation::LoadMaskedValidity, "LOAD_MASKED"), diff --git a/rust/cuda-coop/examples/scan/.gitignore b/rust/cuda-coop/examples/scan/.gitignore new file mode 100644 index 000000000000..74bffa8af4a4 --- /dev/null +++ b/rust/cuda-coop/examples/scan/.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/scan/Cargo.lock b/rust/cuda-coop/examples/scan/Cargo.lock new file mode 100644 index 000000000000..565676cddbda --- /dev/null +++ b/rust/cuda-coop/examples/scan/Cargo.lock @@ -0,0 +1,1019 @@ +# 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-scan" +version = "0.0.0" +dependencies = [ + "coop-scan-helper", + "cuda-async", + "cuda-coop", + "cuda-core", + "cuda-device", + "cuda-host", +] + +[[package]] +name = "coop-scan-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=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "libnvvm-sys", + "nvjitlink-sys", + "serde", + "sha2", + "thiserror 2.0.21", +] + +[[package]] +name = "cuda-async" +version = "0.5.0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +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.5.0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +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.5.0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "anyhow", + "cuda-bindings", + "cuda-core-derive", + "half", + "oxide-artifacts", +] + +[[package]] +name = "cuda-core-derive" +version = "0.5.0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +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=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "cuda-macros", +] + +[[package]] +name = "cuda-host" +version = "0.2.1" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "cuda-artifact-finalizer", + "cuda-async", + "cuda-core", + "cuda-device", + "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=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +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=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" + +[[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=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +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=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +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=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "object", +] + +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +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=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" + +[[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=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" + +[[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/scan/Cargo.toml b/rust/cuda-coop/examples/scan/Cargo.toml new file mode 100644 index 000000000000..fe413c7daac5 --- /dev/null +++ b/rust/cuda-coop/examples/scan/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-scan" +version = "0.0.0" +edition = "2024" +publish = false + +[workspace] +members = ["helper"] + +[dependencies] +cuda-coop = { path = "../.." } +coop-scan-helper = { path = "helper" } +cuda-device = { git = "https://github.com/tpn/cuda-rust.git", rev = "96caed46c776eb49226137eac849e795b83b7ce0" } +cuda-host = { git = "https://github.com/tpn/cuda-rust.git", rev = "96caed46c776eb49226137eac849e795b83b7ce0", features = ["async"] } +cuda-core = { git = "https://github.com/tpn/cuda-rust.git", rev = "96caed46c776eb49226137eac849e795b83b7ce0" } +cuda-async = { git = "https://github.com/tpn/cuda-rust.git", rev = "96caed46c776eb49226137eac849e795b83b7ce0" } diff --git a/rust/cuda-coop/examples/scan/helper/Cargo.toml b/rust/cuda-coop/examples/scan/helper/Cargo.toml new file mode 100644 index 000000000000..9300b682f9fe --- /dev/null +++ b/rust/cuda-coop/examples/scan/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-scan-helper" +version = "0.0.0" +edition = "2024" +publish = false + +[dependencies] +cuda-coop = { path = "../../.." } diff --git a/rust/cuda-coop/examples/scan/helper/src/lib.rs b/rust/cuda-coop/examples/scan/helper/src/lib.rs new file mode 100644 index 000000000000..6b36174a4290 --- /dev/null +++ b/rust/cuda-coop/examples/scan/helper/src/lib.rs @@ -0,0 +1,39 @@ +// Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +//! A reusable, element-type-preserving block prefix helper. +#![no_std] +use cuda_coop as coop; + +/// Write exclusive sums for the valid prefix of one block tile. +/// +/// `zero` fills the missing tail before Scan, which requires every input slot +/// initialized. The caller supplies the additive identity for `T`. +pub fn prefix( + block: coop::ThreadGroup, + input: &[T], + output: &mut coop::DisjointSlice<'_, T, coop::CollectiveTiles>, + offset: usize, + valid: usize, + zero: T, +) { + let mut values = coop::ThreadData::<_, ITEMS_PER_THREAD>::new(); + coop::load( + block, + input, + &mut values, + coop::LoadOptions::default() + .offset(offset) + .valid_items(valid) + .oob_default(zero), + ); + let prefixes = coop::scan(block, &values, Default::default()); + coop::store( + block, + output, + &prefixes, + coop::StoreOptions::default() + .offset(offset) + .valid_items(valid), + ); +} diff --git a/rust/cuda-coop/examples/scan/src/main.rs b/rust/cuda-coop/examples/scan/src/main.rs new file mode 100644 index 000000000000..5fa3b659352d --- /dev/null +++ b/rust/cuda-coop/examples/scan/src/main.rs @@ -0,0 +1,75 @@ +// 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 prefix( + input: &[i32], + 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 + }; + coop_scan_helper::prefix::<_, ITEMS_PER_THREAD>( + block, + input, + &mut output, + offset, + remaining.min(tile), + 0_i32, + ); + } +} + +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: the typed declarations and embedded kernels compile together. + let module = unsafe { kernels::load_async(0)? }; + macro_rules! run { + ($items:literal) => {{ + let tile = 128 * $items; + let input: Vec = (0..3 * tile - 7).map(|i| (i % 7) as i32).collect(); + let mut expected = Vec::with_capacity(input.len()); + for chunk in input.chunks(tile) { + let mut prefix = 0; + for &value in chunk { + expected.push(prefix); + prefix += value; + } + } + let source = DeviceBuffer::from_host(&stream, &input)?; + let target = DeviceBuffer::from_host(&stream, &vec![-1_i32; input.len()])?; + let plan = module.prepare_prefix::<$items>(LaunchConfig1D::auto_shared(3, 128))?; + let (_, target) = module + .prefix_async_owned::<$items, _, _>(&plan, source, target) + .sync()?; + let mut output = vec![0_i32; input.len()]; + target.copy_to_host(&stream, &mut output)?; + stream.synchronize()?; + assert_eq!(output, expected); + println!("items/thread={}: {} prefixes checked", $items, output.len()); + }}; + } + run!(1); + run!(4); + Ok(()) +} diff --git a/rust/cuda-coop/src/group.rs b/rust/cuda-coop/src/group.rs index 2b5ad8bcfbca..67dcb34dc0ec 100644 --- a/rust/cuda-coop/src/group.rs +++ b/rust/cuda-coop/src/group.rs @@ -412,6 +412,7 @@ pub enum WarpsWithinBlock {} impl ThreadsWithinWarp { const POWER_OF_TWO: bool = N > 0 && N <= 32 && N.is_power_of_two(); const COMPLETE: bool = N > 0 && N <= 32 && 32 % N == 0; + const PREFIX: bool = !E && N >= 17 && N <= 31; const VALID: () = { assert!(N > 0 && N <= 32, "lane partition size must be in 1..=32"); assert!( @@ -430,6 +431,8 @@ impl sealed::Scope for ThreadsWithinWarp { fn member() -> u8 { if Self::COMPLETE { 1 + } else if Self::PREFIX { + library::warp_prefix_member_v1::() } else { ((thread_rank_block() % 32) < (32 / N as u32) * N as u32) as u8 } diff --git a/rust/cuda-coop/src/lib.rs b/rust/cuda-coop/src/lib.rs index 99b223393877..3df7818cdd5d 100644 --- a/rust/cuda-coop/src/lib.rs +++ b/rust/cuda-coop/src/lib.rs @@ -39,6 +39,11 @@ //! compiler-planned shared scratch. [`TempStorage`] selects explicit Block reuse //! and completion policies; [`StorageReservation`] exposes typed scratch through //! raw pointers. +//! [`scan`] and its four named forms return fresh numeric prefixes. Blocks +//! accept scalars or borrowed payloads; physical/logical Warp accepts scalars. +//! Mode, built-in operator and Block algorithm are selected statically through +//! [`ScanOptions`], with exclusive sum as the default. Scan preserves its input +//! and uses the same compiler-checked scratch and participation rules. //! Host payload tests need no GPU; //! executing hardware queries requires an Oxide-compiled CUDA kernel. @@ -47,6 +52,7 @@ mod alignment; mod group; mod load_store; +mod scan; mod temp_storage; mod thread_data; @@ -56,6 +62,12 @@ pub use load_store::{ Algorithm, Direct, LoadOptions, LoadSource, LoadStoreGroup, StoreOptions, StoreValue, Striped, Transpose, Vectorize, WarpTranspose, WarpTransposeTimesliced, load, store, }; +pub use scan::{ + BitAnd, BitOr, BitXor, Exclusive, ExclusiveScanOptions, Inclusive, InclusiveScanOptions, Max, + Min, Multiplies, Raking, RakingMemoize, ScanAlgorithm, ScanGroup, ScanMode, ScanOperator, + ScanOptions, ScanValue, Sum, SumScanOptions, WarpScans, exclusive_scan, exclusive_sum, + inclusive_scan, inclusive_sum, scan, +}; pub use temp_storage::{StorageReservation, TempStorage, TempStorageLike}; pub use thread_data::{Numeric, ThreadData, ThreadDataLike}; diff --git a/rust/cuda-coop/src/scan.rs b/rust/cuda-coop/src/scan.rs new file mode 100644 index 000000000000..78e2db1111f9 --- /dev/null +++ b/rust/cuda-coop/src/scan.rs @@ -0,0 +1,622 @@ +// Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +use core::marker::PhantomData; + +use crate::{ + Alignment, Block, GroupScope, Numeric, SupportedAlignment, TempStorage, TempStorageLike, + ThreadData, ThreadGroup, ThreadsWithinWarp, Warp, +}; +use cuda_device::library; + +mod boundary; + +/// Exclusive prefix order: the first result is the initial value. +#[derive(Clone, Copy, Debug, Default)] +pub struct Exclusive; + +/// Inclusive prefix order: each result includes the current input item. +#[derive(Clone, Copy, Debug, Default)] +pub struct Inclusive; + +macro_rules! operators { + ($(($name:ident, $kind:literal, $doc:literal)),+ $(,)?) => {$ ( + #[doc = $doc] + #[derive(Clone, Copy, Debug, Default)] + pub struct $name; + impl sealed::Operator for $name { const KIND: u8 = $kind; } + )+}; +} +operators!( + ( + Sum, + 0, + "Addition in the input numeric type; the default Scan operator." + ), + (Multiplies, 1, "Multiplication in the input numeric type."), + (Min, 2, "CUB minimum in the input numeric type."), + (Max, 3, "CUB maximum in the input numeric type."), + (BitAnd, 4, "Bitwise AND for integer Scan inputs."), + (BitOr, 5, "Bitwise OR for integer Scan inputs."), + (BitXor, 6, "Bitwise XOR for integer Scan inputs."), +); + +/// A sealed static Scan operator for numeric input type `T`. +/// +/// [`Sum`], [`Multiplies`], [`Min`] and [`Max`] accept every [`Numeric`] type. +/// [`BitAnd`], [`BitOr`] and [`BitXor`] accept the eight integer types only. +/// Generic helpers may carry this bound. Custom operators are not yet supported. +pub trait ScanOperator: sealed::Operator {} +impl ScanOperator for Sum {} +impl ScanOperator for Multiplies {} +impl ScanOperator for Min {} +impl ScanOperator for Max {} +macro_rules! integer_operators { + ($($t:ty),+ $(,)?) => {$ ( + impl ScanOperator<$t> for BitAnd {} + impl ScanOperator<$t> for BitOr {} + impl ScanOperator<$t> for BitXor {} + )+}; +} +integer_operators!(i8, u8, i16, u16, i32, u32, i64, u64); + +macro_rules! algorithms { + ($(($name:ident, $kind:literal, $doc:literal)),+ $(,)?) => {$ ( + #[doc = $doc] + #[derive(Clone, Copy, Debug, Default)] + pub struct $name; + impl sealed::Algorithm for $name { const KIND: u8 = $kind; } + impl ScanAlgorithm for $name {} + )+}; +} +algorithms!( + ( + Raking, + 1, + "CUB's raking Block Scan; selected when the algorithm is omitted." + ), + (RakingMemoize, 2, "CUB's memoized raking Block Scan."), + ( + WarpScans, + 3, + "CUB's warp-based Block Scan; requires a block size divisible by 32." + ), +); + +/// A sealed Block Scan algorithm or omitted selection `()`. +/// +/// Block resolves omission to [`Raking`]. Warp accepts only omission, so even +/// explicit `Raking` is rejected for Warp. This trait is separate from the +/// Load/Store [`Algorithm`](crate::Algorithm) selectors. +pub trait ScanAlgorithm: sealed::Algorithm {} +impl sealed::Algorithm for () { + const KIND: u8 = 0; +} +impl ScanAlgorithm for () {} + +/// A sealed supported combination of Scan mode, operator and initial value. +/// +/// `()` and [`Exclusive`] allow an omitted seed `()` only for [`Sum`], selecting +/// typed zero. A supplied seed has exactly type `T` and is allowed for every +/// supported operator. [`Inclusive`] allows only an omitted seed. Mode `()` is +/// the omitted selection used by default options and the fixed-name functions. +/// Runtime seed uniformity is checked by the device compiler, not this trait. +pub trait ScanMode, Seed>: sealed::Mode {} +macro_rules! exclusive_modes { + ($($mode:ty),+) => {$ ( + impl sealed::Mode for $mode { + type Uniform = (); + const KIND: u8 = 0; + #[inline(always)] fn controls(_: ()) {} + } + impl ScanMode for $mode {} + impl> sealed::Mode for $mode { + type Uniform = (T,); + const KIND: u8 = 1 + Op::KIND; + #[inline(always)] fn controls(seed: T) -> (T,) { (seed,) } + } + impl> ScanMode for $mode {} + )+}; +} +exclusive_modes!((), Exclusive); +impl> sealed::Mode for Inclusive { + type Uniform = (); + const KIND: u8 = 8 + Op::KIND; + #[inline(always)] + fn controls(_: ()) {} +} +impl> ScanMode for Inclusive {} + +/// A sealed group supporting a Scan algorithm and storage choice. +/// +/// Complete Blocks support scalar and payload Scan, with implicit or explicit +/// storage. Physical [`Warp`] and [`ThreadsWithinWarp`] widths 1/2/4/8/16/32 +/// support scalar Scan with omitted algorithm and implicit storage only. Warp +/// requires complete physical warps in the enclosing block. Membership, +/// participation and exact launch geometry are enforced by the compiler. +/// +/// Explicit Block algorithms are not accepted for Warp: +/// ```compile_fail,E0308 +/// use cuda_coop as coop; +/// fn invalid(value: i32) { +/// coop::scan(coop::this_warp(), value, +/// coop::ScanOptions::default().algorithm(coop::Raking)); +/// } +/// ``` +pub trait ScanGroup: GroupScope + sealed::Group {} +impl ScanGroup for Block {} +impl ScanGroup<(), ()> for Warp {} +impl sealed::Group for Block { + const WARP: bool = false; + #[inline(always)] + unsafe fn call &'static str + Copy, S, U: Copy, P: Copy, const I: usize>( + registration: R, + storage: S, + uniform: U, + payload: P, + ) { + // SAFETY: the private boundary supplies live typed payloads and the exact recipe. + unsafe { + library::checked_block_payload_call_v1::( + registration, + storage, + uniform, + payload, + ) + } + } +} +macro_rules! warp_group { + ($scope:ty, $width:expr $(, $e:ident)?) => { + impl<$(const $e: bool)?> sealed::Group for $scope { + const WARP: bool = true; + #[inline(always)] + unsafe fn call &'static str + Copy, S, U: Copy, P: Copy, const I: usize>( + registration: R, _storage: S, uniform: U, payload: P, + ) { + // SAFETY: ScanGroup permits implicit storage only; the compiler enforces the canonical width. + unsafe { library::checked_warp_payload_call_v1::(registration, uniform, payload) } + } + } + }; +} +warp_group!(Warp, 32); +macro_rules! logical_groups { + ($($width:literal),+) => {$ ( + impl ScanGroup<(), ()> for ThreadsWithinWarp<$width, E> {} + warp_group!(ThreadsWithinWarp<$width, E>, $width, E); + )+}; +} +logical_groups!(1, 2, 4, 8, 16, 32); + +/// Static selections and runtime seed/storage for [`scan`] and its named forms. +/// +/// `T` is inferred from the input. `M`, `Op` and `A` select mode, operator and +/// algorithm; `Seed` is either `()` or exactly `T`; `S` is implicit `()` or a +/// mutable [`TempStorage`] borrow. Defaults are omitted mode (exclusive for +/// `scan`), [`Sum`], omitted algorithm (Block [`Raking`]), zero seed and managed +/// storage. Only this canonical specialization implements Default, keeping +/// `Default::default()` unambiguous at all five entry points. +/// +/// Builders consume options and preserve other fields. Fixed-name functions +/// select their own mode and accept only omitted M. Warp rejects an explicit +/// algorithm or storage. Inclusive forms reject seeds; non-sum exclusive forms +/// require one. A seed must be uniform within the participating group. +/// +/// ```no_run +/// use cuda_coop as coop; +/// fn prefixes( +/// values: &coop::ThreadData, seed: i32, +/// ) -> coop::ThreadData { +/// coop::exclusive_scan(coop::this_block(), values, +/// coop::ScanOptions::default().scan_op(coop::Max).initial_value(seed)) +/// } +/// ``` +/// +/// All fixed-name defaults preserve a generic scalar's type: +/// ```no_run +/// use cuda_coop as coop; +/// fn scalar(value: T) { +/// let group = coop::this_warp(); +/// let _: T = coop::scan(group, value, Default::default()); +/// let _: T = coop::exclusive_sum(group, value, Default::default()); +/// let _: T = coop::inclusive_sum(group, value, Default::default()); +/// let _: T = coop::exclusive_scan(group, value, Default::default()); +/// let _: T = coop::inclusive_scan(group, value, Default::default()); +/// } +/// ``` +/// +/// Non-sum exclusive Scan requires a seed: +/// ```compile_fail,E0277 +/// use cuda_coop as coop; +/// fn invalid(value: i32) { +/// coop::scan(coop::this_block(), value, +/// coop::ScanOptions::default().scan_op(coop::Max)); +/// } +/// ``` +/// Inclusive Scan forbids a seed: +/// ```compile_fail,E0277 +/// use cuda_coop as coop; +/// fn invalid(value: i32) { +/// coop::scan(coop::this_block(), value, +/// coop::ScanOptions::default().mode(coop::Inclusive).initial_value(0)); +/// } +/// ``` +/// Bitwise operators require integer inputs: +/// ```compile_fail,E0277 +/// use cuda_coop as coop; +/// fn invalid(value: f32) { +/// coop::inclusive_scan(coop::this_block(), value, +/// coop::ScanOptions::default().scan_op(coop::BitAnd)); +/// } +/// ``` +/// Seeds preserve the input type: +/// ```compile_fail,E0308 +/// use cuda_coop as coop; +/// fn invalid(value: i16) { +/// coop::exclusive_scan(coop::this_block(), value, +/// coop::ScanOptions::default().initial_value(0_i32)); +/// } +/// ``` +pub struct ScanOptions { + scan_op: Op, + initial_value: Seed, + storage: S, + selection: PhantomData<(T, M, A)>, +} +impl Default for ScanOptions { + fn default() -> Self { + Self { + scan_op: Sum, + initial_value: (), + storage: (), + selection: PhantomData, + } + } +} +impl ScanOptions { + /// Select [`Exclusive`] or [`Inclusive`] for [`scan`]. + pub fn mode(self, _mode: M2) -> ScanOptions { + ScanOptions { + scan_op: self.scan_op, + initial_value: self.initial_value, + storage: self.storage, + selection: PhantomData, + } + } + /// Select a built-in operator, retaining its actual generic value. + pub fn scan_op>(self, op: Op2) -> ScanOptions { + ScanOptions { + scan_op: op, + initial_value: self.initial_value, + storage: self.storage, + selection: PhantomData, + } + } + /// Select a Block algorithm. Warp accepts only omitted selection `()`. + pub fn algorithm( + self, + _algorithm: A2, + ) -> ScanOptions { + ScanOptions { + scan_op: self.scan_op, + initial_value: self.initial_value, + storage: self.storage, + selection: PhantomData, + } + } + /// Supply the exact-T, group-uniform seed required by a non-sum exclusive Scan. + pub fn initial_value(self, value: T) -> ScanOptions { + ScanOptions { + scan_op: self.scan_op, + initial_value: value, + storage: self.storage, + selection: PhantomData, + } + } + /// Borrow explicit Block scratch with its existing reuse/completion policy. + /// + /// Warp rejects this choice. The operation consumes this options value, + /// permitting a fresh mutable reborrow for the next operation. + pub fn temp_storage( + self, + storage: &mut TempStorage, + ) -> ScanOptions> { + ScanOptions { + scan_op: self.scan_op, + initial_value: self.initial_value, + storage, + selection: PhantomData, + } + } +} + +/// Options shared by [`exclusive_sum`] and [`inclusive_sum`]. +pub type SumScanOptions = ScanOptions; +/// Operator, seed, algorithm and storage options for [`exclusive_scan`]. +pub type ExclusiveScanOptions = + ScanOptions; +/// Operator, algorithm and storage options for [`inclusive_scan`]; no seed. +pub type InclusiveScanOptions = ScanOptions; + +/// Scan numeric values in group order, returning a fresh fully initialized result. +/// +/// `T` and the result shape are inferred from `value`. A scalar returns T. +/// A shared or mutable Block payload borrow returns a fresh payload with the +/// same dtype, extent and alignment; input values and validity are preserved. +/// Block order is x-fastest thread rank, then each thread's items in index order. +/// Warp accepts a scalar only, with prefixes restarting within each logical group. +/// Every required input slot must be initialized; no valid-count option exists. +/// +/// Defaults select exclusive sum with typed-zero seed and Block Raking. Inclusive +/// mode includes the current item and forbids a seed. Non-sum exclusive mode +/// requires an exact-T, group-uniform seed. All ten [`Numeric`] widths are +/// supported, except bitwise operators require integers. CUB may regroup +/// associative arithmetic; overflow and floating behavior follow the input +/// dtype and underlying primitive, not sequential host arithmetic guarantees. +/// +/// Every member must reach matching operations. Oxide enforces the exact launch +/// geometry, canonical group participation and seed uniformity. Block supports +/// all three algorithms and implicit or explicit storage; Warp supports only +/// omitted algorithm/storage. [`WarpScans`] requires a block size divisible by 32. +/// Implicit scratch inserts completion before dirty reuse; it is not a trailing +/// completion guarantee. Explicit manual scratch requires completion before +/// conflicting reuse, and auto-sync storage inserts a trailing block barrier. +/// Scan alone is not completion of prior global-memory accesses. The operation +/// uses the existing checked native provider and requires the pinned Oxide toolchain. +/// +/// # Panics +/// An unwritten required input traps before its numeric read. Unsupported +/// type/options, group participation, geometry, or scratch reuse fail compilation. +/// +/// ```no_run +/// use cuda_coop as coop; +/// fn prefixes( +/// values: &coop::ThreadData, +/// ) -> coop::ThreadData { +/// coop::scan(coop::this_block(), values, Default::default()) +/// } +/// ``` +#[inline(always)] +pub fn scan< + T: Numeric, + V: ScanValue, + G: ScanGroup, + M: ScanMode, + Op: ScanOperator, + A: ScanAlgorithm, + Seed, + S: TempStorageLike, +>( + group: ThreadGroup, + value: V, + options: ScanOptions, +) -> V::Output { + value.scan(group, options) +} + +/// Exclusive sum with typed-zero seed; input/result and completion follow [`scan`]. +/// +/// Scalars return scalars; borrowed Block payloads return fresh payloads. Options +/// select only Block algorithm/storage. Default options also support scalar Warp. +#[inline(always)] +pub fn exclusive_sum< + T: Numeric, + V: ScanValue, + G: ScanGroup, + A: ScanAlgorithm, + S: TempStorageLike, +>( + group: ThreadGroup, + value: V, + options: SumScanOptions, +) -> V::Output { + scan(group, value, options) +} + +/// Inclusive sum, including the current item; ownership and completion follow [`scan`]. +/// +/// No seed or replacement operator is accepted. Default options support Block +/// scalar/payload inputs and scalar Warp inputs. +#[inline(always)] +pub fn inclusive_sum< + T: Numeric, + V: ScanValue, + G: ScanGroup, + A: ScanAlgorithm, + S: TempStorageLike, +>( + group: ThreadGroup, + value: V, + options: SumScanOptions, +) -> V::Output { + scan(group, value, options.mode(Inclusive)) +} + +/// Exclusive Scan with a static operator; result ownership and completion follow [`scan`]. +/// +/// Defaults select sum/typed zero. Every other operator requires an exact-T, +/// group-uniform [`initial_value`](ScanOptions::initial_value). The first result +/// is the seed and every later result combines it with preceding inputs. +#[inline(always)] +pub fn exclusive_scan< + T: Numeric, + V: ScanValue, + G: ScanGroup, + Op: ScanOperator, + A: ScanAlgorithm, + Seed, + S: TempStorageLike, +>( + group: ThreadGroup, + value: V, + options: ExclusiveScanOptions, +) -> V::Output +where + (): ScanMode, +{ + scan(group, value, options) +} + +/// Inclusive Scan with a static operator and no seed; ownership/completion follow [`scan`]. +/// +/// Defaults select sum. The result includes the current item, and every output +/// slot is defined. Block algorithm/storage choices and scalar Warp restrictions +/// are the same as for [`scan`]. +#[inline(always)] +pub fn inclusive_scan< + T: Numeric, + V: ScanValue, + G: ScanGroup, + Op: ScanOperator, + A: ScanAlgorithm, + S: TempStorageLike, +>( + group: ThreadGroup, + value: V, + options: InclusiveScanOptions, +) -> V::Output { + scan(group, value, options.mode(Inclusive)) +} + +/// A sealed Scan input with its inferred fresh output type. +/// +/// Numeric scalars return T. Shared/mutable payload borrows are supported only +/// for Block and return `ThreadData` without changing the input. +/// Warp payload inputs and owned-payload-consuming forms are not provided. +/// +/// ```compile_fail,E0277 +/// use cuda_coop as coop; +/// fn invalid(values: &coop::ThreadData) { +/// coop::scan(coop::this_warp(), values, Default::default()); +/// } +/// ``` +pub trait ScanValue: sealed::Value { + /// Scalar T or a fresh payload with the input's dtype, extent and alignment. + type Output; + #[doc(hidden)] + fn scan< + M: ScanMode, + Op: ScanOperator, + A: ScanAlgorithm, + Seed, + S: TempStorageLike, + >( + self, + group: ThreadGroup, + options: ScanOptions, + ) -> Self::Output + where + G: ScanGroup; +} +impl sealed::Value for T {} +impl ScanValue for T { + type Output = T; + #[inline(always)] + fn scan< + M: ScanMode, + Op: ScanOperator, + A: ScanAlgorithm, + Seed, + S: TempStorageLike, + >( + self, + _group: ThreadGroup, + options: ScanOptions, + ) -> T + where + G: ScanGroup, + { + let mut input = ThreadData::<_, 1>::new(); + input.write(0, self); + let output = boundary::scan::(&input, options); + let (pointer, _) = output.block_io_read_parts(); + // SAFETY: the synchronous selected recipe initializes every output before returning. + unsafe { pointer.read() } + } +} +impl sealed::Value + for &ThreadData +where + Alignment: SupportedAlignment, +{ +} +impl ScanValue + for &ThreadData +where + Alignment: SupportedAlignment, +{ + type Output = ThreadData; + #[inline(always)] + fn scan< + M: ScanMode, + Op: ScanOperator, + A: ScanAlgorithm, + Seed, + S: TempStorageLike, + >( + self, + _group: ThreadGroup, + options: ScanOptions, + ) -> Self::Output { + boundary::scan::(self, options) + } +} +impl sealed::Value + for &mut ThreadData +where + Alignment: SupportedAlignment, +{ +} +impl ScanValue + for &mut ThreadData +where + Alignment: SupportedAlignment, +{ + type Output = ThreadData; + #[inline(always)] + fn scan< + M: ScanMode, + Op: ScanOperator, + A: ScanAlgorithm, + Seed, + S: TempStorageLike, + >( + self, + group: ThreadGroup, + options: ScanOptions, + ) -> Self::Output { + (&*self).scan(group, options) + } +} + +mod sealed { + pub trait Operator { + const KIND: u8; + } + pub trait Algorithm { + const KIND: u8; + } + pub trait Mode { + type Uniform: Copy; + const KIND: u8; + fn controls(seed: Seed) -> Self::Uniform; + } + pub trait Value {} + pub trait Group { + const WARP: bool; + unsafe fn call< + T: Copy, + R: Fn() -> &'static str + Copy, + S, + U: Copy, + P: Copy, + const I: usize, + >( + registration: R, + storage: S, + uniform: U, + payload: P, + ); + } +} diff --git a/rust/cuda-coop/src/scan/boundary.rs b/rust/cuda-coop/src/scan/boundary.rs new file mode 100644 index 000000000000..ae4434e9e220 --- /dev/null +++ b/rust/cuda-coop/src/scan/boundary.rs @@ -0,0 +1,334 @@ +// Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +use super::{ScanAlgorithm, ScanGroup, ScanMode, ScanOperator, ScanOptions}; +use crate::{Alignment, Numeric, SupportedAlignment, TempStorageLike, ThreadData}; +use cuda_device::library; + +// Each static recipe keeps a concrete function-item registration identity. +macro_rules! registrations { + ($(($name:ident, $variable:literal)),+ $(,)?) => {$ ( + fn $name() -> &'static str { library::registration!($variable) } + )+}; +} +registrations!( + (raking_exclusive_zero, "COOP_SCAN_RAKING_EXCLUSIVE_ZERO"), + ( + raking_memoize_exclusive_zero, + "COOP_SCAN_RAKING_MEMOIZE_EXCLUSIVE_ZERO" + ), + ( + warp_scans_exclusive_zero, + "COOP_SCAN_WARP_SCANS_EXCLUSIVE_ZERO" + ), + (warp_exclusive_zero, "COOP_SCAN_WARP_EXCLUSIVE_ZERO"), + (raking_exclusive_sum, "COOP_SCAN_RAKING_EXCLUSIVE_SUM"), + ( + raking_memoize_exclusive_sum, + "COOP_SCAN_RAKING_MEMOIZE_EXCLUSIVE_SUM" + ), + ( + warp_scans_exclusive_sum, + "COOP_SCAN_WARP_SCANS_EXCLUSIVE_SUM" + ), + (warp_exclusive_sum, "COOP_SCAN_WARP_EXCLUSIVE_SUM"), + ( + raking_exclusive_multiplies, + "COOP_SCAN_RAKING_EXCLUSIVE_MULTIPLIES" + ), + ( + raking_memoize_exclusive_multiplies, + "COOP_SCAN_RAKING_MEMOIZE_EXCLUSIVE_MULTIPLIES" + ), + ( + warp_scans_exclusive_multiplies, + "COOP_SCAN_WARP_SCANS_EXCLUSIVE_MULTIPLIES" + ), + ( + warp_exclusive_multiplies, + "COOP_SCAN_WARP_EXCLUSIVE_MULTIPLIES" + ), + (raking_exclusive_min, "COOP_SCAN_RAKING_EXCLUSIVE_MIN"), + ( + raking_memoize_exclusive_min, + "COOP_SCAN_RAKING_MEMOIZE_EXCLUSIVE_MIN" + ), + ( + warp_scans_exclusive_min, + "COOP_SCAN_WARP_SCANS_EXCLUSIVE_MIN" + ), + (warp_exclusive_min, "COOP_SCAN_WARP_EXCLUSIVE_MIN"), + (raking_exclusive_max, "COOP_SCAN_RAKING_EXCLUSIVE_MAX"), + ( + raking_memoize_exclusive_max, + "COOP_SCAN_RAKING_MEMOIZE_EXCLUSIVE_MAX" + ), + ( + warp_scans_exclusive_max, + "COOP_SCAN_WARP_SCANS_EXCLUSIVE_MAX" + ), + (warp_exclusive_max, "COOP_SCAN_WARP_EXCLUSIVE_MAX"), + ( + raking_exclusive_bit_and, + "COOP_SCAN_RAKING_EXCLUSIVE_BIT_AND" + ), + ( + raking_memoize_exclusive_bit_and, + "COOP_SCAN_RAKING_MEMOIZE_EXCLUSIVE_BIT_AND" + ), + ( + warp_scans_exclusive_bit_and, + "COOP_SCAN_WARP_SCANS_EXCLUSIVE_BIT_AND" + ), + (warp_exclusive_bit_and, "COOP_SCAN_WARP_EXCLUSIVE_BIT_AND"), + (raking_exclusive_bit_or, "COOP_SCAN_RAKING_EXCLUSIVE_BIT_OR"), + ( + raking_memoize_exclusive_bit_or, + "COOP_SCAN_RAKING_MEMOIZE_EXCLUSIVE_BIT_OR" + ), + ( + warp_scans_exclusive_bit_or, + "COOP_SCAN_WARP_SCANS_EXCLUSIVE_BIT_OR" + ), + (warp_exclusive_bit_or, "COOP_SCAN_WARP_EXCLUSIVE_BIT_OR"), + ( + raking_exclusive_bit_xor, + "COOP_SCAN_RAKING_EXCLUSIVE_BIT_XOR" + ), + ( + raking_memoize_exclusive_bit_xor, + "COOP_SCAN_RAKING_MEMOIZE_EXCLUSIVE_BIT_XOR" + ), + ( + warp_scans_exclusive_bit_xor, + "COOP_SCAN_WARP_SCANS_EXCLUSIVE_BIT_XOR" + ), + (warp_exclusive_bit_xor, "COOP_SCAN_WARP_EXCLUSIVE_BIT_XOR"), + (raking_inclusive_sum, "COOP_SCAN_RAKING_INCLUSIVE_SUM"), + ( + raking_memoize_inclusive_sum, + "COOP_SCAN_RAKING_MEMOIZE_INCLUSIVE_SUM" + ), + ( + warp_scans_inclusive_sum, + "COOP_SCAN_WARP_SCANS_INCLUSIVE_SUM" + ), + (warp_inclusive_sum, "COOP_SCAN_WARP_INCLUSIVE_SUM"), + ( + raking_inclusive_multiplies, + "COOP_SCAN_RAKING_INCLUSIVE_MULTIPLIES" + ), + ( + raking_memoize_inclusive_multiplies, + "COOP_SCAN_RAKING_MEMOIZE_INCLUSIVE_MULTIPLIES" + ), + ( + warp_scans_inclusive_multiplies, + "COOP_SCAN_WARP_SCANS_INCLUSIVE_MULTIPLIES" + ), + ( + warp_inclusive_multiplies, + "COOP_SCAN_WARP_INCLUSIVE_MULTIPLIES" + ), + (raking_inclusive_min, "COOP_SCAN_RAKING_INCLUSIVE_MIN"), + ( + raking_memoize_inclusive_min, + "COOP_SCAN_RAKING_MEMOIZE_INCLUSIVE_MIN" + ), + ( + warp_scans_inclusive_min, + "COOP_SCAN_WARP_SCANS_INCLUSIVE_MIN" + ), + (warp_inclusive_min, "COOP_SCAN_WARP_INCLUSIVE_MIN"), + (raking_inclusive_max, "COOP_SCAN_RAKING_INCLUSIVE_MAX"), + ( + raking_memoize_inclusive_max, + "COOP_SCAN_RAKING_MEMOIZE_INCLUSIVE_MAX" + ), + ( + warp_scans_inclusive_max, + "COOP_SCAN_WARP_SCANS_INCLUSIVE_MAX" + ), + (warp_inclusive_max, "COOP_SCAN_WARP_INCLUSIVE_MAX"), + ( + raking_inclusive_bit_and, + "COOP_SCAN_RAKING_INCLUSIVE_BIT_AND" + ), + ( + raking_memoize_inclusive_bit_and, + "COOP_SCAN_RAKING_MEMOIZE_INCLUSIVE_BIT_AND" + ), + ( + warp_scans_inclusive_bit_and, + "COOP_SCAN_WARP_SCANS_INCLUSIVE_BIT_AND" + ), + (warp_inclusive_bit_and, "COOP_SCAN_WARP_INCLUSIVE_BIT_AND"), + (raking_inclusive_bit_or, "COOP_SCAN_RAKING_INCLUSIVE_BIT_OR"), + ( + raking_memoize_inclusive_bit_or, + "COOP_SCAN_RAKING_MEMOIZE_INCLUSIVE_BIT_OR" + ), + ( + warp_scans_inclusive_bit_or, + "COOP_SCAN_WARP_SCANS_INCLUSIVE_BIT_OR" + ), + (warp_inclusive_bit_or, "COOP_SCAN_WARP_INCLUSIVE_BIT_OR"), + ( + raking_inclusive_bit_xor, + "COOP_SCAN_RAKING_INCLUSIVE_BIT_XOR" + ), + ( + raking_memoize_inclusive_bit_xor, + "COOP_SCAN_RAKING_MEMOIZE_INCLUSIVE_BIT_XOR" + ), + ( + warp_scans_inclusive_bit_xor, + "COOP_SCAN_WARP_SCANS_INCLUSIVE_BIT_XOR" + ), + (warp_inclusive_bit_xor, "COOP_SCAN_WARP_INCLUSIVE_BIT_XOR"), +); + +#[inline(always)] +pub(super) fn scan< + T: Numeric, + const ITEMS: usize, + const ALIGN: usize, + G: ScanGroup, + M: ScanMode, + Op: ScanOperator, + A: ScanAlgorithm, + Seed, + S: TempStorageLike, +>( + input: &ThreadData, + options: ScanOptions, +) -> ThreadData +where + Alignment: SupportedAlignment, +{ + let mut output = ThreadData::new(); + let (input, input_bits) = input.block_io_read_parts(); + let (output_pointer, output_bits) = output.block_io_parts(); + let payload = (input, output_pointer, input_bits, output_bits); + let uniform = M::controls(options.initial_value); + let storage = options.storage.into_raw(); + let _operator = options.scan_op; + macro_rules! call { + ($raking:ident, $memoize:ident, $warps:ident, $warp:ident) => { + if G::WARP { + G::call::($warp, storage, uniform, payload); + } else { + match A::KIND { + 0 | 1 => G::call::($raking, storage, uniform, payload), + 2 => G::call::($memoize, storage, uniform, payload), + 3 => G::call::($warps, storage, uniform, payload), + _ => panic!(), + } + } + }; + } + // SAFETY: input and fresh output are distinct live aligned lane-local + // payloads. The exact native recipe validates every input bit before + // numeric reads, preserves input, and initializes all outputs before + // setting their bits. The pointers cannot escape the synchronous call. + // Sealed static selections establish ABI/mode/operator compatibility; + // the compiler enforces participation, seed uniformity, exact layout and + // scratch policy through the payload-only checked marker. No fabricated + // global allocation witness or memory completion is introduced. + unsafe { + match M::KIND { + 0 => call!( + raking_exclusive_zero, + raking_memoize_exclusive_zero, + warp_scans_exclusive_zero, + warp_exclusive_zero + ), + 1 => call!( + raking_exclusive_sum, + raking_memoize_exclusive_sum, + warp_scans_exclusive_sum, + warp_exclusive_sum + ), + 2 => call!( + raking_exclusive_multiplies, + raking_memoize_exclusive_multiplies, + warp_scans_exclusive_multiplies, + warp_exclusive_multiplies + ), + 3 => call!( + raking_exclusive_min, + raking_memoize_exclusive_min, + warp_scans_exclusive_min, + warp_exclusive_min + ), + 4 => call!( + raking_exclusive_max, + raking_memoize_exclusive_max, + warp_scans_exclusive_max, + warp_exclusive_max + ), + 5 => call!( + raking_exclusive_bit_and, + raking_memoize_exclusive_bit_and, + warp_scans_exclusive_bit_and, + warp_exclusive_bit_and + ), + 6 => call!( + raking_exclusive_bit_or, + raking_memoize_exclusive_bit_or, + warp_scans_exclusive_bit_or, + warp_exclusive_bit_or + ), + 7 => call!( + raking_exclusive_bit_xor, + raking_memoize_exclusive_bit_xor, + warp_scans_exclusive_bit_xor, + warp_exclusive_bit_xor + ), + 8 => call!( + raking_inclusive_sum, + raking_memoize_inclusive_sum, + warp_scans_inclusive_sum, + warp_inclusive_sum + ), + 9 => call!( + raking_inclusive_multiplies, + raking_memoize_inclusive_multiplies, + warp_scans_inclusive_multiplies, + warp_inclusive_multiplies + ), + 10 => call!( + raking_inclusive_min, + raking_memoize_inclusive_min, + warp_scans_inclusive_min, + warp_inclusive_min + ), + 11 => call!( + raking_inclusive_max, + raking_memoize_inclusive_max, + warp_scans_inclusive_max, + warp_inclusive_max + ), + 12 => call!( + raking_inclusive_bit_and, + raking_memoize_inclusive_bit_and, + warp_scans_inclusive_bit_and, + warp_inclusive_bit_and + ), + 13 => call!( + raking_inclusive_bit_or, + raking_memoize_inclusive_bit_or, + warp_scans_inclusive_bit_or, + warp_inclusive_bit_or + ), + 14 => call!( + raking_inclusive_bit_xor, + raking_memoize_inclusive_bit_xor, + warp_scans_inclusive_bit_xor, + warp_inclusive_bit_xor + ), + _ => panic!(), + } + } + output +} diff --git a/rust/cuda-coop/tests/check_load_store.py b/rust/cuda-coop/tests/check_load_store.py index 3e71d7f4bf4a..b205e730a4e6 100644 --- a/rust/cuda-coop/tests/check_load_store.py +++ b/rust/cuda-coop/tests/check_load_store.py @@ -1,7 +1,7 @@ # 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. +"""Compile public Load/Store and Scan participation/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 @@ -25,6 +25,9 @@ ) parser.add_argument("--cuda-root", type=Path, default=Path("/usr/local/cuda")) parser.add_argument("--arch", default="sm_120") +parser.add_argument( + "--case", action="append", help="Run only this named case; repeat to select several" +) args = parser.parse_args() output = args.output.resolve() output.mkdir(parents=True, exist_ok=False) @@ -831,6 +834,111 @@ ) +# Scan returns fresh lane-private output while retaining the existing group and +# storage proofs. These cases exercise the public API, not private markers. +scan = "values=coop::scan(block,&values,Default::default());" +scan_manual = "values=coop::scan(block,&values,coop::ScanOptions::default().temp_storage(&mut scratch));" +scan_warp = "let warp=coop::this_warp().group_by::<8>();" +cases.update( + scan_defaults=( + load + + scan + + "values=coop::exclusive_sum(block,&values,Default::default());" + + "values=coop::inclusive_sum(block,&values,Default::default());" + + "values=coop::exclusive_scan(block,&values,Default::default());" + + "values=coop::inclusive_scan(block,&values,Default::default());" + + store, + None, + ), + scan_generic_seed=( + load + "fn seeded,const N:usize>(" + "group:coop::ThreadGroup,input:&coop::ThreadData,seed:T,op:Op)" + "->coop::ThreadData where ():coop::ScanMode {" + "coop::exclusive_scan(group,input,coop::ScanOptions::default().scan_op(op).initial_value(seed))}" + "values=seeded(block,&values,valid as i32,coop::Min);" + store, + None, + ), + scan_varying_seed=( + load + + "values=coop::exclusive_scan(block,&values,coop::ScanOptions::default().initial_value(thread::threadIdx_x() as i32));", + "not proved block-uniform", + ), + scan_varying_participation=( + load + "if thread::threadIdx_x()&1==0 {" + scan + "}", + "lane-varying", + ), + scan_warp_scalar=( + load + scan_warp + "let seed=warp.rank_at::() as i32;" + "let result=coop::exclusive_scan(warp,thread::threadIdx_x() as i32,coop::ScanOptions::default().initial_value(seed));" + "values.write(0,result);", + None, + ), + scan_warp_payload=( + load + scan_warp + "let _result=coop::scan(warp,&values,Default::default());", + "ScanValue", + ), + scan_warp_algorithm=( + load + + scan_warp + + "let _result=coop::scan(warp,1_i32,coop::ScanOptions::default().algorithm(coop::Raking));", + "expected `()`, found `Raking`", + ), + scan_warp_storage=( + load + + scan_warp + + new_storage + + "let _result=coop::scan(warp,1_i32,coop::ScanOptions::default().temp_storage(&mut scratch));", + "expected struct `ScanOptions`", + ), + scan_missing_seed=( + load + + "let _result=coop::scan(block,&values,coop::ScanOptions::default().scan_op(coop::Min));", + "ScanMode", + ), + scan_inclusive_seed=( + load + + "let _result=coop::scan(block,&values,coop::ScanOptions::default().mode(coop::Inclusive).initial_value(0));", + "ScanMode", + ), + scan_seed_dtype=( + load + + "let _result=coop::exclusive_scan(block,&values,coop::ScanOptions::default().initial_value(0_i16));", + "mismatched types", + ), + scan_float_bitwise=( + load + + "let _result=coop::inclusive_scan(block,1_f32,coop::ScanOptions::default().scan_op(coop::BitAnd));", + "ScanOperator", + ), + scan_storage_reuse=( + load + new_storage + scan_manual + sync + scan_manual + store, + None, + ), + scan_storage_missing_barrier=( + load + new_storage + scan_manual + scan_manual, + "requires a full-block barrier", + ), + scan_no_global_completion=( + load + store + scan + reload, + "requires a barrier covering both accesses", + ), + scan_storage_auto_reuse=( + load + + new_storage.replace("new()", "new().auto_sync()") + + scan_manual + + scan_manual + + store, + None, + ), +) + +if args.case: + unknown = set(args.case).difference(cases) + if unknown: + parser.error(f"unknown cases: {', '.join(sorted(unknown))}") + cases = {name: case for name, case in cases.items() if name in args.case} + + command = [ str(args.cargo_oxide.resolve()), "oxide", diff --git a/rust/cuda-coop/tests/scan/.gitignore b/rust/cuda-coop/tests/scan/.gitignore new file mode 100644 index 000000000000..74bffa8af4a4 --- /dev/null +++ b/rust/cuda-coop/tests/scan/.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/scan/Cargo.lock b/rust/cuda-coop/tests/scan/Cargo.lock new file mode 100644 index 000000000000..0809eccb39af --- /dev/null +++ b/rust/cuda-coop/tests/scan/Cargo.lock @@ -0,0 +1,1011 @@ +# 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-scan-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=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "libnvvm-sys", + "nvjitlink-sys", + "serde", + "sha2", + "thiserror 2.0.21", +] + +[[package]] +name = "cuda-async" +version = "0.5.0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +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.5.0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +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.5.0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "anyhow", + "cuda-bindings", + "cuda-core-derive", + "half", + "oxide-artifacts", +] + +[[package]] +name = "cuda-core-derive" +version = "0.5.0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +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=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "cuda-macros", +] + +[[package]] +name = "cuda-host" +version = "0.2.1" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "cuda-artifact-finalizer", + "cuda-async", + "cuda-core", + "cuda-device", + "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=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +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=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" + +[[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=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +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=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +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=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "object", +] + +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +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=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" + +[[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=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" + +[[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/scan/Cargo.toml b/rust/cuda-coop/tests/scan/Cargo.toml new file mode 100644 index 000000000000..05aab44b5af6 --- /dev/null +++ b/rust/cuda-coop/tests/scan/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-scan-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 = "96caed46c776eb49226137eac849e795b83b7ce0" } +cuda-host = { git = "https://github.com/tpn/cuda-rust.git", rev = "96caed46c776eb49226137eac849e795b83b7ce0", features = ["async"] } +cuda-core = { git = "https://github.com/tpn/cuda-rust.git", rev = "96caed46c776eb49226137eac849e795b83b7ce0" } +cuda-async = { git = "https://github.com/tpn/cuda-rust.git", rev = "96caed46c776eb49226137eac849e795b83b7ce0" } diff --git a/rust/cuda-coop/tests/scan/src/main.rs b/rust/cuda-coop/tests/scan/src/main.rs new file mode 100644 index 000000000000..ceba29760c27 --- /dev/null +++ b/rust/cuda-coop/tests/scan/src/main.rs @@ -0,0 +1,540 @@ +// Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +//! Public Scan values, validity, preservation, groups, and scratch composition. +use cuda_coop as coop; +use cuda_device::{kernel, launch_contract}; +use cuda_host::cuda_module; + +#[cuda_module] +mod kernels { + use super::*; + + #[inline(never)] + fn observe( + input: &coop::ThreadData, + before: &coop::ThreadData, + result: &coop::ThreadData, + ) -> bool { + let mut matches = true; + let mut i = 0; + while i < ITEMS_PER_THREAD { + matches &= input.get(i) == before.get(i) && result.get(i).is_some(); + i += 1; + } + matches && result.get(ITEMS_PER_THREAD).is_none() + } + + #[kernel] + #[launch_contract(domain = 3, block = (4, 4, 4), dynamic_shared = inferred)] + pub fn scan_sum( + input: &[T], + mut output: coop::DisjointSlice>, + valid: usize, + zero: T, + invalid: bool, + ) { + let block = coop::this_block(); + let offset = block.rank_at::() as usize * 64 * ITEMS_PER_THREAD; + let mut values = coop::ThreadData::<_, ITEMS_PER_THREAD>::new(); + let options = coop::LoadOptions::default() + .offset(offset) + .valid_items(valid); + if invalid { + coop::load(block, input, &mut values, options); + } else { + coop::load(block, input, &mut values, options.oob_default(zero)); + } + let before = values; + let mut result = coop::scan(block, &values, Default::default()); + let matches = observe(&values, &before, &result); + // Store every result, including prefixes whose input was default-filled. + coop::store( + block, + &mut output, + &result, + coop::StoreOptions::default().offset(offset), + ); + result.write(0, zero); + assert!(matches && observe(&values, &before, &result)); + } + + #[kernel] + #[launch_contract(domain = 1, block = (64, 1, 1), dynamic_shared = inferred)] + pub fn scan_seeded( + input: &[T], + mut output: coop::DisjointSlice>, + valid: usize, + zero: T, + seed: T, + ) { + let block = coop::this_block(); + let mut values = coop::ThreadData::<_, ITEMS_PER_THREAD>::new(); + coop::load( + block, + input, + &mut values, + coop::LoadOptions::default() + .valid_items(valid) + .oob_default(zero), + ); + let before = values; + let result = coop::exclusive_scan( + block, + &mut values, + coop::ScanOptions::default().initial_value(seed), + ); + let matches = observe(&values, &before, &result); + coop::store(block, &mut output, &result, Default::default()); + assert!(matches && observe(&values, &before, &result)); + } + + #[kernel] + #[launch_contract(domain = 3, block = (4, 4, 4), dynamic_shared = inferred)] + pub fn scan_operators( + input: &[i32], + mut output: coop::DisjointSlice>, + mode: u32, + ) { + let block = coop::this_block(); + let mut values = coop::ThreadData::<_, 4>::new(); + coop::load(block, input, &mut values, Default::default()); + let before = values; + let result = match mode { + 0 => coop::exclusive_sum(block, &values, Default::default()), + 1 => coop::inclusive_sum(block, &values, Default::default()), + 2 => coop::exclusive_scan( + block, + &values, + coop::ScanOptions::default() + .scan_op(coop::Multiplies) + .initial_value(1), + ), + 3 => coop::inclusive_scan( + block, + &values, + coop::ScanOptions::default().scan_op(coop::Min), + ), + 4 => coop::scan( + block, + &values, + coop::ScanOptions::default() + .mode(coop::Inclusive) + .scan_op(coop::Max), + ), + 5 => coop::exclusive_scan( + block, + &values, + coop::ScanOptions::default() + .scan_op(coop::BitAnd) + .initial_value(255), + ), + 6 => coop::inclusive_scan( + block, + &values, + coop::ScanOptions::default().scan_op(coop::BitOr), + ), + 7 => coop::inclusive_scan( + block, + &values, + coop::ScanOptions::default().scan_op(coop::BitXor), + ), + 8 => coop::exclusive_scan( + block, + &values, + coop::ScanOptions::default() + .initial_value(3) + .algorithm(coop::RakingMemoize), + ), + _ => coop::exclusive_scan( + block, + &values, + coop::ScanOptions::default() + .initial_value(3) + .algorithm(coop::WarpScans), + ), + }; + let matches = observe(&values, &before, &result); + coop::store(block, &mut output, &result, Default::default()); + assert!(matches && observe(&values, &before, &result)); + } + + #[inline(never)] + fn with_storage( + input: &[i32], + output: &mut coop::DisjointSlice>, + valid: usize, + scratch: &mut coop::TempStorage<0, 0, AUTO, false>, + ) { + let block = coop::this_block(); + let mut values = coop::ThreadData::<_, 4>::new(); + coop::load( + block, + input, + &mut values, + coop::LoadOptions::default() + .valid_items(valid) + .oob_default(0) + .algorithm(coop::Transpose) + .temp_storage(scratch), + ); + if !AUTO { + cuda_device::thread::sync_threads(); + } + let before = values; + let mut matches = true; + let mut iteration = 0; + while iteration < 2 { + let result = coop::exclusive_scan( + block, + &values, + coop::ScanOptions::default() + .initial_value(5) + .temp_storage(scratch), + ); + matches &= observe(&values, &before, &result); + if !AUTO { + cuda_device::thread::sync_threads(); + } + coop::store(block, &mut *output, &result, Default::default()); + // Complete the global Store before repeating the same output region. + cuda_device::thread::sync_threads(); + iteration += 1; + } + assert!(matches); + } + + #[kernel] + #[launch_contract(domain = 1, block = (64, 1, 1), dynamic_shared = inferred)] + pub fn scan_manual( + input: &[i32], + mut output: coop::DisjointSlice>, + valid: usize, + ) { + let mut scratch = coop::TempStorage::new(); + with_storage(input, &mut output, valid, &mut scratch); + } + + #[kernel] + #[launch_contract(domain = 1, block = (64, 1, 1), dynamic_shared = inferred)] + pub fn scan_auto( + input: &[i32], + mut output: coop::DisjointSlice>, + valid: usize, + ) { + let mut scratch = coop::TempStorage::new().auto_sync(); + with_storage(input, &mut output, valid, &mut scratch); + } + + #[kernel] + #[launch_contract(domain = 1, block = (64, 1, 1), dynamic_shared = inferred)] + pub fn scan_warp( + input: &[i32], + mut output: coop::DisjointSlice>, + seeded: bool, + ) { + let group = coop::this_warp().group_by::<8>(); + let block = coop::this_block(); + let offset = block.rank_at::() as usize * 64; + let mut values = coop::ThreadData::<_, 1>::new(); + coop::load( + group, + input, + &mut values, + coop::LoadOptions::default().offset(offset), + ); + let value = values.get(0).unwrap_or(0); + let q = block.rank_at::() as i32 * 4 + group.rank_at::() as i32; + let result = if seeded { + coop::exclusive_scan( + group, + value, + coop::ScanOptions::default().initial_value(q + 5), + ) + } else { + coop::exclusive_sum(group, value, Default::default()) + }; + let matches = values.get(0) == Some(value); + coop::store( + group, + &mut output, + result, + coop::StoreOptions::default().offset(offset), + ); + assert!(matches && values.get(0) == Some(value)); + } + + #[kernel] + #[launch_contract(domain = 1, block = (64, 1, 1), dynamic_shared = inferred)] + pub fn scan_physical_warp( + input: &[i32], + mut output: coop::DisjointSlice>, + ) { + let group = coop::this_warp(); + let offset = coop::this_block().rank_at::() as usize * 64; + let mut values = coop::ThreadData::<_, 1>::new(); + coop::load( + group, + input, + &mut values, + coop::LoadOptions::default().offset(offset), + ); + let value = values.get(0).unwrap_or(0); + let result = coop::exclusive_sum(group, value, Default::default()); + coop::store( + group, + &mut output, + result, + coop::StoreOptions::default().offset(offset), + ); + assert!(values.get(0) == Some(value)); + } + + #[kernel] + #[launch_contract(domain = 1, block = (64, 1, 1), dynamic_shared = inferred)] + pub fn scan_block_scalar( + input: &[i32], + mut output: coop::DisjointSlice>, + ) { + let group = coop::this_block(); + let offset = group.rank_at::() as usize * 64; + let mut values = coop::ThreadData::<_, 1>::new(); + coop::load( + group, + input, + &mut values, + coop::LoadOptions::default().offset(offset), + ); + let value = values.get(0).unwrap_or(0); + let result = coop::scan(group, value, Default::default()); + coop::store( + group, + &mut output, + result, + coop::StoreOptions::default().offset(offset), + ); + assert!(values.get(0) == Some(value)); + } +} + +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, LaunchConfig3D}; + + init_device_contexts(0, 1)?; + let context = with_cuda_context(0, |context| context.clone())?; + let stream = context.default_stream(); + // SAFETY: the typed declarations and embedded kernels compile together. + let module = unsafe { kernels::load_async(0)? }; + if std::env::args().any(|arg| arg == "invalid-input") { + let source = DeviceBuffer::from_host(&stream, &[1_i32; 256])?; + let target = DeviceBuffer::from_host(&stream, &[0_i32; 256])?; + let plan = + module.prepare_scan_sum::(LaunchConfig3D::auto_shared((1, 1, 1), (4, 4, 4)))?; + module + .scan_sum_async_owned::(&plan, source, target, 1, 0, true) + .sync()?; + stream.synchronize()?; + panic!("unwritten Scan input unexpectedly completed"); + } + let mut cases = 0; + macro_rules! sum_cases { + ($ty:ty, $items:literal) => {{ + let tile = 64 * $items; + let input: Vec<$ty> = (0..2 * tile).map(|i| (i % 8 == 0) as u8 as $ty).collect(); + let plan = module.prepare_scan_sum::<$ty, $items>(LaunchConfig3D::auto_shared( + (2, 1, 1), + (4, 4, 4), + ))?; + for valid in [0, 1, tile - 1, tile] { + let source = DeviceBuffer::from_host(&stream, &input)?; + let target = DeviceBuffer::from_host(&stream, &vec![0 as $ty; input.len()])?; + let (_, target) = module + .scan_sum_async_owned::<$ty, $items, _, _>( + &plan, source, target, valid, 0 as $ty, false, + ) + .sync()?; + let mut actual = vec![0 as $ty; input.len()]; + target.copy_to_host(&stream, &mut actual)?; + stream.synchronize()?; + let mut expected = Vec::with_capacity(input.len()); + for chunk in input.chunks(tile) { + let mut prefix = 0 as $ty; + for (i, &value) in chunk.iter().enumerate() { + expected.push(prefix); + if i < valid { + prefix += value; + } + } + } + assert_eq!( + actual, + expected, + "{} N{} valid{}", + stringify!($ty), + $items, + valid + ); + cases += 1; + } + }}; + } + sum_cases!(i8, 1); + sum_cases!(i8, 4); + sum_cases!(u8, 1); + sum_cases!(u8, 4); + sum_cases!(i16, 1); + sum_cases!(i16, 4); + sum_cases!(u16, 1); + sum_cases!(u16, 4); + sum_cases!(i32, 1); + sum_cases!(i32, 4); + sum_cases!(u32, 1); + sum_cases!(u32, 4); + sum_cases!(i64, 1); + sum_cases!(i64, 4); + sum_cases!(u64, 1); + sum_cases!(u64, 4); + sum_cases!(f32, 1); + sum_cases!(f32, 4); + sum_cases!(f64, 1); + sum_cases!(f64, 4); + sum_cases!(i32, 9); + macro_rules! seeded_cases { + ($ty:ty, $items:literal, $seed:expr) => {{ + let tile = 64 * $items; + let input: Vec<$ty> = (0..tile).map(|i| (i % 8 == 0) as u8 as $ty).collect(); + let plan = + module.prepare_scan_seeded::<$ty, $items>(LaunchConfig1D::auto_shared(1, 64))?; + for valid in [tile - 1, tile] { + let source = DeviceBuffer::from_host(&stream, &input)?; + let target = DeviceBuffer::from_host(&stream, &vec![0 as $ty; tile])?; + let (_, target) = module + .scan_seeded_async_owned::<$ty, $items, _, _>( + &plan, source, target, valid, 0 as $ty, $seed, + ) + .sync()?; + let mut actual = vec![0 as $ty; tile]; + target.copy_to_host(&stream, &mut actual)?; + stream.synchronize()?; + let mut prefix: $ty = $seed; + for (i, (&value, &got)) in input.iter().zip(&actual).enumerate() { + assert_eq!(got, prefix); + if i < valid { + prefix += value; + } + } + cases += 1; + } + }}; + } + seeded_cases!(i8, 4, -3_i8); + seeded_cases!(u16, 1, 7_u16); + let plan = module.prepare_scan_operators(LaunchConfig3D::auto_shared((1, 1, 1), (4, 4, 4)))?; + for mode in 0..10 { + let input: Vec = (0..256) + .map(|i| { + if mode == 2 { + if i == 0 { 2 } else { 1 } + } else { + i % 7 + 1 + } + }) + .collect(); + let source = DeviceBuffer::from_host(&stream, &input)?; + let target = DeviceBuffer::from_host(&stream, &[0_i32; 256])?; + let (_, target) = module + .scan_operators_async_owned(&plan, source, target, mode) + .sync()?; + let mut actual = vec![0_i32; 256]; + target.copy_to_host(&stream, &mut actual)?; + stream.synchronize()?; + let inclusive = matches!(mode, 1 | 3 | 4 | 6 | 7); + let mut prefix = match mode { + 2 => 1, + 3 => i32::MAX, + 4 => i32::MIN, + 5 => 255, + 8 | 9 => 3, + _ => 0, + }; + for (&value, &got) in input.iter().zip(&actual) { + let next = match mode { + 2 => prefix * value, + 3 => prefix.min(value), + 4 => prefix.max(value), + 5 => prefix & value, + 6 => prefix | value, + 7 => prefix ^ value, + _ => prefix + value, + }; + assert_eq!(got, if inclusive { next } else { prefix }, "mode {mode}"); + prefix = next; + } + cases += 1; + } + macro_rules! storage_cases { + ($prepare:ident, $launch:ident) => {{ + let input: Vec = (0..256).map(|i| (i % 8 == 0) as i32).collect(); + let plan = module.$prepare(LaunchConfig1D::auto_shared(1, 64))?; + for valid in [255, 256] { + let source = DeviceBuffer::from_host(&stream, &input)?; + let target = DeviceBuffer::from_host(&stream, &[0_i32; 256])?; + let (_, target) = module.$launch(&plan, source, target, valid).sync()?; + let mut actual = vec![0_i32; 256]; + target.copy_to_host(&stream, &mut actual)?; + stream.synchronize()?; + let mut prefix = 5; + for (i, (&value, &got)) in input.iter().zip(&actual).enumerate() { + assert_eq!(got, prefix); + if i < valid { + prefix += value; + } + } + cases += 1; + } + }}; + } + storage_cases!(prepare_scan_manual, scan_manual_async_owned); + storage_cases!(prepare_scan_auto, scan_auto_async_owned); + let plan = module.prepare_scan_warp(LaunchConfig1D::auto_shared(3, 64))?; + for seeded in [false, true] { + let source = DeviceBuffer::from_host(&stream, &[1_i32; 192])?; + let target = DeviceBuffer::from_host(&stream, &[0_i32; 192])?; + let (_, target) = module + .scan_warp_async_owned(&plan, source, target, seeded) + .sync()?; + let mut actual = vec![0_i32; 192]; + target.copy_to_host(&stream, &mut actual)?; + stream.synchronize()?; + for (i, &got) in actual.iter().enumerate() { + let seed = if seeded { (i % 64 / 8) as i32 + 5 } else { 0 }; + assert_eq!(got, seed + (i % 8) as i32); + } + cases += 1; + } + macro_rules! scalar_case { + ($prepare:ident, $launch:ident, $width:literal) => {{ + let plan = module.$prepare(LaunchConfig1D::auto_shared(3, 64))?; + let source = DeviceBuffer::from_host(&stream, &[1_i32; 192])?; + let target = DeviceBuffer::from_host(&stream, &[0_i32; 192])?; + let (_, target) = module.$launch(&plan, source, target).sync()?; + let mut actual = vec![0_i32; 192]; + target.copy_to_host(&stream, &mut actual)?; + stream.synchronize()?; + for (i, &got) in actual.iter().enumerate() { + assert_eq!(got, (i % $width) as i32); + } + cases += 1; + }}; + } + scalar_case!( + prepare_scan_physical_warp, + scan_physical_warp_async_owned, + 32 + ); + scalar_case!(prepare_scan_block_scalar, scan_block_scalar_async_owned, 64); + println!("PASS: {cases} public Scan cases"); + Ok(()) +} From 5d4f9c3d28df6cb854a67dac2ee2ca2a2138b9dd Mon Sep 17 00:00:00 2001 From: Trent Nelson Date: Sat, 24 Oct 2026 17:40:52 -0700 Subject: [PATCH 3/3] Add Rust Reduce and Sum with supported binary callables Provide scalar and payload reductions with validity-aware results. Include checked stateless and uniform i32 capture callbacks at their first public consumer, together with the native callback ABI. Signed-off-by: Trent Nelson --- docs/rust/index.rst | 6 +- docs/rust/reduce.rst | 149 +++ rust/Cargo.lock | 8 +- rust/cuda-coop-codegen/src/lib.rs | 79 +- rust/cuda-coop-codegen/src/oxide.rs | 460 +++++++- rust/cuda-coop-codegen/src/reduce.rs | 391 +++++++ rust/cuda-coop-codegen/tests/native.rs | 167 ++- rust/cuda-coop/Cargo.toml | 2 +- rust/cuda-coop/build.rs | 49 +- rust/cuda-coop/examples/load_store/Cargo.lock | 30 +- rust/cuda-coop/examples/load_store/Cargo.toml | 8 +- rust/cuda-coop/examples/reduce/.gitignore | 7 + rust/cuda-coop/examples/reduce/Cargo.lock | 1019 +++++++++++++++++ rust/cuda-coop/examples/reduce/Cargo.toml | 19 + .../examples/reduce/helper/Cargo.toml | 11 + .../examples/reduce/helper/src/lib.rs | 53 + rust/cuda-coop/examples/reduce/src/main.rs | 195 ++++ rust/cuda-coop/examples/scan/Cargo.lock | 30 +- rust/cuda-coop/examples/scan/Cargo.toml | 8 +- rust/cuda-coop/src/lib.rs | 10 + rust/cuda-coop/src/reduce.rs | 534 +++++++++ rust/cuda-coop/src/reduce/boundary.rs | 473 ++++++++ rust/cuda-coop/src/scan.rs | 10 +- rust/cuda-coop/tests/check_load_store.py | 137 ++- .../cuda-coop/tests/device_queries/Cargo.lock | 28 +- .../cuda-coop/tests/device_queries/Cargo.toml | 6 +- rust/cuda-coop/tests/load_store/Cargo.lock | 30 +- rust/cuda-coop/tests/load_store/Cargo.toml | 8 +- rust/cuda-coop/tests/reduce/.gitignore | 7 + rust/cuda-coop/tests/reduce/Cargo.lock | 1011 ++++++++++++++++ rust/cuda-coop/tests/reduce/Cargo.toml | 17 + rust/cuda-coop/tests/reduce/src/main.rs | 527 +++++++++ rust/cuda-coop/tests/scan/Cargo.lock | 30 +- rust/cuda-coop/tests/scan/Cargo.toml | 8 +- 34 files changed, 5388 insertions(+), 139 deletions(-) create mode 100644 docs/rust/reduce.rst create mode 100644 rust/cuda-coop-codegen/src/reduce.rs create mode 100644 rust/cuda-coop/examples/reduce/.gitignore create mode 100644 rust/cuda-coop/examples/reduce/Cargo.lock create mode 100644 rust/cuda-coop/examples/reduce/Cargo.toml create mode 100644 rust/cuda-coop/examples/reduce/helper/Cargo.toml create mode 100644 rust/cuda-coop/examples/reduce/helper/src/lib.rs create mode 100644 rust/cuda-coop/examples/reduce/src/main.rs create mode 100644 rust/cuda-coop/src/reduce.rs create mode 100644 rust/cuda-coop/src/reduce/boundary.rs create mode 100644 rust/cuda-coop/tests/reduce/.gitignore create mode 100644 rust/cuda-coop/tests/reduce/Cargo.lock create mode 100644 rust/cuda-coop/tests/reduce/Cargo.toml create mode 100644 rust/cuda-coop/tests/reduce/src/main.rs diff --git a/docs/rust/index.rst b/docs/rust/index.rst index ad383f5f66e5..53926726b2d0 100644 --- a/docs/rust/index.rst +++ b/docs/rust/index.rst @@ -43,8 +43,9 @@ as the remaining input length capped at the tile size. The kernel above shows that idiom, including a short final tile. Counts and offsets describe the whole group's tile, and must be uniform across that group. -The public operations include :doc:`load_store` and :doc:`scan`. Scan returns -fresh per-group prefixes while preserving its input. Additional primitive +The public operations include :doc:`load_store`, :doc:`scan` and :doc:`reduce`. +Scan returns fresh per-group prefixes; Reduce returns a total only at each +group's root. Both preserve their input. Additional primitive families are under development. The crate is unpublished and its interface can change. @@ -114,6 +115,7 @@ Read next groups load_store scan + reduce provider Generate the API reference with ``cargo doc --workspace --no-deps`` from diff --git a/docs/rust/reduce.rst b/docs/rust/reduce.rst new file mode 100644 index 000000000000..2b897ff5c321 --- /dev/null +++ b/docs/rust/reduce.rst @@ -0,0 +1,149 @@ +.. Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. +.. SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +Block and Warp Reduce +===================== + +``reduce`` combines a group's inputs and returns ``Option``: ``Some(total)`` +at group rank zero and ``None`` at every other participating thread. ``sum`` +fixes the operator to ``Sum``. Both functions take the group first, preserve +their inputs, and keep the input element type. Defaults select sum, all inputs, +and the Block ``WarpReductions`` algorithm. + +Block and Warp accept either a scalar or a shared or mutable borrow of +``ThreadData``. Payload items are ordered by x-fastest thread rank, then by +each thread's item index. Block supports complete one-, two- and +three-dimensional shapes. Warp supports physical warps and logical groups; +each group's total is independent. These operations do not reduce a whole +device-wide array into one value. + +A complete program +------------------ + +The example uses separate generic helpers, inferred payload elements, and +``ITEMS_PER_THREAD`` propagated from the host through the kernel. It checks +one and four items per thread over three blocks, including a short input tail +and full, short, and empty output buffers. A full-input path combines built-in +``Sum``, function-item wrapping addition, a stateless maximum closure, and +addition with a captured bias using one auto-sync scratch descriptor. It checks +all four results and preserves the inputs. Two runtime bias values use the same +compiled kernel; reducing a tile with bias adds ``(tile_size - 1) * bias`` to its +ordinary sum. + +.. literalinclude:: ../../rust/cuda-coop/examples/reduce/src/main.rs + :language: rust + :start-after: // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +The sum helper fills the missing input tail with an explicit zero. The callable +helper carries ``Op: ReduceOperator`` and passes the actual operator value +through the consuming ``.binary_op(op)`` builder. Every thread calls each +reduction. Only afterwards does the kernel inspect the retained results and +obtain block-output witnesses. + +.. literalinclude:: ../../rust/cuda-coop/examples/reduce/helper/src/lib.rs + :language: rust + :start-after: // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +Use the Oxide revision selected by the Cargo manifests and +``1.99.1``. From a complete CCCL checkout with that toolchain installed: + +.. code-block:: shell + + cd rust/cuda-coop/examples/reduce + 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. The compiler discovers the concrete CUB +specializations and scratch layout; the consumer needs no C++ shim, build script, +manual linker inputs, or scratch-byte calculation. + +Results and safe output +----------------------- + +``Some`` establishes that this thread has a valid numeric result. It supplies +neither global-memory ownership nor collective completion. For one output per +block, ``thread::block_index()`` grants a unique witness only to thread +``(0, 0, 0)``. ``DisjointSlice::get_mut`` checks the +output bounds before returning a mutable reference. A short or empty output +therefore leaves out-of-range blocks unwritten. This witness also supports +x-fastest indexing of two- and three-dimensional grids. + +Keep the reduction outside the terminal result branch so that all members +participate. A Warp result belongs to its own group's first lane; a block-output +witness does not provide a separate slot for every logical group. Results are +not broadcast, and nonroot numeric values are never read or replaced by zero. + +Operators, algorithms, and counts +--------------------------------- + +``ReduceOptions::default()`` and ``Default::default()`` infer the element type. +Consuming builders select static operators and Block algorithms: + +.. code-block:: rust + + let maximum = coop::reduce(block, &values, + coop::ReduceOptions::default().binary_op(coop::Max)); + let total = coop::sum(block, &values, + coop::SumOptions::default().algorithm(coop::Raking)); + let prefix_total = coop::sum(block, scalar, + coop::SumOptions::default().valid_items(count)); + +All ten numeric types (``i8``, ``u8``, ``i16``, ``u16``, ``i32``, ``u32``, +``i64``, ``u64``, ``f32``, ``f64``) support ``Sum``, ``Multiplies``, ``Min`` and +``Max``. Integer types additionally support ``BitAnd``, ``BitOr`` and +``BitXor``. Operators retain the input dtype; narrow integers are not widened. +CUB may regroup associative arithmetic, so floating-point totals need not +match a sequential host reduction exactly. A generic helper can carry +``Op: ReduceOperator`` and pass the actual operator value through +``.binary_op(op)``. + +Function items and closures implementing +``Fn(i32, i32) -> i32 + Copy`` support Block reduction with explicit ``Raking`` +and full input. The operator must be associative. Scalar input uses the same +callable path as a one-item payload. A closure may capture one ``i32`` by value +when the compiler proves that its value is identical for every thread in the +block at the call. The example captures a kernel scalar argument with ``move`` +after the preceding native operations. A lane-derived value or an unknown +value loaded from global memory does not establish this proof. Other captures, +function pointers, other types, Warp groups, other algorithms, and partial +counts are unsupported. + +Block algorithms are ``WarpReductions`` (default), ``Raking`` and +``RakingCommutativeOnly``. The last requires a commutative operator; all +currently supported builtins qualify. Warp uses its native algorithm and +rejects an explicit Block algorithm or temporary-storage descriptor. + +``.valid_items(count)`` is available only for scalar input, even when a payload +happens to contain one item. It selects the first ``count`` scalar contributors +and requires ``1 <= count <= group.count()``. Every group member still calls +the operation with the same count. Zero and oversized counts trap before +reduction. Payload input requires every slot initialized; an unwritten slot +traps before its numeric value is read. For a partial payload, fill the missing +items explicitly with a value appropriate to the operator. An empty Load at +offset zero with an explicit zero default produces a valid all-zero payload +sum; this differs from passing a zero scalar count. + +Warp membership and temporary storage +------------------------------------- + +The enclosing block must contain complete physical warps. Logical widths +1, 2, 4, 8, 16 and 32 divide each physical warp into independent groups. +Non-exhaustive widths 17 through 31 select the low-width prefix of each +physical warp. Use the genuine membership guard before the operation: + +.. code-block:: rust + + let group = coop::this_warp().group_by_non_exhaustive::<23>(); + if group.is_member() == 1 { + let total = coop::sum(group, &values, Default::default()); + // Consume this group's root result here. + } + +Block operations accept ``.temp_storage(&mut scratch)`` with the same +``TempStorage`` descriptor used by :doc:`load_store` and :doc:`scan`. +Omitted storage requests managed synchronization before dirty scratch reuse. +Explicit manual storage requires completion before conflicting reuse; +``auto_sync()`` adds a trailing Block barrier for scratch-using operations. +No reduction implicitly completes unrelated global Load/Store accesses. +Rust borrowing preserves local ownership; the matching compiler checks +group participation, uniform counts, geometry and scratch use. diff --git a/rust/Cargo.lock b/rust/Cargo.lock index 8861258fbb13..b4f1a410d583 100644 --- a/rust/Cargo.lock +++ b/rust/Cargo.lock @@ -57,7 +57,7 @@ dependencies = [ [[package]] name = "cuda-device" version = "0.2.1" -source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=732bec8e2166c33b9b851860363c81baa475203e#732bec8e2166c33b9b851860363c81baa475203e" dependencies = [ "cuda-macros", ] @@ -65,7 +65,7 @@ dependencies = [ [[package]] name = "cuda-macros" version = "0.2.1" -source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=732bec8e2166c33b9b851860363c81baa475203e#732bec8e2166c33b9b851860363c81baa475203e" dependencies = [ "oxide-tile-export", "proc-macro2", @@ -127,7 +127,7 @@ checksum = "cf8baf1c55e62ffcace7a9f06f4bd9cd3f0c4beb022d3b367256b91b87513d98" [[package]] name = "oxide-tile-export" version = "0.2.1" -source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=732bec8e2166c33b9b851860363c81baa475203e#732bec8e2166c33b9b851860363c81baa475203e" dependencies = [ "serde", "serde_json", @@ -155,7 +155,7 @@ dependencies = [ [[package]] name = "reserved-oxide-symbols" version = "0.2.1" -source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=732bec8e2166c33b9b851860363c81baa475203e#732bec8e2166c33b9b851860363c81baa475203e" [[package]] name = "serde" diff --git a/rust/cuda-coop-codegen/src/lib.rs b/rust/cuda-coop-codegen/src/lib.rs index ede8ed4450b6..aa62e89e6b2b 100644 --- a/rust/cuda-coop-codegen/src/lib.rs +++ b/rust/cuda-coop-codegen/src/lib.rs @@ -3,7 +3,7 @@ //! Native, host-only CUB specialization for Rust cooperative algorithms. //! -//! This private development crate compiles block and Warp Load/Store and Scan providers with +//! This private development crate compiles block and Warp Load/Store, Scan and Reduce 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. @@ -34,7 +34,12 @@ //! per lane with implicit algorithm selection. Scan checks every input validity //! bit, preserves the input, and initializes every fresh output slot. Seeded //! exclusive forms carry one runtime value of exactly the input dtype. -//! [`Provider::compile_native`] batches both families in one native compilation. +//! [`ReduceRequest`] preserves scalar or fixed-payload inputs and produces one +//! initialized scalar only at group rank zero. Scalar forms optionally take a +//! nonempty valid prefix; payload forms include every item. Blocks select +//! RakingCommutativeOnly, Raking or WarpReductions. Warp reductions also support +//! guarded low-lane prefixes of widths 17..31 within each physical warp. +//! [`Provider::compile_native`] batches these families in one native compilation. //! //! [`Provider::from_checkout`] snapshots the selected CCCL and CUDA headers; //! [`Provider::compile`] uses those exact bytes and loaded NVRTC to emit an @@ -64,7 +69,9 @@ #![deny(unsafe_code)] mod headers; +mod reduce; mod scan; +pub use reduce::{ReduceAlgorithm, ReduceForm, ReduceOperator, ReduceRecipe, ReduceRequest}; pub use scan::{ScanAlgorithm, ScanOperation, ScanOperator, ScanRecipe, ScanRequest}; #[allow(unsafe_code)] mod nvrtc; @@ -448,6 +455,8 @@ pub enum NativeRequest { LoadStore(Request), /// Numeric Scan with preserved input and fresh initialized output. Scan(ScanRequest), + /// Numeric reduction with preserved input and a fresh root-only scalar. + Reduce(ReduceRequest), } impl NativeRequest { @@ -455,6 +464,7 @@ impl NativeRequest { match self { Self::LoadStore(request) => source(request, symbol), Self::Scan(request) => request.source(symbol), + Self::Reduce(request) => request.source(symbol), } } @@ -462,13 +472,14 @@ impl NativeRequest { match self { Self::LoadStore(request) => request.operation.token().into(), Self::Scan(request) => request.recipe.operation().token(), + Self::Reduce(request) => request.recipe.token(), } } fn scratch_free(self) -> bool { match self { Self::LoadStore(request) => request.algorithm.scratch_free(), - Self::Scan(_) => false, + Self::Scan(_) | Self::Reduce(_) => false, } } @@ -477,8 +488,8 @@ impl NativeRequest { Self::LoadStore(request) => { !request.operation.is_load() && !request.algorithm.scratch_free() } - // Scan writes a fresh output, never rearranges its input payload. - Self::Scan(_) => false, + // Scan and Reduce write fresh outputs, never rearranging inputs. + Self::Scan(_) | Self::Reduce(_) => false, } } } @@ -497,8 +508,8 @@ pub struct Layout { /// # Device ABI obligations /// /// The memory-tile and payload-layout requirements below describe Load/Store. -/// Scan uses lane-private input and fresh output; its ABI and validity contract -/// are documented by [`ScanRequest`]. +/// Scan and Reduce use lane-private input and fresh output; their ABI and +/// validity contracts are documented by [`ScanRequest`] and [`ReduceRequest`]. /// /// Every thread in the selected group executes the provider in converged control /// flow. The group is the exact block, or one aligned Warp of the selected width. @@ -649,16 +660,32 @@ impl Provider { self.compile_native(&requests, target) } - /// Compile one ordered mixed Load/Store and Scan batch in a single NVRTC program. + /// Compile one ordered mixed Load/Store, Scan and Reduce batch in one NVRTC program. /// /// Per-specialization identity is independent of unrelated batch members. - /// Scan scratch describes one group and includes the actual CUB TempStorage, + /// Scan/Reduce scratch describes one group and includes the actual CUB TempStorage, /// even when its C++ empty type occupies one byte. The caller plans replication. /// /// # Errors /// Rejects empty/duplicate requests and malformed targets as [`Self::compile`] /// does; native compilation and layout decoding errors retain their diagnostics. pub fn compile_native(&self, requests: &[NativeRequest], target: &str) -> Result { + self.compile_native_with_callbacks(requests, target, None) + } + + fn compile_native_with_callbacks( + &self, + requests: &[NativeRequest], + target: &str, + callbacks: Option<&[oxide::Callback]>, + ) -> Result { + if let Some(callbacks) = callbacks { + if callbacks.len() != requests.len() { + return Err(Error::message( + "callback count must match specialization count", + )); + } + } if !target.starts_with("compute_") || !target[8..].chars().all(|c| c.is_ascii_alphanumeric()) || target.len() <= 8 @@ -671,7 +698,14 @@ impl Provider { return Err(Error::message("at least one specialization is required")); } for (index, request) in requests.iter().enumerate() { - if requests[..index].contains(request) { + if requests[..index] + .iter() + .enumerate() + .any(|(prior, candidate)| { + candidate == request + && callbacks.is_none_or(|callbacks| callbacks[prior] == callbacks[index]) + }) + { return Err(Error::message("duplicate specialization in one batch")); } } @@ -692,6 +726,7 @@ impl Provider { include_bytes!("nvrtc.rs"), include_bytes!("block_io.cuh"), include_bytes!("scan.rs"), + include_bytes!("reduce.rs"), ] { hash_field(&mut inputs, bytes); } @@ -703,17 +738,27 @@ impl Provider { 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 { + for (index, &request) in requests.iter().enumerate() { + let callback = callbacks.map(|callbacks| &callbacks[index]); + let source_for = |symbol: &str| match (request, callback) { + (NativeRequest::Reduce(reduce), Some(callback)) => { + reduce.source_with_callback(symbol, Some(&callback.symbol)) + } + (_, None) => request.source(symbol), + _ => unreachable!("callback recipes are validated before native compilation"), + }; let mut hash = inputs.clone(); - hash_field( - &mut hash, - request.source("cuda_coop_specialization").as_bytes(), - ); + hash_field(&mut hash, source_for("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.token(), suffix); - source.push_str(&request.source(&symbol)); + let token = if callback.is_some() { + "reduce_callable".to_owned() + } else { + request.token() + }; + let symbol = format!("cuda_coop_{token}_{suffix}"); + source.push_str(&source_for(&symbol)); expressions.push(format!("&layout_probe")); expressions.push(format!("&layout_probe")); specializations.push(Specialization { diff --git a/rust/cuda-coop-codegen/src/oxide.rs b/rust/cuda-coop-codegen/src/oxide.rs index 12807ec40b78..ba64ee917b78 100644 --- a/rust/cuda-coop-codegen/src/oxide.rs +++ b/rust/cuda-coop-codegen/src/oxide.rs @@ -1,23 +1,29 @@ // 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. +//! 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-request-v2` serves compiler-admitted callable Reduce requests; //! `--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 raw i32 and typed validity Load/Store recipes documented -//! by [`Operation`], and numeric Block/scalar Warp Scan through [`ScanRecipe`]. +//! by [`Operation`], numeric Block/scalar Warp Scan through [`ScanRecipe`], and +//! root-result scalar/payload Reduce through [`ReduceRecipe`]. //! Device calls retain all participation, validity and scratch obligations of //! [`crate::Specialization`]. No allocation or synchronization policy is supplied. +//! Callable Reduce uses i32 full-input Block Raking and the same payload validity +//! and root-result ABI. Its thunk and borrowed environment come from Oxide's +//! checked callback marker; the provider does not select or execute Rust code. use crate::{ - Algorithm, Artifact, Dtype, Error, Layout, NativeRequest, Operation, Provider, Request, Result, - ScanRecipe, ScanRequest, + Algorithm, Artifact, Dtype, Error, Layout, NativeRequest, Operation, Provider, ReduceAlgorithm, + ReduceForm, ReduceOperator, ReduceRecipe, ReduceRequest, Request, Result, ScanRecipe, + ScanRequest, }; use serde::{Deserialize, Serialize}; use std::io::{Read, Write}; @@ -99,11 +105,75 @@ enum Argument { #[derive(Debug, Deserialize)] #[serde(deny_unknown_fields)] -struct BatchRequest { +struct BatchRequest { version: u32, target: String, cuda_root: PathBuf, - requests: Vec, + requests: Vec, +} + +/// The compiler batches callback requests separately under protocol version two. +#[derive(Debug, Deserialize, PartialEq)] +#[serde(deny_unknown_fields)] +pub(crate) struct Callback { + pub(crate) symbol: String, + arguments: Vec, + result: Scalar, +} + +#[derive(Debug, Deserialize)] +#[serde(deny_unknown_fields)] +struct CallbackRequest { + operation: String, + algorithm: String, + dtype: Scalar, + block: [u32; 3], + items_per_thread: u32, + #[serde(default)] + warp_width: Option, + callback: Callback, +} + +impl CallbackRequest { + fn native(&self) -> Result { + if self.operation != "reduce_callable_payload_validity_bits_v1" + || self.algorithm != "raking" + || self.dtype != Scalar::Signed(32) + || self.warp_width.is_some() + { + return Err(Error::message( + "callable Reduce requires the i32 Block payload raking recipe", + )); + } + if self.callback.symbol.is_empty() + || !self + .callback + .symbol + .bytes() + .enumerate() + .all(|(i, b)| b == b'_' || b.is_ascii_alphabetic() || (i > 0 && b.is_ascii_digit())) + || self.callback.arguments + != [ + Argument::ConstPointer(Scalar::Unsigned(8)), + Argument::Value(Scalar::Signed(32)), + Argument::Value(Scalar::Signed(32)), + ] + || self.callback.result != Scalar::Signed(32) + { + return Err(Error::message( + "callable Reduce requires a C identifier and the (const environment*, i32, i32) -> i32 callback ABI", + )); + } + ReduceRequest::block( + Dtype::I32, + ReduceOperator::Sum, + ReduceAlgorithm::Raking, + ReduceForm::Payload, + self.block, + self.items_per_thread, + ) + .map(NativeRequest::Reduce) + } } #[derive(Debug, Deserialize)] @@ -120,6 +190,25 @@ struct SpecializationRequest { impl SpecializationRequest { fn native(&self) -> Result { + if self.operation.starts_with("reduce_") { + let recipe = ReduceRecipe::parse(&self.operation, &self.algorithm)?; + let dtype = self.dtype.dtype()?; + let request = match (recipe, self.warp_width) { + (ReduceRecipe::Block(op, algorithm, form), None) => ReduceRequest::block( + dtype, + op, + algorithm, + form, + self.block, + self.items_per_thread, + )?, + (ReduceRecipe::Warp(op, form), Some(width)) => { + ReduceRequest::warp(dtype, op, form, self.block, width, self.items_per_thread)? + } + _ => return Err(Error::message("Reduce recipe disagrees with Warp scope")), + }; + return Ok(NativeRequest::Reduce(request)); + } if self.operation.starts_with("scan_") { let recipe = ScanRecipe::parse(&self.operation, &self.algorithm)?; let dtype = self.dtype.dtype()?; @@ -226,6 +315,19 @@ struct BatchResponse { fn arguments(request: NativeRequest) -> Vec { let request = match request { + NativeRequest::Reduce(request) => { + let dtype = Scalar::from(request.dtype); + let mut arguments = vec![ + Argument::ConstPointer(dtype), + Argument::MutPointer(dtype), + Argument::ConstPointer(Scalar::Unsigned(8)), + Argument::MutPointer(Scalar::Unsigned(8)), + ]; + if request.recipe.form() == ReduceForm::ScalarPrefix { + arguments.push(Argument::Value(Scalar::Signed(32))); + } + return arguments; + } NativeRequest::Scan(request) => { let dtype = Scalar::from(request.dtype); let mut arguments = vec![ @@ -295,6 +397,14 @@ pub fn build_script() -> Result<()> { Some("--oxide-library-request-v1") => { return serve(&checkout, std::io::stdin().lock(), std::io::stdout().lock()); } + Some("--oxide-library-request-v2") => { + return serve_version( + &checkout, + std::io::stdin().lock(), + std::io::stdout().lock(), + 2, + ); + } Some("--oxide-library-inputs-v1") => { return serve_inputs(&checkout, std::io::stdin().lock(), std::io::stdout().lock()); } @@ -373,6 +483,50 @@ pub fn scan_registration(executable: &Path, recipe: ScanRecipe) -> Result 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: "reduce_callable_payload_validity_bits_v1", + algorithm: "raking", + }) + .map_err(|error| Error::message(format!("cannot encode provider registration: {error}"))) +} + +/// Encode one Reduce recipe for the existing dependency provider executable. +/// +/// ScalarPrefix adds one signed 32-bit uniform count. Scalar and Payload add +/// no controls; their distinct identities preserve scalar-only count support. +/// Warp recipes use implicit algorithm selection. Dtype, geometry and input +/// extent arrive with specialization; the output always has one slot. +/// +/// # Errors +/// Rejects a nonabsolute executable or a path that cannot be encoded as JSON. +pub fn reduce_registration(executable: &Path, recipe: ReduceRecipe) -> 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: &recipe.token(), + algorithm: recipe.algorithm_token(), + }) + .map_err(|error| Error::message(format!("cannot encode provider registration: {error}"))) +} + #[derive(Deserialize)] #[serde(deny_unknown_fields)] struct InputRequest { @@ -462,7 +616,16 @@ pub fn serve_inputs(checkout: &Path, input: impl Read, mut output: impl Write) - /// 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<()> { +pub fn serve(checkout: &Path, input: impl Read, output: impl Write) -> Result<()> { + serve_version(checkout, input, output, VERSION) +} + +fn serve_version( + checkout: &Path, + input: impl Read, + mut output: impl Write, + version: u32, +) -> Result<()> { let mut bytes = Vec::new(); input .take((MAX_HEADER + 1) as u64) @@ -470,6 +633,30 @@ pub fn serve(checkout: &Path, input: impl Read, mut output: impl Write) -> Resul if bytes.len() > MAX_HEADER { return Err(Error::message("provider request exceeds 1 MiB")); } + if version == 2 { + let request: BatchRequest = + serde_json::from_slice(&bytes).map_err(|error| { + Error::message(format!("invalid callback provider request: {error}")) + })?; + if request.version != version { + return Err(Error::message( + "callback provider requires protocol version 2", + )); + } + let requests = request + .requests + .iter() + .map(CallbackRequest::native) + .collect::>>()?; + let callbacks = request + .requests + .into_iter() + .map(|request| request.callback) + .collect::>(); + let artifact = Provider::from_checkout(checkout, &request.cuda_root)? + .compile_native_with_callbacks(&requests, &request.target, Some(&callbacks))?; + return write_response(&requests, artifact, &mut output, version); + } let request: BatchRequest = serde_json::from_slice(&bytes) .map_err(|error| Error::message(format!("invalid provider request: {error}")))?; if request.version != VERSION { @@ -485,13 +672,14 @@ pub fn serve(checkout: &Path, input: impl Read, mut output: impl Write) -> Resul .collect::>()?; let artifact = Provider::from_checkout(checkout, &request.cuda_root)? .compile_native(&requests, &request.target)?; - write_response(&requests, artifact, &mut output) + write_response(&requests, artifact, &mut output, version) } fn write_response( requests: &[NativeRequest], artifact: Artifact, output: &mut impl Write, + version: u32, ) -> Result<()> { if artifact.ltoir.is_empty() || artifact.ltoir.len() > MAX_LTO { return Err(Error::message( @@ -509,6 +697,8 @@ fn write_response( .map_err(|_| Error::message("request index exceeds u32"))?, identity: entry.identity, symbol: entry.symbol, + // Both protocols describe only the caller's payload/uniform + // tuple. Oxide inserts scratch and the V2 environment itself. arguments: arguments(*request), cub_layout: entry.cub_layout.try_into()?, scratch: entry.scratch.try_into()?, @@ -518,7 +708,7 @@ fn write_response( }) .collect::>>()?; let header = serde_json::to_vec(&BatchResponse { - version: VERSION, + version, batch_identity: artifact.identity, specializations, }) @@ -526,7 +716,7 @@ fn write_response( if header.len() > MAX_HEADER { return Err(Error::message("provider response header exceeds 1 MiB")); } - output.write_all(MAGIC)?; + output.write_all(if version == 2 { b"OXDLIB02" } else { 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)?; @@ -536,8 +726,256 @@ fn write_response( #[cfg(test)] mod tests { + fn callable_request(symbol: &str, items: u32) -> serde_json::Value { + serde_json::json!({ + "operation":"reduce_callable_payload_validity_bits_v1", "algorithm":"raking", + "dtype":{"kind":"signed","bits":32}, "block":[128,1,1], "items_per_thread":items, + "callback":{ + "symbol":symbol, + "arguments":[ + {"kind":"const_pointer","scalar":{"kind":"unsigned","bits":8}}, + {"kind":"value","scalar":{"kind":"signed","bits":32}}, + {"kind":"value","scalar":{"kind":"signed","bits":32}} + ], + "result":{"kind":"signed","bits":32} + } + }) + } + + #[test] + fn callable_protocol_rejects_wrong_recipe_or_thunk_before_toolkit() { + let base = callable_request("rust_reduce_thunk", 4); + assert!(serde_json::from_value::(base.clone()).is_err()); + let mut missing = base.clone(); + missing.as_object_mut().unwrap().remove("callback"); + assert!(serde_json::from_value::(missing).is_err()); + let parsed: CallbackRequest = serde_json::from_value(base.clone()).unwrap(); + assert!(parsed.native().is_ok()); + for (path, value) in [ + ( + "/operation", + serde_json::json!("reduce_sum_payload_validity_bits_v1"), + ), + ("/algorithm", serde_json::json!("warp_reductions")), + ("/dtype/bits", serde_json::json!(64)), + ("/block/0", serde_json::json!(0)), + ("/items_per_thread", serde_json::json!(0)), + ("/callback/symbol", serde_json::json!("bad;symbol")), + ( + "/callback/arguments/0/kind", + serde_json::json!("mut_pointer"), + ), + ("/callback/result/bits", serde_json::json!(64)), + ] { + let mut request = base.clone(); + *request.pointer_mut(path).unwrap() = value; + let request = serde_json::json!({"version":2, "target":"compute_120", "cuda_root":"/absent", "requests":[request]}); + let mut output = Vec::new(); + let error = serve_version( + Path::new("/absent"), + request.to_string().as_bytes(), + &mut output, + 2, + ) + .unwrap_err(); + assert!(!error.to_string().contains("No such file"), "{error}"); + assert!(output.is_empty()); + } + } + + #[test] + #[ignore = "requires CUDA_HOME and a Linux CUDA toolkit with NVRTC LTO support"] + fn native_callable_reduce_identity_and_v2_response() + -> std::result::Result<(), Box> { + let requests = [ + callable_request("rust_add", 1), + callable_request("rust_add", 4), + callable_request("rust_max", 4), + ] + .into_iter() + .map(serde_json::from_value::) + .collect::, _>>()?; + let native = requests + .iter() + .map(CallbackRequest::native) + .collect::>>()?; + let callbacks = requests + .into_iter() + .map(|request| request.callback) + .collect::>(); + let checkout = Path::new(env!("CARGO_MANIFEST_DIR")).join("../.."); + let provider = Provider::from_checkout(&checkout, std::env::var("CUDA_HOME")?)?; + let target = + std::env::var("CUDA_COOP_TEST_TARGET").unwrap_or_else(|_| "compute_120".into()); + let artifact = + provider.compile_native_with_callbacks(&native, &target, Some(&callbacks))?; + let identities = artifact + .specializations + .iter() + .map(|s| s.identity) + .collect::>(); + assert_eq!(identities.len(), 3); + assert!( + artifact + .specializations + .iter() + .all(|s| !s.scratch_free && !s.mutates_payload && s.scratch.size > 0) + ); + let builtin = provider.compile_native(&native[..1], &target)?; + assert!(!identities.contains(&builtin.specializations[0].identity)); + let mut response = Vec::new(); + write_response(&native, artifact, &mut response, 2)?; + assert_eq!(&response[..8], b"OXDLIB02"); + let length = u64::from_le_bytes(response[8..16].try_into()?) as usize; + let header: serde_json::Value = serde_json::from_slice(&response[24..24 + length])?; + assert_eq!(header["version"], 2); + for entry in header["specializations"].as_array().unwrap() { + assert_eq!( + entry["arguments"], + serde_json::json!([ + {"kind":"const_pointer","scalar":{"kind":"signed","bits":32}}, + {"kind":"mut_pointer","scalar":{"kind":"signed","bits":32}}, + {"kind":"const_pointer","scalar":{"kind":"unsigned","bits":8}}, + {"kind":"mut_pointer","scalar":{"kind":"unsigned","bits":8}} + ]) + ); + } + Ok(()) + } + use super::*; - use crate::{ScanAlgorithm, ScanOperation, ScanOperator}; + use crate::{ReduceAlgorithm, ReduceOperator, ScanAlgorithm, ScanOperation, ScanOperator}; + + #[test] + fn reduce_registration_preserves_scalar_count_and_payload_abi() { + for (dtype, recipe, items) in [ + ( + Dtype::I8, + ReduceRecipe::Block( + ReduceOperator::Sum, + ReduceAlgorithm::WarpReductions, + ReduceForm::Payload, + ), + 1, + ), + ( + Dtype::U16, + ReduceRecipe::Block( + ReduceOperator::Min, + ReduceAlgorithm::Raking, + ReduceForm::ScalarPrefix, + ), + 1, + ), + ( + Dtype::F64, + ReduceRecipe::Warp(ReduceOperator::Max, ReduceForm::Scalar), + 1, + ), + ( + Dtype::I32, + ReduceRecipe::Warp(ReduceOperator::BitXor, ReduceForm::Payload), + 9, + ), + ] { + let registration: serde_json::Value = + serde_json::from_str(&reduce_registration(Path::new("/provider"), recipe).unwrap()) + .unwrap(); + let scalar = Scalar::from(dtype); + let request = SpecializationRequest { + operation: registration["operation"].as_str().unwrap().into(), + algorithm: registration["algorithm"].as_str().unwrap().into(), + dtype: scalar, + block: [4, 4, 4], + items_per_thread: items, + warp_width: matches!(recipe, ReduceRecipe::Warp(..)).then_some(17), + } + .native() + .unwrap(); + let mut expected = vec![ + Argument::ConstPointer(scalar), + Argument::MutPointer(scalar), + Argument::ConstPointer(Scalar::Unsigned(8)), + Argument::MutPointer(Scalar::Unsigned(8)), + ]; + if recipe.form() == ReduceForm::ScalarPrefix { + expected.push(Argument::Value(Scalar::Signed(32))); + } + assert_eq!(arguments(request), expected); + assert_eq!(registration["provider"], "cuda_coop"); + } + } + + #[test] + fn rejects_invalid_reduce_before_opening_toolkit() { + let base = serde_json::json!({"version":1,"target":"compute_120","cuda_root":"/absent", "requests":[{ + "operation":"reduce_bit_xor_payload_validity_bits_v1","algorithm":"raking", + "dtype":{"kind":"signed","bits":32},"block":[8,4,2],"items_per_thread":9,"warp_width":null + }]}); + for (path, value, diagnostic) in [ + ( + "/requests/0/dtype", + serde_json::json!({"kind":"float","bits":32}), + "integer", + ), + ( + "/requests/0/algorithm", + serde_json::json!("raking_memoize"), + "Reduce algorithm", + ), + ( + "/requests/0/operation", + serde_json::json!("reduce_sum_prefix_validity_bits_v1"), + "exactly one", + ), + ( + "/requests/0/operation", + serde_json::json!("reduce_sum_payload_prefix_validity_bits_v1"), + "Reduce operator", + ), + ( + "/requests/0/items_per_thread", + serde_json::json!(0), + "positive items", + ), + ( + "/requests/0/warp_width", + serde_json::json!(17), + "Warp scope", + ), + ( + "/requests/0/block", + serde_json::json!([u32::MAX, 4, 4]), + "positive axes", + ), + ] { + let mut request = base.clone(); + *request.pointer_mut(path).unwrap() = value; + let error = serve( + Path::new("/absent"), + request.to_string().as_bytes(), + Vec::new(), + ) + .unwrap_err(); + assert!(error.to_string().contains(diagnostic), "{error}"); + } + for (block, width, diagnostic) in [ + ([48, 1, 1], 17, "complete physical warps"), + ([64, 1, 1], 15, "guarded prefix"), + ] { + let mut request = base.clone(); + request["requests"][0]["algorithm"] = "warp".into(); + request["requests"][0]["block"] = serde_json::json!(block); + request["requests"][0]["warp_width"] = width.into(); + let error = serve( + Path::new("/absent"), + request.to_string().as_bytes(), + Vec::new(), + ) + .unwrap_err(); + assert!(error.to_string().contains(diagnostic), "{error}"); + } + } #[test] fn scan_registration_preserves_typed_seed_and_empty_controls() { diff --git a/rust/cuda-coop-codegen/src/reduce.rs b/rust/cuda-coop-codegen/src/reduce.rs new file mode 100644 index 000000000000..6c6c9dbcd415 --- /dev/null +++ b/rust/cuda-coop-codegen/src/reduce.rs @@ -0,0 +1,391 @@ +// Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +//! Validated root-result Reduce recipes, separate from Scan and Load/Store. + +use crate::{Dtype, Error, Result}; + +/// Static associative operator used by a native Reduce specialization. +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub enum ReduceOperator { + /// Addition in the input dtype. + Sum, + /// Multiplication in the input dtype. + Multiplies, + /// Minimum. + Min, + /// Maximum. + Max, + /// Integer bitwise AND. + BitAnd, + /// Integer bitwise OR. + BitOr, + /// Integer bitwise XOR. + BitXor, +} + +impl ReduceOperator { + fn token(self) -> &'static str { + match self { + Self::Sum => "sum", + Self::Multiplies => "multiplies", + Self::Min => "min", + Self::Max => "max", + Self::BitAnd => "bit_and", + Self::BitOr => "bit_or", + Self::BitXor => "bit_xor", + } + } + + fn parse(token: &str) -> Result { + match token { + "sum" => Ok(Self::Sum), + "multiplies" => Ok(Self::Multiplies), + "min" => Ok(Self::Min), + "max" => Ok(Self::Max), + "bit_and" => Ok(Self::BitAnd), + "bit_or" => Ok(Self::BitOr), + "bit_xor" => Ok(Self::BitXor), + _ => Err(Error::message("unsupported Reduce operator")), + } + } + + fn functor(self) -> &'static str { + match self { + Self::Sum => "cuda::std::plus", + Self::Multiplies => "cuda::std::multiplies", + Self::Min => "cuda::minimum", + Self::Max => "cuda::maximum", + Self::BitAnd => "cuda::std::bit_and", + Self::BitOr => "cuda::std::bit_or", + Self::BitXor => "cuda::std::bit_xor", + } + } +} + +/// Input form; a one-item payload remains distinct from a scalar. +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub enum ReduceForm { + /// One initialized scalar from every member. + Scalar, + /// One initialized scalar per member, with a group-uniform valid count. + ScalarPrefix, + /// Every initialized item in a fixed per-thread payload. + Payload, +} + +/// CUB algorithm selected for a complete Block reduction. +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub enum ReduceAlgorithm { + /// Raking specialized for the supported commutative builtins. + RakingCommutativeOnly, + /// Raking reduction. + Raking, + /// Combine physical-warp reductions, the default Block algorithm. + WarpReductions, +} + +/// Registration identity independent of dtype, geometry and input extent. +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub enum ReduceRecipe { + /// Complete Block operation with its selected algorithm and input form. + Block(ReduceOperator, ReduceAlgorithm, ReduceForm), + /// Physical/logical Warp operation with implicit algorithm selection. + Warp(ReduceOperator, ReduceForm), +} + +impl ReduceRecipe { + /// Builtin selected by this recipe. + pub fn operator(self) -> ReduceOperator { + match self { + Self::Block(op, _, _) | Self::Warp(op, _) => op, + } + } + + /// Scalar or fixed-payload input form. + pub fn form(self) -> ReduceForm { + match self { + Self::Block(_, _, form) | Self::Warp(_, form) => form, + } + } + + pub(crate) fn token(self) -> String { + let form = match self.form() { + ReduceForm::Scalar => "scalar", + ReduceForm::ScalarPrefix => "prefix", + ReduceForm::Payload => "payload", + }; + format!("reduce_{}_{form}_validity_bits_v1", self.operator().token()) + } + + pub(crate) fn algorithm_token(self) -> &'static str { + match self { + Self::Block(_, ReduceAlgorithm::RakingCommutativeOnly, _) => "raking_commutative_only", + Self::Block(_, ReduceAlgorithm::Raking, _) => "raking", + Self::Block(_, ReduceAlgorithm::WarpReductions, _) => "warp_reductions", + Self::Warp(_, _) => "warp", + } + } + + pub(crate) fn parse(operation: &str, algorithm: &str) -> Result { + let operation = operation + .strip_prefix("reduce_") + .and_then(|s| s.strip_suffix("_validity_bits_v1")) + .ok_or_else(|| Error::message("unsupported Reduce operation or input form"))?; + let (operator, form) = operation + .rsplit_once('_') + .ok_or_else(|| Error::message("missing Reduce input form"))?; + let op = ReduceOperator::parse(operator)?; + let form = match form { + "scalar" => ReduceForm::Scalar, + "prefix" => ReduceForm::ScalarPrefix, + "payload" => ReduceForm::Payload, + _ => return Err(Error::message("unsupported Reduce input form")), + }; + match algorithm { + "raking_commutative_only" => Ok(Self::Block( + op, + ReduceAlgorithm::RakingCommutativeOnly, + form, + )), + "raking" => Ok(Self::Block(op, ReduceAlgorithm::Raking, form)), + "warp_reductions" => Ok(Self::Block(op, ReduceAlgorithm::WarpReductions, form)), + "warp" => Ok(Self::Warp(op, form)), + _ => Err(Error::message("unsupported Reduce algorithm")), + } + } +} + +/// Validated reduction with an initialized scalar result only at group rank zero. +/// +/// Native calls preserve all input values and packed validity bits. The ABI is +/// `(const T* input, T* output, u8* scratch, const u8* input_bits, +/// u8* output_bits[, i32 valid])`. Input has the declared item extent and +/// `ceil(items/8)` low-bit-first validity bytes; output has exactly one T slot +/// and one validity byte. All pointers are live, aligned lane-local storage; +/// output is fresh and disjoint from input. Every input bit is checked before +/// numeric input is read. ScalarPrefix requires a group-uniform count in +/// `1..=group_size`, and every scalar is initialized even outside that prefix. +/// Only rank zero writes the result and sets its byte to one. Other members' +/// result bytes are zero and their numeric output slots remain uninitialized. +/// All members participate; non-power-of-two Warps mean exactly the low-width +/// prefix of each physical warp. The caller owns scratch replication and reuse +/// completion. The provider adds no completion barrier or result broadcast. +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub struct ReduceRequest { + pub(crate) recipe: ReduceRecipe, + pub(crate) dtype: Dtype, + block: [u32; 3], + items: u32, + width: Option, +} + +impl ReduceRequest { + /// Validate a complete Block reduction with exact XYZ geometry. + /// + /// # Errors + /// Rejects invalid axes, more than 1024 threads, invalid item extents, + /// scalar forms with more than one item, and floating bitwise operators. + pub fn block( + dtype: Dtype, + operator: ReduceOperator, + algorithm: ReduceAlgorithm, + form: ReduceForm, + block: [u32; 3], + items: u32, + ) -> Result { + Self::validate( + dtype, + ReduceRecipe::Block(operator, algorithm, form), + block, + items, + None, + ) + } + + /// Validate a Warp reduction within complete physical warps. + /// + /// Widths 1/2/4/8/16/32 select aligned partitions; widths 17..31 select one + /// participating low-lane prefix per physical warp. This native validation + /// does not establish the caller's required membership or participation. + /// + /// # Errors + /// Rejects invalid Block geometry, incomplete physical warps, unsupported + /// widths, invalid input extents and floating bitwise operators. + pub fn warp( + dtype: Dtype, + operator: ReduceOperator, + form: ReduceForm, + block: [u32; 3], + width: u32, + items: u32, + ) -> Result { + Self::validate( + dtype, + ReduceRecipe::Warp(operator, form), + block, + items, + Some(width), + ) + } + + fn validate( + dtype: Dtype, + recipe: ReduceRecipe, + block: [u32; 3], + items: u32, + width: Option, + ) -> Result { + let threads = block + .into_iter() + .try_fold(1_u32, u32::checked_mul) + .filter(|n| (1..=1024).contains(n)) + .ok_or_else(|| { + Error::message("Reduce requires 1..=1024 block threads with positive axes") + })?; + if items == 0 || u64::from(threads) * u64::from(items) > i32::MAX as u64 { + return Err(Error::message( + "Reduce requires positive items and tile extent <= i32::MAX", + )); + } + if recipe.form() != ReduceForm::Payload && items != 1 { + return Err(Error::message( + "scalar Reduce forms require exactly one input item", + )); + } + if let Some(width) = width { + if threads % 32 != 0 { + return Err(Error::message( + "Warp Reduce requires complete physical warps", + )); + } + if width == 0 || width > 32 || (!width.is_power_of_two() && width < 17) { + return Err(Error::message( + "Warp Reduce requires power-of-two width 1..=32 or guarded prefix 17..=31", + )); + } + } + if matches!(dtype, Dtype::F32 | Dtype::F64) + && matches!( + recipe.operator(), + ReduceOperator::BitAnd | ReduceOperator::BitOr | ReduceOperator::BitXor + ) + { + return Err(Error::message( + "bitwise Reduce operators require integer elements", + )); + } + Ok(Self { + recipe, + dtype, + block, + items, + width, + }) + } + + pub(crate) fn source(self, symbol: &str) -> String { + self.source_with_callback(symbol, None) + } + + pub(crate) fn source_with_callback(self, symbol: &str, callback: Option<&str>) -> String { + let (dtype, _) = self.dtype.native(); + let [x, y, z] = self.block; + let items = self.items; + let primitive = match self.recipe { + ReduceRecipe::Block(_, algorithm, _) => { + let algorithm = match algorithm { + ReduceAlgorithm::RakingCommutativeOnly => { + "BLOCK_REDUCE_RAKING_COMMUTATIVE_ONLY" + } + ReduceAlgorithm::Raking => "BLOCK_REDUCE_RAKING", + ReduceAlgorithm::WarpReductions => "BLOCK_REDUCE_WARP_REDUCTIONS", + }; + format!("cub::BlockReduce") + } + ReduceRecipe::Warp(_, _) => { + format!("cub::WarpReduce", self.width.unwrap()) + } + }; + let counted = self.recipe.form() == ReduceForm::ScalarPrefix; + let controls = if callback.is_some() { + ", const void* environment" + } else if counted { + ", int valid" + } else { + "" + }; + let group_size = self.width.unwrap_or(x * y * z); + let validate_count = if counted { + format!("if (valid < 1 || valid > {group_size}) __trap();") + } else { + String::new() + }; + let operand = if self.recipe.form() == ReduceForm::Payload { + "local_input" + } else { + "local_input[0]" + }; + let count = if counted { ", valid" } else { "" }; + let callback_declaration = callback.map_or_else(String::new, |name| { + format!("extern \"C\" __device__ int {name}(const void*, int, int);\n") + }); + let callback_functor = callback.map_or_else(String::new, |name| { + format!( + "struct binary_op {{ const void* environment;\n __device__ int operator()(int a, int b) const {{ return {name}(environment, a, b); }}\n}};" + ) + }); + let call = if callback.is_some() { + format!("operation.Reduce({operand}, binary_op{{environment}})") + } else if self.recipe.operator() == ReduceOperator::Sum { + format!("operation.Sum({operand}{count})") + } else { + format!( + "operation.Reduce({operand}, {}{{}}{count})", + self.recipe.operator().functor() + ) + }; + let root = match self.width { + None => "thread_rank == 0".into(), + Some(w) => format!( + "thread_rank % {} == 0", + if w.is_power_of_two() { w } else { 32 } + ), + }; + format!( + r#" +#include +#include +#include +#include +#include +{callback_declaration} +namespace {symbol}_detail {{ +using value_type = {dtype}; +using primitive = {primitive}; +using scratch_storage = primitive::TempStorage; +{callback_functor} +}} +extern "C" __device__ void {symbol}(const {dtype}* input, {dtype}* output, + unsigned char* scratch, const unsigned char* input_bits, unsigned char* output_bits{controls}) {{ + using namespace {symbol}_detail; + const unsigned int thread_rank = cub::RowMajorTid({x}, {y}, {z}); + output_bits[0] = 0; + {validate_count} +#pragma unroll + for (int i = 0; i < {items}; ++i) + if ((input_bits[i / 8] & (1u << (i % 8))) == 0) __trap(); + value_type local_input[{items}]; +#pragma unroll + for (int i = 0; i < {items}; ++i) local_input[i] = input[i]; + primitive operation{{*reinterpret_cast(scratch)}}; + const value_type aggregate = {call}; + if ({root}) {{ + output[0] = aggregate; + output_bits[0] = 1; + }} +}} +"# + ) + } +} diff --git a/rust/cuda-coop-codegen/tests/native.rs b/rust/cuda-coop-codegen/tests/native.rs index 59acafc02b5f..d8ebbd209a13 100644 --- a/rust/cuda-coop-codegen/tests/native.rs +++ b/rust/cuda-coop-codegen/tests/native.rs @@ -4,11 +4,174 @@ //! Real native compiler coverage; ordinary host tests need no CUDA installation. use cuda_coop_codegen::{ - Algorithm, Dtype, Layout, NativeRequest, Operation, Provider, Request, ScanAlgorithm, - ScanOperation, ScanOperator, ScanRequest, + Algorithm, Dtype, Layout, NativeRequest, Operation, Provider, ReduceAlgorithm, ReduceForm, + ReduceOperator, ReduceRequest, Request, ScanAlgorithm, ScanOperation, ScanOperator, + ScanRequest, }; use std::{env, path::Path, time::Instant}; +#[test] +#[ignore = "requires CUDA_HOME and a Linux CUDA toolkit with NVRTC LTO support"] +fn native_mixed_reduce_recipes() -> 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 requests = [1, 4].map(|items| { + ReduceRequest::block( + Dtype::I32, + ReduceOperator::Sum, + ReduceAlgorithm::WarpReductions, + ReduceForm::Payload, + [128, 1, 1], + items, + ) + .map(NativeRequest::Reduce) + }); + let mut requests = requests + .into_iter() + .collect::>>()?; + // Distribute types, operators, forms and algorithms rather than multiplying + // independent dimensions. Payload extents cross the packed-mask byte boundary. + let types = [ + Dtype::I8, + Dtype::U8, + Dtype::I16, + Dtype::U16, + Dtype::I32, + Dtype::U32, + Dtype::I64, + Dtype::U64, + Dtype::F32, + Dtype::F64, + ]; + let operators = [ + ReduceOperator::Sum, + ReduceOperator::Multiplies, + ReduceOperator::BitAnd, + ReduceOperator::BitOr, + ReduceOperator::BitXor, + ReduceOperator::Min, + ReduceOperator::Max, + ]; + let algorithms = [ + ReduceAlgorithm::RakingCommutativeOnly, + ReduceAlgorithm::Raking, + ReduceAlgorithm::WarpReductions, + ]; + let forms = [ + ReduceForm::Scalar, + ReduceForm::ScalarPrefix, + ReduceForm::Payload, + ]; + for (index, dtype) in types.into_iter().enumerate() { + let operator = match dtype { + Dtype::F32 => ReduceOperator::Min, + Dtype::F64 => ReduceOperator::Max, + _ => operators[index % operators.len()], + }; + let form = forms[index % forms.len()]; + requests.push(NativeRequest::Reduce(ReduceRequest::block( + dtype, + operator, + algorithms[index % algorithms.len()], + form, + if index % 2 == 0 { + [4, 4, 4] + } else { + [11, 3, 1] + }, + if form == ReduceForm::Payload { + if index % 2 == 0 { 4 } else { 9 } + } else { + 1 + }, + )?)); + } + for (index, width) in [1, 2, 4, 8, 16].into_iter().chain(17..=32).enumerate() { + let dtype = types[index % types.len()]; + let operator = match dtype { + Dtype::F32 => ReduceOperator::Min, + Dtype::F64 => ReduceOperator::Max, + _ => operators[index % operators.len()], + }; + let form = forms[index % forms.len()]; + requests.push(NativeRequest::Reduce(ReduceRequest::warp( + dtype, + operator, + form, + [4, 4, 4], + width, + if form == ReduceForm::Payload { + if index % 2 == 0 { 4 } else { 9 } + } else { + 1 + }, + )?)); + } + requests.push(NativeRequest::LoadStore(Request::new_typed( + Dtype::I32, + Operation::LoadFullValidity, + Algorithm::Direct, + 64, + 4, + )?)); + requests.push(NativeRequest::Scan(ScanRequest::block( + Dtype::I32, + ScanOperation::ExclusiveSum, + ScanAlgorithm::Raking, + [4, 4, 4], + 4, + )?)); + let artifact = provider.compile_native(&requests, &target)?; + assert_eq!(artifact.specializations.len(), requests.len()); + assert!(!artifact.ltoir.is_empty()); + let mut identities = std::collections::HashSet::new(); + let mut symbols = std::collections::HashSet::new(); + for (request, entry) in requests.iter().zip(&artifact.specializations) { + assert!(identities.insert(entry.identity)); + assert!(symbols.insert(&entry.symbol)); + if matches!(request, NativeRequest::LoadStore(_)) { + assert!(entry.scratch_free); + assert_eq!(entry.scratch.size, 0); + } else { + assert_eq!(entry.cub_layout, entry.scratch); + assert!(entry.scratch.size > 0 && entry.scratch.alignment.is_power_of_two()); + assert!(!entry.scratch_free && !entry.mutates_payload); + } + eprintln!( + "{request:?}: CUB {:?}, scratch {:?}", + entry.cub_layout, entry.scratch + ); + } + let single = provider.compile_native(&requests[..1], &target)?; + assert_eq!( + single.specializations[0].identity, + artifact.specializations[0].identity + ); + assert_eq!( + single.specializations[0].symbol, + artifact.specializations[0].symbol + ); + assert_eq!( + single.specializations[0].scratch, + artifact.specializations[0].scratch + ); + assert!( + provider + .compile_native(&[requests[0]; 2], &target) + .unwrap_err() + .to_string() + .contains("duplicate") + ); + eprintln!( + "{} mixed Reduce/Scan/Load specializations, {} LTO bytes", + requests.len(), + artifact.ltoir.len() + ); + Ok(()) +} + #[test] #[ignore = "requires CUDA_HOME and a Linux CUDA toolkit with NVRTC LTO support"] fn native_mixed_scan_recipes() -> Result<(), Box> { diff --git a/rust/cuda-coop/Cargo.toml b/rust/cuda-coop/Cargo.toml index 45682aa419bd..8580098be28d 100644 --- a/rust/cuda-coop/Cargo.toml +++ b/rust/cuda-coop/Cargo.toml @@ -9,7 +9,7 @@ license.workspace = true publish = false [dependencies] -cuda-device = { git = "https://github.com/tpn/cuda-rust.git", rev = "96caed46c776eb49226137eac849e795b83b7ce0" } +cuda-device = { git = "https://github.com/tpn/cuda-rust.git", rev = "732bec8e2166c33b9b851860363c81baa475203e" } [build-dependencies] cuda-coop-codegen = { path = "../cuda-coop-codegen" } diff --git a/rust/cuda-coop/build.rs b/rust/cuda-coop/build.rs index 6d85c64cc097..89a9817518b9 100644 --- a/rust/cuda-coop/build.rs +++ b/rust/cuda-coop/build.rs @@ -1,20 +1,63 @@ // 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 Load, Store and Scan. +//! Register the owning host provider for native Load, Store, Scan and Reduce. use cuda_coop_codegen::{ - Algorithm, Operation, ScanAlgorithm, ScanOperation, ScanOperator, ScanRecipe, oxide, + Algorithm, Operation, ReduceAlgorithm, ReduceForm, ReduceOperator, ReduceRecipe, ScanAlgorithm, + ScanOperation, ScanOperator, ScanRecipe, 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") + Some( + "--oxide-library-request-v1" + | "--oxide-library-request-v2" + | "--oxide-library-inputs-v1" + ) ) { return; } let executable = std::env::current_exe().expect("cuda.coop provider executable"); + let callable = oxide::reduce_callable_registration(&executable) + .expect("cuda.coop callable Reduce registration"); + println!("cargo::rustc-env=COOP_REDUCE_RAKING_CALLABLE_PAYLOAD={callable}"); + for (operator, name) in [ + (ReduceOperator::Sum, "SUM"), + (ReduceOperator::Multiplies, "MULTIPLIES"), + (ReduceOperator::Min, "MIN"), + (ReduceOperator::Max, "MAX"), + (ReduceOperator::BitAnd, "BIT_AND"), + (ReduceOperator::BitOr, "BIT_OR"), + (ReduceOperator::BitXor, "BIT_XOR"), + ] { + for (form, suffix) in [ + (ReduceForm::Scalar, "SCALAR"), + (ReduceForm::ScalarPrefix, "PREFIX"), + (ReduceForm::Payload, "PAYLOAD"), + ] { + for (recipe, prefix) in [ + ( + ReduceRecipe::Block(operator, ReduceAlgorithm::RakingCommutativeOnly, form), + "RAKING_COMMUTATIVE_ONLY", + ), + ( + ReduceRecipe::Block(operator, ReduceAlgorithm::Raking, form), + "RAKING", + ), + ( + ReduceRecipe::Block(operator, ReduceAlgorithm::WarpReductions, form), + "WARP_REDUCTIONS", + ), + (ReduceRecipe::Warp(operator, form), "WARP"), + ] { + let value = oxide::reduce_registration(&executable, recipe) + .expect("cuda.coop reduce registration"); + println!("cargo::rustc-env=COOP_REDUCE_{prefix}_{name}_{suffix}={value}"); + } + } + } for (operation, name) in [ (ScanOperation::ExclusiveSum, "EXCLUSIVE_ZERO"), (ScanOperation::Exclusive(ScanOperator::Sum), "EXCLUSIVE_SUM"), diff --git a/rust/cuda-coop/examples/load_store/Cargo.lock b/rust/cuda-coop/examples/load_store/Cargo.lock index 98b133dcd9dd..785c01edc7dc 100644 --- a/rust/cuda-coop/examples/load_store/Cargo.lock +++ b/rust/cuda-coop/examples/load_store/Cargo.lock @@ -143,7 +143,7 @@ dependencies = [ [[package]] name = "cuda-artifact-finalizer" version = "0.2.1" -source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=732bec8e2166c33b9b851860363c81baa475203e#732bec8e2166c33b9b851860363c81baa475203e" dependencies = [ "libnvvm-sys", "nvjitlink-sys", @@ -155,7 +155,7 @@ dependencies = [ [[package]] name = "cuda-async" version = "0.5.0" -source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=732bec8e2166c33b9b851860363c81baa475203e#732bec8e2166c33b9b851860363c81baa475203e" dependencies = [ "anyhow", "cuda-bindings", @@ -171,7 +171,7 @@ dependencies = [ [[package]] name = "cuda-bindings" version = "0.5.0" -source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=732bec8e2166c33b9b851860363c81baa475203e#732bec8e2166c33b9b851860363c81baa475203e" dependencies = [ "bindgen", "libloading", @@ -202,7 +202,7 @@ dependencies = [ [[package]] name = "cuda-core" version = "0.5.0" -source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=732bec8e2166c33b9b851860363c81baa475203e#732bec8e2166c33b9b851860363c81baa475203e" dependencies = [ "anyhow", "cuda-bindings", @@ -214,7 +214,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" version = "0.5.0" -source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=732bec8e2166c33b9b851860363c81baa475203e#732bec8e2166c33b9b851860363c81baa475203e" dependencies = [ "proc-macro2", "quote", @@ -224,7 +224,7 @@ dependencies = [ [[package]] name = "cuda-device" version = "0.2.1" -source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=732bec8e2166c33b9b851860363c81baa475203e#732bec8e2166c33b9b851860363c81baa475203e" dependencies = [ "cuda-macros", ] @@ -232,7 +232,7 @@ dependencies = [ [[package]] name = "cuda-host" version = "0.2.1" -source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=732bec8e2166c33b9b851860363c81baa475203e#732bec8e2166c33b9b851860363c81baa475203e" dependencies = [ "cuda-artifact-finalizer", "cuda-async", @@ -250,7 +250,7 @@ dependencies = [ [[package]] name = "cuda-macros" version = "0.2.1" -source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=732bec8e2166c33b9b851860363c81baa475203e#732bec8e2166c33b9b851860363c81baa475203e" dependencies = [ "oxide-tile-export", "proc-macro2", @@ -264,7 +264,7 @@ dependencies = [ [[package]] name = "cuda-target-spec" version = "0.1.0" -source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=732bec8e2166c33b9b851860363c81baa475203e#732bec8e2166c33b9b851860363c81baa475203e" [[package]] name = "dashmap" @@ -519,7 +519,7 @@ dependencies = [ [[package]] name = "libnvvm-sys" version = "0.2.1" -source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=732bec8e2166c33b9b851860363c81baa475203e#732bec8e2166c33b9b851860363c81baa475203e" dependencies = [ "cuda-target-spec", "libloading", @@ -572,7 +572,7 @@ dependencies = [ [[package]] name = "nvjitlink-sys" version = "0.2.1" -source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=732bec8e2166c33b9b851860363c81baa475203e#732bec8e2166c33b9b851860363c81baa475203e" dependencies = [ "libloading", "thiserror 2.0.21", @@ -599,7 +599,7 @@ checksum = "9f7c3e4beb33f85d45ae3e3a1792185706c8e16d043238c593331cc7cd313b50" [[package]] name = "oxide-artifacts" version = "0.2.1" -source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=732bec8e2166c33b9b851860363c81baa475203e#732bec8e2166c33b9b851860363c81baa475203e" dependencies = [ "object", ] @@ -607,7 +607,7 @@ dependencies = [ [[package]] name = "oxide-tile-export" version = "0.2.1" -source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=732bec8e2166c33b9b851860363c81baa475203e#732bec8e2166c33b9b851860363c81baa475203e" dependencies = [ "serde", "serde_json", @@ -655,7 +655,7 @@ dependencies = [ [[package]] name = "ptx-parse" version = "0.1.0" -source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=732bec8e2166c33b9b851860363c81baa475203e#732bec8e2166c33b9b851860363c81baa475203e" [[package]] name = "quote" @@ -707,7 +707,7 @@ checksum = "d6f6ff9a378485b298a5286656da665ba74413d36db0979633275d2e708145d4" [[package]] name = "reserved-oxide-symbols" version = "0.2.1" -source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=732bec8e2166c33b9b851860363c81baa475203e#732bec8e2166c33b9b851860363c81baa475203e" [[package]] name = "rustc-hash" diff --git a/rust/cuda-coop/examples/load_store/Cargo.toml b/rust/cuda-coop/examples/load_store/Cargo.toml index 45b0a1df3f9b..4b4a0d8c327b 100644 --- a/rust/cuda-coop/examples/load_store/Cargo.toml +++ b/rust/cuda-coop/examples/load_store/Cargo.toml @@ -13,7 +13,7 @@ members = ["helper"] [dependencies] cuda-coop = { path = "../.." } coop-load-store-helper = { path = "helper" } -cuda-device = { git = "https://github.com/tpn/cuda-rust.git", rev = "96caed46c776eb49226137eac849e795b83b7ce0" } -cuda-host = { git = "https://github.com/tpn/cuda-rust.git", rev = "96caed46c776eb49226137eac849e795b83b7ce0", features = ["async"] } -cuda-core = { git = "https://github.com/tpn/cuda-rust.git", rev = "96caed46c776eb49226137eac849e795b83b7ce0" } -cuda-async = { git = "https://github.com/tpn/cuda-rust.git", rev = "96caed46c776eb49226137eac849e795b83b7ce0" } +cuda-device = { git = "https://github.com/tpn/cuda-rust.git", rev = "732bec8e2166c33b9b851860363c81baa475203e" } +cuda-host = { git = "https://github.com/tpn/cuda-rust.git", rev = "732bec8e2166c33b9b851860363c81baa475203e", features = ["async"] } +cuda-core = { git = "https://github.com/tpn/cuda-rust.git", rev = "732bec8e2166c33b9b851860363c81baa475203e" } +cuda-async = { git = "https://github.com/tpn/cuda-rust.git", rev = "732bec8e2166c33b9b851860363c81baa475203e" } diff --git a/rust/cuda-coop/examples/reduce/.gitignore b/rust/cuda-coop/examples/reduce/.gitignore new file mode 100644 index 000000000000..74bffa8af4a4 --- /dev/null +++ b/rust/cuda-coop/examples/reduce/.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/reduce/Cargo.lock b/rust/cuda-coop/examples/reduce/Cargo.lock new file mode 100644 index 000000000000..78b9cabaa435 --- /dev/null +++ b/rust/cuda-coop/examples/reduce/Cargo.lock @@ -0,0 +1,1019 @@ +# 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-reduce" +version = "0.0.0" +dependencies = [ + "coop-reduce-helper", + "cuda-async", + "cuda-coop", + "cuda-core", + "cuda-device", + "cuda-host", +] + +[[package]] +name = "coop-reduce-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=732bec8e2166c33b9b851860363c81baa475203e#732bec8e2166c33b9b851860363c81baa475203e" +dependencies = [ + "libnvvm-sys", + "nvjitlink-sys", + "serde", + "sha2", + "thiserror 2.0.21", +] + +[[package]] +name = "cuda-async" +version = "0.5.0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=732bec8e2166c33b9b851860363c81baa475203e#732bec8e2166c33b9b851860363c81baa475203e" +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.5.0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=732bec8e2166c33b9b851860363c81baa475203e#732bec8e2166c33b9b851860363c81baa475203e" +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.5.0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=732bec8e2166c33b9b851860363c81baa475203e#732bec8e2166c33b9b851860363c81baa475203e" +dependencies = [ + "anyhow", + "cuda-bindings", + "cuda-core-derive", + "half", + "oxide-artifacts", +] + +[[package]] +name = "cuda-core-derive" +version = "0.5.0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=732bec8e2166c33b9b851860363c81baa475203e#732bec8e2166c33b9b851860363c81baa475203e" +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=732bec8e2166c33b9b851860363c81baa475203e#732bec8e2166c33b9b851860363c81baa475203e" +dependencies = [ + "cuda-macros", +] + +[[package]] +name = "cuda-host" +version = "0.2.1" +source = "git+https://github.com/tpn/cuda-rust.git?rev=732bec8e2166c33b9b851860363c81baa475203e#732bec8e2166c33b9b851860363c81baa475203e" +dependencies = [ + "cuda-artifact-finalizer", + "cuda-async", + "cuda-core", + "cuda-device", + "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=732bec8e2166c33b9b851860363c81baa475203e#732bec8e2166c33b9b851860363c81baa475203e" +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=732bec8e2166c33b9b851860363c81baa475203e#732bec8e2166c33b9b851860363c81baa475203e" + +[[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=732bec8e2166c33b9b851860363c81baa475203e#732bec8e2166c33b9b851860363c81baa475203e" +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=732bec8e2166c33b9b851860363c81baa475203e#732bec8e2166c33b9b851860363c81baa475203e" +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=732bec8e2166c33b9b851860363c81baa475203e#732bec8e2166c33b9b851860363c81baa475203e" +dependencies = [ + "object", +] + +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +source = "git+https://github.com/tpn/cuda-rust.git?rev=732bec8e2166c33b9b851860363c81baa475203e#732bec8e2166c33b9b851860363c81baa475203e" +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=732bec8e2166c33b9b851860363c81baa475203e#732bec8e2166c33b9b851860363c81baa475203e" + +[[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=732bec8e2166c33b9b851860363c81baa475203e#732bec8e2166c33b9b851860363c81baa475203e" + +[[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/reduce/Cargo.toml b/rust/cuda-coop/examples/reduce/Cargo.toml new file mode 100644 index 000000000000..a5d0545bee73 --- /dev/null +++ b/rust/cuda-coop/examples/reduce/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-reduce" +version = "0.0.0" +edition = "2024" +publish = false + +[workspace] +members = ["helper"] + +[dependencies] +cuda-coop = { path = "../.." } +coop-reduce-helper = { path = "helper" } +cuda-device = { git = "https://github.com/tpn/cuda-rust.git", rev = "732bec8e2166c33b9b851860363c81baa475203e" } +cuda-host = { git = "https://github.com/tpn/cuda-rust.git", rev = "732bec8e2166c33b9b851860363c81baa475203e", features = ["async"] } +cuda-core = { git = "https://github.com/tpn/cuda-rust.git", rev = "732bec8e2166c33b9b851860363c81baa475203e" } +cuda-async = { git = "https://github.com/tpn/cuda-rust.git", rev = "732bec8e2166c33b9b851860363c81baa475203e" } diff --git a/rust/cuda-coop/examples/reduce/helper/Cargo.toml b/rust/cuda-coop/examples/reduce/helper/Cargo.toml new file mode 100644 index 000000000000..7c39f6ba2cdb --- /dev/null +++ b/rust/cuda-coop/examples/reduce/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-reduce-helper" +version = "0.0.0" +edition = "2024" +publish = false + +[dependencies] +cuda-coop = { path = "../../.." } diff --git a/rust/cuda-coop/examples/reduce/helper/src/lib.rs b/rust/cuda-coop/examples/reduce/helper/src/lib.rs new file mode 100644 index 000000000000..ac1edfb73095 --- /dev/null +++ b/rust/cuda-coop/examples/reduce/helper/src/lib.rs @@ -0,0 +1,53 @@ +// Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +//! A reusable block sum that preserves the input element type. +#![no_std] +use cuda_coop as coop; + +/// Sum one block tile, returning its total only at block rank zero. +/// +/// Every block member calls this helper. `zero` fills the missing input tail +/// before reduction, which requires every payload slot initialized. The caller +/// supplies the additive identity for `T`. +pub fn sum_tile( + block: coop::ThreadGroup, + input: &[T], + offset: usize, + valid: usize, + zero: T, +) -> Option { + let mut values = coop::ThreadData::<_, ITEMS_PER_THREAD>::new(); + coop::load( + block, + input, + &mut values, + coop::LoadOptions::default() + .offset(offset) + .valid_items(valid) + .oob_default(zero), + ); + coop::sum(block, &values, Default::default()) +} + +/// Reduce a fully initialized i32 tile with a built-in or callable operator. +/// +/// Every block member calls this helper. The associative operator is retained +/// as its concrete type. Auto-sync scratch completes each reduction before +/// the caller reuses it; the result is valid only at block rank zero. +/// A captured i32 must have the same proved value throughout the block. +pub fn reduce_values, const ITEMS_PER_THREAD: usize>( + block: coop::ThreadGroup, + values: &coop::ThreadData, + op: Op, + scratch: &mut coop::TempStorage<0, 0, true, false>, +) -> Option { + coop::reduce( + block, + values, + coop::ReduceOptions::default() + .binary_op(op) + .algorithm(coop::Raking) + .temp_storage(scratch), + ) +} diff --git a/rust/cuda-coop/examples/reduce/src/main.rs b/rust/cuda-coop/examples/reduce/src/main.rs new file mode 100644 index 000000000000..dc9910eda3a0 --- /dev/null +++ b/rust/cuda-coop/examples/reduce/src/main.rs @@ -0,0 +1,195 @@ +// 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, thread}; +use cuda_host::cuda_module; + +#[cuda_module] +mod kernels { + use super::*; + + fn wrapping_add(left: i32, right: i32) -> i32 { + left.wrapping_add(right) + } + + #[kernel] + #[launch_contract(domain = 1, block = (128, 1, 1), dynamic_shared = inferred)] + pub fn block_callable( + input: &[i32], + mut sums: coop::DisjointSlice, + mut callable_sums: coop::DisjointSlice, + mut maxima: coop::DisjointSlice, + mut biased_sums: coop::DisjointSlice, + bias: i32, + ) { + let block = coop::this_block(); + let offset = block.rank_at::() as usize * 128 * ITEMS_PER_THREAD; + let mut values = coop::ThreadData::<_, ITEMS_PER_THREAD>::new(); + coop::load( + block, + input, + &mut values, + coop::LoadOptions::default().offset(offset), + ); + let before = values; + let mut scratch = coop::TempStorage::new().auto_sync(); + let sum = coop_reduce_helper::reduce_values(block, &values, coop::Sum, &mut scratch); + let callable_sum = + coop_reduce_helper::reduce_values(block, &values, wrapping_add, &mut scratch); + let maximum = coop_reduce_helper::reduce_values( + block, + &values, + |left, right| if left > right { left } else { right }, + &mut scratch, + ); + let biased_sum = coop_reduce_helper::reduce_values( + block, + &values, + move |left, right| wrapping_add(wrapping_add(left, right), bias), + &mut scratch, + ); + for index in 0..ITEMS_PER_THREAD { + assert!(values.get(index) == before.get(index)); + } + // All four collectives finish before any root-only output branch. + if let (Some(sum), Some(callable_sum), Some(maximum), Some(biased_sum)) = + (sum, callable_sum, maximum, biased_sum) + { + if let Some(index) = thread::block_index() + && let Some(slot) = sums.get_mut(index) + { + *slot = sum; + } + if let Some(index) = thread::block_index() + && let Some(slot) = callable_sums.get_mut(index) + { + *slot = callable_sum; + } + if let Some(index) = thread::block_index() + && let Some(slot) = maxima.get_mut(index) + { + *slot = maximum; + } + if let Some(index) = thread::block_index() + && let Some(slot) = biased_sums.get_mut(index) + { + *slot = biased_sum; + } + } + } + + #[kernel] + #[launch_contract(domain = 1, block = (128, 1, 1), dynamic_shared = inferred)] + pub fn block_sum( + input: &[i32], + 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 total = coop_reduce_helper::sum_tile::<_, ITEMS_PER_THREAD>( + block, + input, + offset, + remaining.min(tile), + 0_i32, + ); + // Every thread reduces before this terminal output branch. The witness + // proves ownership of one block's output; get_mut checks its bounds. + if let Some(total) = total + && let Some(index) = thread::block_index() + && let Some(slot) = output.get_mut(index) + { + *slot = total; + } + } +} + +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: the typed declarations and embedded kernels compile together. + let module = unsafe { kernels::load_async(0)? }; + macro_rules! run { + ($items:literal) => {{ + let tile = 128 * $items; + let input: Vec = (0..3 * tile - 7).map(|i| (i % 7) as i32).collect(); + let expected: Vec = input.chunks(tile).map(|chunk| chunk.iter().sum()).collect(); + let plan = module.prepare_block_sum::<$items>(LaunchConfig1D::auto_shared(3, 128))?; + for output_len in [3, 2, 0] { + let source = DeviceBuffer::from_host(&stream, &input)?; + let target = DeviceBuffer::from_host(&stream, &vec![-1_i32; output_len])?; + let (_, target) = module + .block_sum_async_owned::<$items, _, _>(&plan, source, target) + .sync()?; + let output = if target.is_empty() { + stream.synchronize()?; + Vec::new() + } else { + target.to_host_vec(&stream)? + }; + assert_eq!(output, expected[..output_len]); + println!( + "items/thread={}: {output_len} block sums checked (3 blocks launched)", + $items + ); + } + + let input: Vec = (0..3 * tile).map(|i| (i % 29) as i32 - 14).collect(); + let expected_sums: Vec = input.chunks(tile).map(|v| v.iter().sum()).collect(); + let expected_maxima: Vec = input + .chunks(tile) + .map(|v| *v.iter().max().expect("nonempty block tile")) + .collect(); + let plan = + module.prepare_block_callable::<$items>(LaunchConfig1D::auto_shared(3, 128))?; + // The same compiled specialization receives two different capture values. + for bias in [-3_i32, 7] { + let expected_biased: Vec = expected_sums + .iter() + .map(|&sum| sum + (tile as i32 - 1) * bias) + .collect(); + let source = DeviceBuffer::from_host(&stream, &input)?; + let sums = DeviceBuffer::from_host(&stream, &[-1_i32; 3])?; + let callable_sums = DeviceBuffer::from_host(&stream, &[-1_i32; 3])?; + let maxima = DeviceBuffer::from_host(&stream, &[-1_i32; 3])?; + let biased_sums = DeviceBuffer::from_host(&stream, &[-1_i32; 3])?; + let (source, sums, callable_sums, maxima, biased_sums) = module + .block_callable_async_owned::<$items, _, _, _, _, _>( + &plan, + source, + sums, + callable_sums, + maxima, + biased_sums, + bias, + ) + .sync()?; + assert_eq!(source.to_host_vec(&stream)?, input); + assert_eq!(sums.to_host_vec(&stream)?, expected_sums); + assert_eq!(callable_sums.to_host_vec(&stream)?, expected_sums); + assert_eq!(maxima.to_host_vec(&stream)?, expected_maxima); + assert_eq!(biased_sums.to_host_vec(&stream)?, expected_biased); + println!( + "items/thread={}, bias={bias}: 12 built-in/callable results and {} preserved inputs checked", + $items, + input.len(), + ); + } + }}; + } + run!(1); + run!(4); + Ok(()) +} diff --git a/rust/cuda-coop/examples/scan/Cargo.lock b/rust/cuda-coop/examples/scan/Cargo.lock index 565676cddbda..34ea363f2141 100644 --- a/rust/cuda-coop/examples/scan/Cargo.lock +++ b/rust/cuda-coop/examples/scan/Cargo.lock @@ -143,7 +143,7 @@ dependencies = [ [[package]] name = "cuda-artifact-finalizer" version = "0.2.1" -source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=732bec8e2166c33b9b851860363c81baa475203e#732bec8e2166c33b9b851860363c81baa475203e" dependencies = [ "libnvvm-sys", "nvjitlink-sys", @@ -155,7 +155,7 @@ dependencies = [ [[package]] name = "cuda-async" version = "0.5.0" -source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=732bec8e2166c33b9b851860363c81baa475203e#732bec8e2166c33b9b851860363c81baa475203e" dependencies = [ "anyhow", "cuda-bindings", @@ -171,7 +171,7 @@ dependencies = [ [[package]] name = "cuda-bindings" version = "0.5.0" -source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=732bec8e2166c33b9b851860363c81baa475203e#732bec8e2166c33b9b851860363c81baa475203e" dependencies = [ "bindgen", "libloading", @@ -202,7 +202,7 @@ dependencies = [ [[package]] name = "cuda-core" version = "0.5.0" -source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=732bec8e2166c33b9b851860363c81baa475203e#732bec8e2166c33b9b851860363c81baa475203e" dependencies = [ "anyhow", "cuda-bindings", @@ -214,7 +214,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" version = "0.5.0" -source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=732bec8e2166c33b9b851860363c81baa475203e#732bec8e2166c33b9b851860363c81baa475203e" dependencies = [ "proc-macro2", "quote", @@ -224,7 +224,7 @@ dependencies = [ [[package]] name = "cuda-device" version = "0.2.1" -source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=732bec8e2166c33b9b851860363c81baa475203e#732bec8e2166c33b9b851860363c81baa475203e" dependencies = [ "cuda-macros", ] @@ -232,7 +232,7 @@ dependencies = [ [[package]] name = "cuda-host" version = "0.2.1" -source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=732bec8e2166c33b9b851860363c81baa475203e#732bec8e2166c33b9b851860363c81baa475203e" dependencies = [ "cuda-artifact-finalizer", "cuda-async", @@ -250,7 +250,7 @@ dependencies = [ [[package]] name = "cuda-macros" version = "0.2.1" -source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=732bec8e2166c33b9b851860363c81baa475203e#732bec8e2166c33b9b851860363c81baa475203e" dependencies = [ "oxide-tile-export", "proc-macro2", @@ -264,7 +264,7 @@ dependencies = [ [[package]] name = "cuda-target-spec" version = "0.1.0" -source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=732bec8e2166c33b9b851860363c81baa475203e#732bec8e2166c33b9b851860363c81baa475203e" [[package]] name = "dashmap" @@ -519,7 +519,7 @@ dependencies = [ [[package]] name = "libnvvm-sys" version = "0.2.1" -source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=732bec8e2166c33b9b851860363c81baa475203e#732bec8e2166c33b9b851860363c81baa475203e" dependencies = [ "cuda-target-spec", "libloading", @@ -572,7 +572,7 @@ dependencies = [ [[package]] name = "nvjitlink-sys" version = "0.2.1" -source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=732bec8e2166c33b9b851860363c81baa475203e#732bec8e2166c33b9b851860363c81baa475203e" dependencies = [ "libloading", "thiserror 2.0.21", @@ -599,7 +599,7 @@ checksum = "9f7c3e4beb33f85d45ae3e3a1792185706c8e16d043238c593331cc7cd313b50" [[package]] name = "oxide-artifacts" version = "0.2.1" -source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=732bec8e2166c33b9b851860363c81baa475203e#732bec8e2166c33b9b851860363c81baa475203e" dependencies = [ "object", ] @@ -607,7 +607,7 @@ dependencies = [ [[package]] name = "oxide-tile-export" version = "0.2.1" -source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=732bec8e2166c33b9b851860363c81baa475203e#732bec8e2166c33b9b851860363c81baa475203e" dependencies = [ "serde", "serde_json", @@ -655,7 +655,7 @@ dependencies = [ [[package]] name = "ptx-parse" version = "0.1.0" -source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=732bec8e2166c33b9b851860363c81baa475203e#732bec8e2166c33b9b851860363c81baa475203e" [[package]] name = "quote" @@ -707,7 +707,7 @@ checksum = "d6f6ff9a378485b298a5286656da665ba74413d36db0979633275d2e708145d4" [[package]] name = "reserved-oxide-symbols" version = "0.2.1" -source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=732bec8e2166c33b9b851860363c81baa475203e#732bec8e2166c33b9b851860363c81baa475203e" [[package]] name = "rustc-hash" diff --git a/rust/cuda-coop/examples/scan/Cargo.toml b/rust/cuda-coop/examples/scan/Cargo.toml index fe413c7daac5..f63c89373ef8 100644 --- a/rust/cuda-coop/examples/scan/Cargo.toml +++ b/rust/cuda-coop/examples/scan/Cargo.toml @@ -13,7 +13,7 @@ members = ["helper"] [dependencies] cuda-coop = { path = "../.." } coop-scan-helper = { path = "helper" } -cuda-device = { git = "https://github.com/tpn/cuda-rust.git", rev = "96caed46c776eb49226137eac849e795b83b7ce0" } -cuda-host = { git = "https://github.com/tpn/cuda-rust.git", rev = "96caed46c776eb49226137eac849e795b83b7ce0", features = ["async"] } -cuda-core = { git = "https://github.com/tpn/cuda-rust.git", rev = "96caed46c776eb49226137eac849e795b83b7ce0" } -cuda-async = { git = "https://github.com/tpn/cuda-rust.git", rev = "96caed46c776eb49226137eac849e795b83b7ce0" } +cuda-device = { git = "https://github.com/tpn/cuda-rust.git", rev = "732bec8e2166c33b9b851860363c81baa475203e" } +cuda-host = { git = "https://github.com/tpn/cuda-rust.git", rev = "732bec8e2166c33b9b851860363c81baa475203e", features = ["async"] } +cuda-core = { git = "https://github.com/tpn/cuda-rust.git", rev = "732bec8e2166c33b9b851860363c81baa475203e" } +cuda-async = { git = "https://github.com/tpn/cuda-rust.git", rev = "732bec8e2166c33b9b851860363c81baa475203e" } diff --git a/rust/cuda-coop/src/lib.rs b/rust/cuda-coop/src/lib.rs index 3df7818cdd5d..2e86271050c6 100644 --- a/rust/cuda-coop/src/lib.rs +++ b/rust/cuda-coop/src/lib.rs @@ -44,6 +44,11 @@ //! Mode, built-in operator and Block algorithm are selected statically through //! [`ScanOptions`], with exclusive sum as the default. Scan preserves its input //! and uses the same compiler-checked scratch and participation rules. +//! [`reduce`] and [`sum`] return `Some(total)` only at group rank zero, with +//! `None` at other participants. They accept scalars or borrowed payloads on +//! Blocks and Warp groups, including guarded non-exhaustive widths 17–31. +//! [`ReduceOptions`] selects the operator, Block algorithm, scalar count and +//! storage without changing the input. //! Host payload tests need no GPU; //! executing hardware queries requires an Oxide-compiled CUDA kernel. @@ -52,6 +57,7 @@ mod alignment; mod group; mod load_store; +mod reduce; mod scan; mod temp_storage; mod thread_data; @@ -62,6 +68,10 @@ pub use load_store::{ Algorithm, Direct, LoadOptions, LoadSource, LoadStoreGroup, StoreOptions, StoreValue, Striped, Transpose, Vectorize, WarpTranspose, WarpTransposeTimesliced, load, store, }; +pub use reduce::{ + RakingCommutativeOnly, ReduceAlgorithm, ReduceGroup, ReduceOperator, ReduceOptions, + ReduceValue, SumOptions, WarpReductions, reduce, sum, +}; pub use scan::{ BitAnd, BitOr, BitXor, Exclusive, ExclusiveScanOptions, Inclusive, InclusiveScanOptions, Max, Min, Multiplies, Raking, RakingMemoize, ScanAlgorithm, ScanGroup, ScanMode, ScanOperator, diff --git a/rust/cuda-coop/src/reduce.rs b/rust/cuda-coop/src/reduce.rs new file mode 100644 index 000000000000..7733ebbd531a --- /dev/null +++ b/rust/cuda-coop/src/reduce.rs @@ -0,0 +1,534 @@ +// Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +use core::marker::PhantomData; + +use crate::{ + Alignment, BitAnd, BitOr, BitXor, Block, GroupScope, Max, Min, Multiplies, Numeric, Raking, + Sum, SupportedAlignment, TempStorage, TempStorageLike, ThreadData, ThreadGroup, + ThreadsWithinWarp, Warp, +}; +use cuda_device::library; + +mod boundary; + +macro_rules! operators { + ($(($name:ty, $kind:literal)),+ $(,)?) => {$ ( + impl sealed::Operator for $name { const KIND: u8 = $kind; } + impl sealed::Dispatch for $name { + #[inline(always)] + fn reduce< + const ITEMS: usize, const ALIGN: usize, G: ReduceGroup, + A: ReduceAlgorithm, C: sealed::Count, S: TempStorageLike, + const PAYLOAD: bool, + >( + input: &ThreadData, + options: ReduceOptions, + ) -> Option + where Alignment: SupportedAlignment { + boundary::reduce::(input, options) + } + } + )+}; +} +operators!( + (Sum, 0), + (Multiplies, 1), + (Min, 2), + (Max, 3), + (BitAnd, 4), + (BitOr, 5), + (BitXor, 6) +); + +/// A sealed static Reduce operator for numeric input type `T`. +/// +/// [`Sum`], [`Multiplies`], [`Min`] and [`Max`] accept all ten [`Numeric`] types. +/// Bitwise operators accept only the eight integer types. Generic helpers may +/// carry this bound and pass the actual operator value. These built-in +/// operators permit the commutative algorithm. +/// +/// Function items and closures implementing `Fn(i32, i32) -> i32 + Copy` +/// support full-input Block reduction with explicit [`Raking`]. The operator +/// must be associative. A closure may capture one `i32` by value when the compiler +/// proves that its value is identical for every thread in the block at the call. +/// Other captures, function pointers, other input types, other algorithms, Warp +/// groups and partial counts are not supported. +pub trait ReduceOperator: sealed::Operator + sealed::Dispatch {} +impl ReduceOperator for Sum {} +impl ReduceOperator for Multiplies {} +impl ReduceOperator for Min {} +impl ReduceOperator for Max {} +macro_rules! integer_operators { + ($($t:ty),+ $(,)?) => {$ ( + impl ReduceOperator<$t> for BitAnd {} + impl ReduceOperator<$t> for BitOr {} + impl ReduceOperator<$t> for BitXor {} + )+}; +} +integer_operators!(i8, u8, i16, u16, i32, u32, i64, u64); + +impl i32 + Copy> ReduceOperator for F {} +impl i32 + Copy> sealed::Operator for F { + const KIND: u8 = panic!("callable Reduce does not support batched reduction"); +} + +struct CallableBounds(PhantomData<(G, A, C)>); +impl CallableBounds { + const VALID: () = { + assert!(!G::WARP, "callable Reduce requires Block"); + assert!(A::KIND == 2, "callable Reduce requires explicit Raking"); + assert!(C::KIND == 0, "callable Reduce requires full input"); + }; +} +impl i32 + Copy> sealed::Dispatch for F { + #[inline(always)] + fn reduce< + const ITEMS: usize, + const ALIGN: usize, + G: ReduceGroup, + A: ReduceAlgorithm, + C: sealed::Count, + S: TempStorageLike, + const PAYLOAD: bool, + >( + input: &ThreadData, + options: ReduceOptions, + ) -> Option + where + Alignment: SupportedAlignment, + { + let () = CallableBounds::::VALID; + boundary::callable(input, options.binary_op, options.storage) + } +} + +/// CUB's raking Block reduction specialized for commutative operators. +#[derive(Clone, Copy, Debug, Default)] +pub struct RakingCommutativeOnly; +/// CUB's warp-based Block reduction; selected when the algorithm is omitted. +#[derive(Clone, Copy, Debug, Default)] +pub struct WarpReductions; + +/// A sealed Block reduction algorithm or omitted selection `()`. +/// +/// Block resolves omission to [`WarpReductions`]. Explicit [`Raking`], +/// [`RakingCommutativeOnly`] and [`WarpReductions`] are Block-only. Warp accepts +/// only the omitted selection, using its own CUB Warp reduction. +pub trait ReduceAlgorithm: sealed::Algorithm {} +macro_rules! algorithms { + ($(($name:ty, $kind:literal)),+ $(,)?) => {$ ( + impl sealed::Algorithm for $name { const KIND: u8 = $kind; } + impl ReduceAlgorithm for $name {} + )+}; +} +algorithms!( + ((), 0), + (RakingCommutativeOnly, 1), + (Raking, 2), + (WarpReductions, 3) +); + +/// A sealed group supporting a Reduce algorithm and storage choice. +/// +/// Complete Blocks support scalar or payload input with implicit or explicit +/// storage. Physical [`Warp`] and logical widths 1/2/4/8/16/32 support both +/// input forms with omitted algorithm and implicit storage only. Non-exhaustive +/// logical widths 17 through 31 support the low-width prefix of each physical +/// warp, inside a genuine [`ThreadGroup::is_member`] guard. Warp operations +/// require complete physical warps in the enclosing block. +/// +/// ```compile_fail,E0308 +/// use cuda_coop as coop; +/// fn invalid(value: i32) { +/// coop::sum(coop::this_warp(), value, +/// coop::ReduceOptions::default().algorithm(coop::Raking)); +/// } +/// ``` +pub trait ReduceGroup: GroupScope + sealed::Group {} +impl ReduceGroup for Block {} +impl ReduceGroup<(), ()> for Warp {} +impl sealed::Group for Block { + const WARP: bool = false; + #[inline(always)] + fn count() -> usize { + crate::this_block().count() as usize + } + #[inline(always)] + unsafe fn call &'static str + Copy, S, U: Copy, P: Copy, const I: usize>( + registration: R, + storage: S, + uniform: U, + payload: P, + ) { + // SAFETY: the private boundary supplies live typed payloads and the exact recipe. + unsafe { + library::checked_block_payload_call_v1::( + registration, + storage, + uniform, + payload, + ) + } + } +} +macro_rules! warp_group { + ($scope:ty, $width:expr $(, $e:ident)?) => { + impl<$(const $e: bool)?> sealed::Group for $scope { + const WARP: bool = true; + #[inline(always)] fn count() -> usize { $width } + #[inline(always)] + unsafe fn call &'static str + Copy, S, U: Copy, P: Copy, const I: usize>( + registration: R, _storage: S, uniform: U, payload: P, + ) { + // SAFETY: ReduceGroup permits implicit storage only; the compiler checks membership and width. + unsafe { library::checked_warp_payload_call_v1::(registration, uniform, payload) } + } + } + }; +} +warp_group!(Warp, 32); +macro_rules! logical_groups { + ($($width:literal),+) => {$ ( + impl ReduceGroup<(), ()> for ThreadsWithinWarp<$width, E> {} + warp_group!(ThreadsWithinWarp<$width, E>, $width, E); + )+}; +} +logical_groups!(1, 2, 4, 8, 16, 32); +macro_rules! prefix_groups { + ($($width:literal),+) => {$ ( + impl ReduceGroup<(), ()> for ThreadsWithinWarp<$width, false> {} + warp_group!(ThreadsWithinWarp<$width, false>, $width); + )+}; +} +prefix_groups!(17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31); + +/// Static operator/algorithm and runtime count/storage for [`reduce`] and [`sum`]. +/// +/// `T` is inferred from the input. Defaults select [`Sum`], omitted algorithm +/// (Block [`WarpReductions`]), every input and managed storage. Only this +/// canonical specialization implements Default, so `Default::default()` is +/// unambiguous. Builders consume options and preserve their other fields. +/// +/// A valid count is supported only for scalar input, including when a payload +/// happens to contain one item. Warp rejects explicit algorithm and storage. +/// +/// ```no_run +/// use cuda_coop as coop; +/// fn total, const ITEMS_PER_THREAD: usize>( +/// values: &coop::ThreadData, op: Op, +/// ) -> Option { +/// coop::reduce(coop::this_block(), values, +/// coop::ReduceOptions::default().binary_op(op)) +/// } +/// ``` +/// ```compile_fail,E0308 +/// use cuda_coop as coop; +/// fn invalid(values: &coop::ThreadData) { +/// coop::sum(coop::this_block(), values, +/// coop::ReduceOptions::default().valid_items(1)); +/// } +/// ``` +/// ```compile_fail,E0277 +/// use cuda_coop as coop; +/// fn invalid(value: f32) { +/// coop::reduce(coop::this_block(), value, +/// coop::ReduceOptions::default().binary_op(coop::BitAnd)); +/// } +/// ``` +pub struct ReduceOptions { + binary_op: Op, + valid_items: Count, + storage: S, + selection: PhantomData<(T, A)>, +} +impl Default for ReduceOptions { + fn default() -> Self { + Self { + binary_op: Sum, + valid_items: (), + storage: (), + selection: PhantomData, + } + } +} +impl ReduceOptions { + /// Select an operator, retaining its actual generic value. + /// + /// `i32` callables require Block, explicit [`Raking`] and full input. + /// See [`ReduceOperator`] for supported captures and their uniformity rule. + pub fn binary_op>(self, op: Op2) -> ReduceOptions { + ReduceOptions { + binary_op: op, + valid_items: self.valid_items, + storage: self.storage, + selection: PhantomData, + } + } + /// Select a Block algorithm. Warp accepts only the omitted selection `()`. + pub fn algorithm( + self, + _algorithm: A2, + ) -> ReduceOptions { + ReduceOptions { + binary_op: self.binary_op, + valid_items: self.valid_items, + storage: self.storage, + selection: PhantomData, + } + } + /// Reduce the first `count` scalar inputs in group rank order. + /// + /// Requires `1 <= count <= group.count()` and a group-uniform count. All + /// members still participate; excluded scalar inputs do not contribute. + /// Payload inputs reject this option, even for a one-item payload. + pub fn valid_items(self, count: usize) -> ReduceOptions { + ReduceOptions { + binary_op: self.binary_op, + valid_items: count, + storage: self.storage, + selection: PhantomData, + } + } + /// Borrow explicit Block scratch with its existing reuse/completion policy. + /// + /// Warp rejects this choice. Consuming the options permits ordinary mutable + /// reborrows for later operations, without requiring storage to be Copy. + pub fn temp_storage( + self, + storage: &mut TempStorage, + ) -> ReduceOptions> { + ReduceOptions { + binary_op: self.binary_op, + valid_items: self.valid_items, + storage, + selection: PhantomData, + } + } +} + +/// Algorithm, scalar count and storage options for [`sum`]; fixes the operator to [`Sum`]. +pub type SumOptions = ReduceOptions; + +/// Reduce numeric values, returning `Some(total)` only at group rank zero. +/// +/// Scalars and shared/mutable payload borrows all return `Option`. Every +/// other participating thread gets `None`; no undefined numeric slot is read, +/// broadcast or replaced by an implicit identity. Input values and validity +/// are preserved. Block order is x-fastest thread rank followed by each +/// thread's item order. Warp uses the corresponding order within each group. +/// +/// Defaults select sum, all inputs and Block [`WarpReductions`]. Scalar-only +/// [`valid_items`](ReduceOptions::valid_items) restricts contributors to a +/// nonempty prefix; every member must still reach the operation with the same +/// count. Payload reduction requires every slot initialized. A tail must be +/// filled explicitly before reduction, for example by a default-filled Load. +/// CUB may regroup associative arithmetic; overflow and floating behavior +/// follow the input dtype and primitive, not a sequential host-order guarantee. +/// +/// All members must participate. [`ReduceGroup`] describes supported groups; +/// prefix widths 17–31 require a genuine membership guard. A result's `Some` +/// branch grants no collective participation or global-memory ownership. +/// For a per-block output, obtain a checked [`cuda_device::thread::BlockIndex`] +/// witness inside the terminal result branch and use a [`crate::DisjointSlice`]. +/// +/// Managed scratch completes before dirty reuse, without a trailing completion +/// promise. Explicit manual scratch needs completion before conflicting reuse; +/// auto-sync storage adds a trailing Block barrier. The operation alone is not +/// completion of prior global-memory accesses. Execution requires pinned Oxide. +/// +/// # Panics +/// An uninitialized required input or a scalar count outside `1..=group.count()` +/// traps before reduction. Invalid participation, options, geometry or scratch +/// reuse fail compilation. +/// +/// ```no_run +/// use cuda_coop as coop; +/// fn total( +/// values: &coop::ThreadData, +/// ) -> Option { +/// coop::reduce(coop::this_block(), values, Default::default()) +/// } +/// ``` +#[inline(always)] +pub fn reduce< + T: Numeric, + V: ReduceValue, + G: ReduceGroup, + Op: ReduceOperator, + A: ReduceAlgorithm, + Count, + S: TempStorageLike, +>( + group: ThreadGroup, + value: V, + options: ReduceOptions, +) -> Option { + value.reduce(group, options) +} + +/// Sum numeric values; root-only validity, input preservation and completion follow [`reduce`]. +/// +/// The operator is fixed to [`Sum`]. Options select Block algorithm/storage and +/// an optional scalar count. Default options infer `T` and payload extent. +#[inline(always)] +pub fn sum< + T: Numeric, + V: ReduceValue, + G: ReduceGroup, + A: ReduceAlgorithm, + Count, + S: TempStorageLike, +>( + group: ThreadGroup, + value: V, + options: SumOptions, +) -> Option { + reduce(group, value, options) +} + +/// A sealed scalar or borrowed payload input with an inferred root-only result. +/// +/// Numeric scalars permit omitted or `usize` count. Shared/mutable +/// `ThreadData` borrows permit only omitted count. Both input +/// forms support every [`ReduceGroup`], preserve input and return `Option`. +/// Owned-payload-consuming forms are not provided. +pub trait ReduceValue: sealed::Value { + #[doc(hidden)] + fn reduce, A: ReduceAlgorithm, S: TempStorageLike>( + self, + group: ThreadGroup, + options: ReduceOptions, + ) -> Option + where + G: ReduceGroup; +} +impl sealed::Value for T {} +impl ReduceValue for T { + #[inline(always)] + fn reduce, A: ReduceAlgorithm, S: TempStorageLike>( + self, + _group: ThreadGroup, + options: ReduceOptions, + ) -> Option + where + G: ReduceGroup, + { + let mut input = ThreadData::<_, 1>::new(); + input.write(0, self); + Op::reduce::<1, 1, G, A, C, S, false>(&input, options) + } +} +impl sealed::Value + for &ThreadData +where + Alignment: SupportedAlignment, +{ +} +impl ReduceValue + for &ThreadData +where + Alignment: SupportedAlignment, +{ + #[inline(always)] + fn reduce, A: ReduceAlgorithm, S: TempStorageLike>( + self, + _group: ThreadGroup, + options: ReduceOptions, + ) -> Option + where + G: ReduceGroup, + { + Op::reduce::(self, options) + } +} +impl sealed::Value + for &mut ThreadData +where + Alignment: SupportedAlignment, +{ +} +impl ReduceValue + for &mut ThreadData +where + Alignment: SupportedAlignment, +{ + #[inline(always)] + fn reduce, A: ReduceAlgorithm, S: TempStorageLike>( + self, + group: ThreadGroup, + options: ReduceOptions, + ) -> Option + where + G: ReduceGroup, + { + (&*self).reduce(group, options) + } +} + +mod sealed { + use super::*; + + pub trait Dispatch: Sized { + fn reduce< + const ITEMS: usize, + const ALIGN: usize, + G: ReduceGroup, + A: ReduceAlgorithm, + C: Count, + S: TempStorageLike, + const PAYLOAD: bool, + >( + input: &ThreadData, + options: ReduceOptions, + ) -> Option + where + Alignment: SupportedAlignment; + } + pub trait Operator { + const KIND: u8; + } + pub trait Algorithm { + const KIND: u8; + } + pub trait Value {} + pub trait Count { + type Uniform: Copy; + const KIND: u8; + fn controls(self) -> Self::Uniform; + } + impl Count for () { + type Uniform = (); + const KIND: u8 = 0; + #[inline(always)] + fn controls(self) {} + } + impl Count for usize { + type Uniform = (i32,); + const KIND: u8 = 1; + #[inline(always)] + fn controls(self) -> (i32,) { + assert!(self > 0); + assert!(self <= G::count()); + // CUDA group sizes are bounded by the maximum threads in a block. + (self as i32,) + } + } + pub trait Group { + const WARP: bool; + fn count() -> usize; + unsafe fn call< + T: Copy, + R: Fn() -> &'static str + Copy, + S, + U: Copy, + P: Copy, + const I: usize, + >( + registration: R, + storage: S, + uniform: U, + payload: P, + ); + } +} diff --git a/rust/cuda-coop/src/reduce/boundary.rs b/rust/cuda-coop/src/reduce/boundary.rs new file mode 100644 index 000000000000..6c52b78e428a --- /dev/null +++ b/rust/cuda-coop/src/reduce/boundary.rs @@ -0,0 +1,473 @@ +// Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +use super::{ReduceAlgorithm, ReduceGroup, ReduceOptions, sealed::Count}; +use crate::{Alignment, Numeric, SupportedAlignment, TempStorageLike, ThreadData}; +use cuda_device::library; + +// Each static recipe keeps a concrete function-item registration identity. +macro_rules! registrations { + ($(($name:ident, $variable:literal)),+ $(,)?) => {$ ( + fn $name() -> &'static str { library::registration!($variable) } + )+}; +} +registrations!( + ( + raking_commutative_only_sum_scalar, + "COOP_REDUCE_RAKING_COMMUTATIVE_ONLY_SUM_SCALAR" + ), + (raking_sum_scalar, "COOP_REDUCE_RAKING_SUM_SCALAR"), + ( + warp_reductions_sum_scalar, + "COOP_REDUCE_WARP_REDUCTIONS_SUM_SCALAR" + ), + (warp_sum_scalar, "COOP_REDUCE_WARP_SUM_SCALAR"), + ( + raking_commutative_only_multiplies_scalar, + "COOP_REDUCE_RAKING_COMMUTATIVE_ONLY_MULTIPLIES_SCALAR" + ), + ( + raking_multiplies_scalar, + "COOP_REDUCE_RAKING_MULTIPLIES_SCALAR" + ), + ( + warp_reductions_multiplies_scalar, + "COOP_REDUCE_WARP_REDUCTIONS_MULTIPLIES_SCALAR" + ), + (warp_multiplies_scalar, "COOP_REDUCE_WARP_MULTIPLIES_SCALAR"), + ( + raking_commutative_only_min_scalar, + "COOP_REDUCE_RAKING_COMMUTATIVE_ONLY_MIN_SCALAR" + ), + (raking_min_scalar, "COOP_REDUCE_RAKING_MIN_SCALAR"), + ( + warp_reductions_min_scalar, + "COOP_REDUCE_WARP_REDUCTIONS_MIN_SCALAR" + ), + (warp_min_scalar, "COOP_REDUCE_WARP_MIN_SCALAR"), + ( + raking_commutative_only_max_scalar, + "COOP_REDUCE_RAKING_COMMUTATIVE_ONLY_MAX_SCALAR" + ), + (raking_max_scalar, "COOP_REDUCE_RAKING_MAX_SCALAR"), + ( + warp_reductions_max_scalar, + "COOP_REDUCE_WARP_REDUCTIONS_MAX_SCALAR" + ), + (warp_max_scalar, "COOP_REDUCE_WARP_MAX_SCALAR"), + ( + raking_commutative_only_bit_and_scalar, + "COOP_REDUCE_RAKING_COMMUTATIVE_ONLY_BIT_AND_SCALAR" + ), + (raking_bit_and_scalar, "COOP_REDUCE_RAKING_BIT_AND_SCALAR"), + ( + warp_reductions_bit_and_scalar, + "COOP_REDUCE_WARP_REDUCTIONS_BIT_AND_SCALAR" + ), + (warp_bit_and_scalar, "COOP_REDUCE_WARP_BIT_AND_SCALAR"), + ( + raking_commutative_only_bit_or_scalar, + "COOP_REDUCE_RAKING_COMMUTATIVE_ONLY_BIT_OR_SCALAR" + ), + (raking_bit_or_scalar, "COOP_REDUCE_RAKING_BIT_OR_SCALAR"), + ( + warp_reductions_bit_or_scalar, + "COOP_REDUCE_WARP_REDUCTIONS_BIT_OR_SCALAR" + ), + (warp_bit_or_scalar, "COOP_REDUCE_WARP_BIT_OR_SCALAR"), + ( + raking_commutative_only_bit_xor_scalar, + "COOP_REDUCE_RAKING_COMMUTATIVE_ONLY_BIT_XOR_SCALAR" + ), + (raking_bit_xor_scalar, "COOP_REDUCE_RAKING_BIT_XOR_SCALAR"), + ( + warp_reductions_bit_xor_scalar, + "COOP_REDUCE_WARP_REDUCTIONS_BIT_XOR_SCALAR" + ), + (warp_bit_xor_scalar, "COOP_REDUCE_WARP_BIT_XOR_SCALAR"), + ( + raking_commutative_only_sum_prefix, + "COOP_REDUCE_RAKING_COMMUTATIVE_ONLY_SUM_PREFIX" + ), + (raking_sum_prefix, "COOP_REDUCE_RAKING_SUM_PREFIX"), + ( + warp_reductions_sum_prefix, + "COOP_REDUCE_WARP_REDUCTIONS_SUM_PREFIX" + ), + (warp_sum_prefix, "COOP_REDUCE_WARP_SUM_PREFIX"), + ( + raking_commutative_only_multiplies_prefix, + "COOP_REDUCE_RAKING_COMMUTATIVE_ONLY_MULTIPLIES_PREFIX" + ), + ( + raking_multiplies_prefix, + "COOP_REDUCE_RAKING_MULTIPLIES_PREFIX" + ), + ( + warp_reductions_multiplies_prefix, + "COOP_REDUCE_WARP_REDUCTIONS_MULTIPLIES_PREFIX" + ), + (warp_multiplies_prefix, "COOP_REDUCE_WARP_MULTIPLIES_PREFIX"), + ( + raking_commutative_only_min_prefix, + "COOP_REDUCE_RAKING_COMMUTATIVE_ONLY_MIN_PREFIX" + ), + (raking_min_prefix, "COOP_REDUCE_RAKING_MIN_PREFIX"), + ( + warp_reductions_min_prefix, + "COOP_REDUCE_WARP_REDUCTIONS_MIN_PREFIX" + ), + (warp_min_prefix, "COOP_REDUCE_WARP_MIN_PREFIX"), + ( + raking_commutative_only_max_prefix, + "COOP_REDUCE_RAKING_COMMUTATIVE_ONLY_MAX_PREFIX" + ), + (raking_max_prefix, "COOP_REDUCE_RAKING_MAX_PREFIX"), + ( + warp_reductions_max_prefix, + "COOP_REDUCE_WARP_REDUCTIONS_MAX_PREFIX" + ), + (warp_max_prefix, "COOP_REDUCE_WARP_MAX_PREFIX"), + ( + raking_commutative_only_bit_and_prefix, + "COOP_REDUCE_RAKING_COMMUTATIVE_ONLY_BIT_AND_PREFIX" + ), + (raking_bit_and_prefix, "COOP_REDUCE_RAKING_BIT_AND_PREFIX"), + ( + warp_reductions_bit_and_prefix, + "COOP_REDUCE_WARP_REDUCTIONS_BIT_AND_PREFIX" + ), + (warp_bit_and_prefix, "COOP_REDUCE_WARP_BIT_AND_PREFIX"), + ( + raking_commutative_only_bit_or_prefix, + "COOP_REDUCE_RAKING_COMMUTATIVE_ONLY_BIT_OR_PREFIX" + ), + (raking_bit_or_prefix, "COOP_REDUCE_RAKING_BIT_OR_PREFIX"), + ( + warp_reductions_bit_or_prefix, + "COOP_REDUCE_WARP_REDUCTIONS_BIT_OR_PREFIX" + ), + (warp_bit_or_prefix, "COOP_REDUCE_WARP_BIT_OR_PREFIX"), + ( + raking_commutative_only_bit_xor_prefix, + "COOP_REDUCE_RAKING_COMMUTATIVE_ONLY_BIT_XOR_PREFIX" + ), + (raking_bit_xor_prefix, "COOP_REDUCE_RAKING_BIT_XOR_PREFIX"), + ( + warp_reductions_bit_xor_prefix, + "COOP_REDUCE_WARP_REDUCTIONS_BIT_XOR_PREFIX" + ), + (warp_bit_xor_prefix, "COOP_REDUCE_WARP_BIT_XOR_PREFIX"), + ( + raking_commutative_only_sum_payload, + "COOP_REDUCE_RAKING_COMMUTATIVE_ONLY_SUM_PAYLOAD" + ), + (raking_sum_payload, "COOP_REDUCE_RAKING_SUM_PAYLOAD"), + ( + warp_reductions_sum_payload, + "COOP_REDUCE_WARP_REDUCTIONS_SUM_PAYLOAD" + ), + (warp_sum_payload, "COOP_REDUCE_WARP_SUM_PAYLOAD"), + ( + raking_commutative_only_multiplies_payload, + "COOP_REDUCE_RAKING_COMMUTATIVE_ONLY_MULTIPLIES_PAYLOAD" + ), + ( + raking_multiplies_payload, + "COOP_REDUCE_RAKING_MULTIPLIES_PAYLOAD" + ), + ( + warp_reductions_multiplies_payload, + "COOP_REDUCE_WARP_REDUCTIONS_MULTIPLIES_PAYLOAD" + ), + ( + warp_multiplies_payload, + "COOP_REDUCE_WARP_MULTIPLIES_PAYLOAD" + ), + ( + raking_commutative_only_min_payload, + "COOP_REDUCE_RAKING_COMMUTATIVE_ONLY_MIN_PAYLOAD" + ), + (raking_min_payload, "COOP_REDUCE_RAKING_MIN_PAYLOAD"), + ( + warp_reductions_min_payload, + "COOP_REDUCE_WARP_REDUCTIONS_MIN_PAYLOAD" + ), + (warp_min_payload, "COOP_REDUCE_WARP_MIN_PAYLOAD"), + ( + raking_commutative_only_max_payload, + "COOP_REDUCE_RAKING_COMMUTATIVE_ONLY_MAX_PAYLOAD" + ), + (raking_max_payload, "COOP_REDUCE_RAKING_MAX_PAYLOAD"), + ( + warp_reductions_max_payload, + "COOP_REDUCE_WARP_REDUCTIONS_MAX_PAYLOAD" + ), + (warp_max_payload, "COOP_REDUCE_WARP_MAX_PAYLOAD"), + ( + raking_commutative_only_bit_and_payload, + "COOP_REDUCE_RAKING_COMMUTATIVE_ONLY_BIT_AND_PAYLOAD" + ), + (raking_bit_and_payload, "COOP_REDUCE_RAKING_BIT_AND_PAYLOAD"), + ( + warp_reductions_bit_and_payload, + "COOP_REDUCE_WARP_REDUCTIONS_BIT_AND_PAYLOAD" + ), + (warp_bit_and_payload, "COOP_REDUCE_WARP_BIT_AND_PAYLOAD"), + ( + raking_commutative_only_bit_or_payload, + "COOP_REDUCE_RAKING_COMMUTATIVE_ONLY_BIT_OR_PAYLOAD" + ), + (raking_bit_or_payload, "COOP_REDUCE_RAKING_BIT_OR_PAYLOAD"), + ( + warp_reductions_bit_or_payload, + "COOP_REDUCE_WARP_REDUCTIONS_BIT_OR_PAYLOAD" + ), + (warp_bit_or_payload, "COOP_REDUCE_WARP_BIT_OR_PAYLOAD"), + ( + raking_commutative_only_bit_xor_payload, + "COOP_REDUCE_RAKING_COMMUTATIVE_ONLY_BIT_XOR_PAYLOAD" + ), + (raking_bit_xor_payload, "COOP_REDUCE_RAKING_BIT_XOR_PAYLOAD"), + ( + warp_reductions_bit_xor_payload, + "COOP_REDUCE_WARP_REDUCTIONS_BIT_XOR_PAYLOAD" + ), + (warp_bit_xor_payload, "COOP_REDUCE_WARP_BIT_XOR_PAYLOAD"), +); + +#[inline(always)] +pub(super) fn reduce< + T: Numeric, + const ITEMS: usize, + const ALIGN: usize, + G: ReduceGroup, + Op: super::sealed::Operator, + A: ReduceAlgorithm, + C: Count, + S: TempStorageLike, + const PAYLOAD: bool, +>( + input: &ThreadData, + options: ReduceOptions, +) -> Option +where + Alignment: SupportedAlignment, +{ + let mut output = ThreadData::::new(); + let (input, input_bits) = input.block_io_read_parts(); + let (output_pointer, output_bits) = output.block_io_parts(); + let payload = (input, output_pointer, input_bits, output_bits); + let uniform = C::controls::(options.valid_items); + let storage = options.storage.into_raw(); + let _operator = options.binary_op; + macro_rules! call { + ($commutative:ident, $raking:ident, $warps:ident, $warp:ident) => { + if G::WARP { + G::call::($warp, storage, uniform, payload); + } else { + match A::KIND { + 0 | 3 => G::call::($warps, storage, uniform, payload), + 1 => G::call::($commutative, storage, uniform, payload), + 2 => G::call::($raking, storage, uniform, payload), + _ => panic!(), + } + } + }; + } + // SAFETY: the const input and fresh output are distinct, aligned, live + // lane-local payloads. Sealed selections establish the exact native ABI. + // The recipe checks required input bits before numeric reads, preserves + // input, clears every output bit, and writes a value then its bit only at + // group rank zero. The pointers cannot escape the synchronous call. + // The checked marker enforces participation, scalar-control uniformity, + // layout and scratch reuse; no global-memory completion is fabricated. + unsafe { + match const { + if PAYLOAD { + 14 + Op::KIND + } else { + 7 * C::KIND + Op::KIND + } + } { + 0 => call!( + raking_commutative_only_sum_scalar, + raking_sum_scalar, + warp_reductions_sum_scalar, + warp_sum_scalar + ), + 1 => call!( + raking_commutative_only_multiplies_scalar, + raking_multiplies_scalar, + warp_reductions_multiplies_scalar, + warp_multiplies_scalar + ), + 2 => call!( + raking_commutative_only_min_scalar, + raking_min_scalar, + warp_reductions_min_scalar, + warp_min_scalar + ), + 3 => call!( + raking_commutative_only_max_scalar, + raking_max_scalar, + warp_reductions_max_scalar, + warp_max_scalar + ), + 4 => call!( + raking_commutative_only_bit_and_scalar, + raking_bit_and_scalar, + warp_reductions_bit_and_scalar, + warp_bit_and_scalar + ), + 5 => call!( + raking_commutative_only_bit_or_scalar, + raking_bit_or_scalar, + warp_reductions_bit_or_scalar, + warp_bit_or_scalar + ), + 6 => call!( + raking_commutative_only_bit_xor_scalar, + raking_bit_xor_scalar, + warp_reductions_bit_xor_scalar, + warp_bit_xor_scalar + ), + 7 => call!( + raking_commutative_only_sum_prefix, + raking_sum_prefix, + warp_reductions_sum_prefix, + warp_sum_prefix + ), + 8 => call!( + raking_commutative_only_multiplies_prefix, + raking_multiplies_prefix, + warp_reductions_multiplies_prefix, + warp_multiplies_prefix + ), + 9 => call!( + raking_commutative_only_min_prefix, + raking_min_prefix, + warp_reductions_min_prefix, + warp_min_prefix + ), + 10 => call!( + raking_commutative_only_max_prefix, + raking_max_prefix, + warp_reductions_max_prefix, + warp_max_prefix + ), + 11 => call!( + raking_commutative_only_bit_and_prefix, + raking_bit_and_prefix, + warp_reductions_bit_and_prefix, + warp_bit_and_prefix + ), + 12 => call!( + raking_commutative_only_bit_or_prefix, + raking_bit_or_prefix, + warp_reductions_bit_or_prefix, + warp_bit_or_prefix + ), + 13 => call!( + raking_commutative_only_bit_xor_prefix, + raking_bit_xor_prefix, + warp_reductions_bit_xor_prefix, + warp_bit_xor_prefix + ), + 14 => call!( + raking_commutative_only_sum_payload, + raking_sum_payload, + warp_reductions_sum_payload, + warp_sum_payload + ), + 15 => call!( + raking_commutative_only_multiplies_payload, + raking_multiplies_payload, + warp_reductions_multiplies_payload, + warp_multiplies_payload + ), + 16 => call!( + raking_commutative_only_min_payload, + raking_min_payload, + warp_reductions_min_payload, + warp_min_payload + ), + 17 => call!( + raking_commutative_only_max_payload, + raking_max_payload, + warp_reductions_max_payload, + warp_max_payload + ), + 18 => call!( + raking_commutative_only_bit_and_payload, + raking_bit_and_payload, + warp_reductions_bit_and_payload, + warp_bit_and_payload + ), + 19 => call!( + raking_commutative_only_bit_or_payload, + raking_bit_or_payload, + warp_reductions_bit_or_payload, + warp_bit_or_payload + ), + 20 => call!( + raking_commutative_only_bit_xor_payload, + raking_bit_xor_payload, + warp_reductions_bit_xor_payload, + warp_bit_xor_payload + ), + _ => panic!(), + } + } + // get checks root-only validity before reading the numerical output. + output.get(0) +} + +fn raking_callable_payload() -> &'static str { + library::registration!("COOP_REDUCE_RAKING_CALLABLE_PAYLOAD") +} + +unsafe extern "C" fn binary_thunk i32>( + environment: *const F, + left: i32, + right: i32, +) -> i32 { + // SAFETY: the synchronous checked call keeps its shared environment alive. + unsafe { (&*environment)(left, right) } +} + +#[inline(always)] +pub(super) fn callable< + const ITEMS: usize, + const ALIGN: usize, + F: Fn(i32, i32) -> i32 + Copy, + S: TempStorageLike, +>( + input: &ThreadData, + operator: F, + storage: S, +) -> Option +where + Alignment: SupportedAlignment, +{ + let mut output = ThreadData::::new(); + let (input, input_bits) = input.block_io_read_parts(); + let (output_pointer, output_bits) = output.block_io_parts(); + let payload = (input, output_pointer, input_bits, output_bits); + let storage = storage.into_raw(); + // SAFETY: input and fresh output are distinct live payloads. The recipe + // checks every input bit, preserves input and writes only the root output. + // The compiler checks Block participation, scratch reuse and the genuine + // callable and any captured value. The environment borrow lives through this call. + unsafe { + library::checked_block_payload_callback_call_v1::( + raking_callable_payload, + storage, + (), + payload, + binary_thunk::, + &operator, + ) + } + output.get(0) +} diff --git a/rust/cuda-coop/src/scan.rs b/rust/cuda-coop/src/scan.rs index 78e2db1111f9..22742cff9b1d 100644 --- a/rust/cuda-coop/src/scan.rs +++ b/rust/cuda-coop/src/scan.rs @@ -31,14 +31,14 @@ operators!( ( Sum, 0, - "Addition in the input numeric type; the default Scan operator." + "Addition in the input numeric type; the default Scan and Reduce operator." ), (Multiplies, 1, "Multiplication in the input numeric type."), (Min, 2, "CUB minimum in the input numeric type."), (Max, 3, "CUB maximum in the input numeric type."), - (BitAnd, 4, "Bitwise AND for integer Scan inputs."), - (BitOr, 5, "Bitwise OR for integer Scan inputs."), - (BitXor, 6, "Bitwise XOR for integer Scan inputs."), + (BitAnd, 4, "Bitwise AND for integer inputs."), + (BitOr, 5, "Bitwise OR for integer inputs."), + (BitXor, 6, "Bitwise XOR for integer inputs."), ); /// A sealed static Scan operator for numeric input type `T`. @@ -73,7 +73,7 @@ algorithms!( ( Raking, 1, - "CUB's raking Block Scan; selected when the algorithm is omitted." + "CUB's raking Block algorithm; default for Scan and an explicit choice for Reduce." ), (RakingMemoize, 2, "CUB's memoized raking Block Scan."), ( diff --git a/rust/cuda-coop/tests/check_load_store.py b/rust/cuda-coop/tests/check_load_store.py index b205e730a4e6..60c3d0d6643f 100644 --- a/rust/cuda-coop/tests/check_load_store.py +++ b/rust/cuda-coop/tests/check_load_store.py @@ -1,7 +1,7 @@ # Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. # SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception -"""Compile public Load/Store and Scan participation/resource-epoch contrasts. +"""Compile public collective participation, result 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 @@ -932,6 +932,141 @@ ), ) +# A cold hint preserves participation and outstanding memory-access facts. +cases.update( + cold_path_before_collective=("core::hint::cold_path();" + load + store, None), + cold_path_divergent_collective=( + "if thread::threadIdx_x()&1==0 {core::hint::cold_path();" + load + "}", + "lane-varying", + ), + cold_path_no_completion=( + load + store + "core::hint::cold_path();" + reload, + "requires a barrier covering both accesses", + ), + cold_path_user_function=( + "fn cold_path() {thread::sync_threads();}" + "if thread::threadIdx_x()&1==0 {cold_path();}" + load, + "block barrier is reached through lane-varying control flow", + ), +) + +# Reduce returns a guarded scalar only at the group root. That local result +# does not establish collective participation or complete global accesses. +reduce_call = "let _result=coop::sum(block,&values,Default::default());" +reduce_manual = "let _result=coop::sum(block,&values,coop::ReduceOptions::default().temp_storage(&mut scratch));" +reduce_prefix = "let group=coop::this_warp().group_by_non_exhaustive::<17>();" +cases.update( + reduce_defaults=( + load + + reduce_call + + "let _result:Option=coop::reduce(block,&mut values,Default::default());" + + "let _result=coop::sum(block,1_i32,coop::ReduceOptions::default().valid_items(valid));" + + "let _result=coop::sum(coop::this_warp(),&values,Default::default());" + + "let _result=coop::sum(coop::this_warp().group_by::<8>(),&values,Default::default());", + None, + ), + reduce_generic_operator=( + load + "fn aggregate,const N:usize>(" + "group:coop::ThreadGroup,input:&coop::ThreadData,op:Op)" + "->Option {coop::reduce(group,input,coop::ReduceOptions::default().binary_op(op))}" + "let _result=aggregate(block,&values,coop::Min);", + None, + ), + reduce_payload_count=( + load + + "let _result=coop::sum(block,&values,coop::ReduceOptions::default().valid_items(valid));", + "expected struct `ReduceOptions`", + ), + reduce_one_item_payload_count=( + "values.write(0,0_i32);let mut one=coop::ThreadData::<_,1>::new();one.write(0,1_i32);" + "let _result=coop::sum(block,&one,coop::ReduceOptions::default().valid_items(valid));", + "expected struct `ReduceOptions`", + ), + reduce_float_bitwise=( + "values.write(0,0_i32);let _result=coop::reduce(block,1_f32,coop::ReduceOptions::default().binary_op(coop::BitAnd));", + "ReduceOperator", + ), + reduce_warp_algorithm=( + "values.write(0,0_i32);let _result=coop::sum(coop::this_warp(),1_i32,coop::ReduceOptions::default().algorithm(coop::Raking));", + "expected `()`, found `Raking`", + ), + reduce_warp_storage=( + "values.write(0,0_i32);" + + new_storage + + "let _result=coop::sum(coop::this_warp(),1_i32,coop::ReduceOptions::default().temp_storage(&mut scratch));", + "expected struct `ReduceOptions`", + ), + reduce_thread_group=( + "values.write(0,0_i32);let _result=coop::sum(coop::this_thread(),1_i32,Default::default());", + "ReduceGroup", + ), + reduce_varying_count=( + "values.write(0,0_i32);let count=thread::threadIdx_x() as usize+1;" + "let _result=coop::sum(block,1_i32,coop::ReduceOptions::default().valid_items(count));", + "not proved block-uniform", + ), + reduce_varying_participation=( + load + "if thread::threadIdx_x()&1==0 {" + reduce_call + "}", + "lane-varying", + ), + reduce_warp_uniform_count=( + "values.write(0,0_i32);let warp=coop::this_warp();let count=warp.rank_at::() as usize+1;" + "let _result=coop::sum(warp,1_i32,coop::ReduceOptions::default().valid_items(count));", + None, + ), + reduce_warp_varying_count=( + "values.write(0,0_i32);let count=(thread::threadIdx_x()&31) as usize+1;" + "let _result=coop::sum(coop::this_warp(),1_i32,coop::ReduceOptions::default().valid_items(count));", + "not proved Warp-uniform", + ), + reduce_prefix_member=( + load + reduce_prefix + "if group.is_member()!=0 {" + "let _result=coop::sum(group,&values,Default::default());}", + None, + ), + reduce_prefix_missing_guard=( + load + + reduce_prefix + + "let _result=coop::sum(group,&values,Default::default());", + "lane-varying", + ), + reduce_prefix_arithmetic_guard=( + load + reduce_prefix + "if thread::threadIdx_x()%32<17 {" + "let _result=coop::sum(group,&values,Default::default());}", + "lane-varying", + ), + reduce_prefix_varying_count=( + "values.write(0,0_i32);" + reduce_prefix + "if group.is_member()!=0 {" + "let count=(thread::threadIdx_x()&31) as usize+1;" + "let _result=coop::sum(group,1_i32,coop::ReduceOptions::default().valid_items(count));}", + "not proved Warp-uniform", + ), + reduce_result_is_not_participation=( + load + "if let Some(total)=coop::sum(block,&values,Default::default()) {" + "let _result=coop::sum(block,total,Default::default());}", + "lane-varying", + ), + reduce_storage_auto_reuse=( + load + + new_storage.replace("new()", "new().auto_sync()") + + reduce_manual + + reduce_manual, + None, + ), + reduce_storage_reuse=( + load + new_storage + reduce_manual + sync + reduce_manual, + None, + ), + reduce_storage_missing_barrier=( + load + new_storage + reduce_manual + reduce_manual, + "requires a full-block barrier", + ), + reduce_no_global_completion=( + load + store + reduce_call + reload, + "requires a barrier covering both accesses", + ), +) + if args.case: unknown = set(args.case).difference(cases) if unknown: diff --git a/rust/cuda-coop/tests/device_queries/Cargo.lock b/rust/cuda-coop/tests/device_queries/Cargo.lock index 069a9b1a1235..2d1b9c982335 100644 --- a/rust/cuda-coop/tests/device_queries/Cargo.lock +++ b/rust/cuda-coop/tests/device_queries/Cargo.lock @@ -118,7 +118,7 @@ dependencies = [ [[package]] name = "cuda-artifact-finalizer" version = "0.2.1" -source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=732bec8e2166c33b9b851860363c81baa475203e#732bec8e2166c33b9b851860363c81baa475203e" dependencies = [ "libnvvm-sys", "nvjitlink-sys", @@ -130,7 +130,7 @@ dependencies = [ [[package]] name = "cuda-bindings" version = "0.5.0" -source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=732bec8e2166c33b9b851860363c81baa475203e#732bec8e2166c33b9b851860363c81baa475203e" dependencies = [ "bindgen", "libloading", @@ -161,7 +161,7 @@ dependencies = [ [[package]] name = "cuda-core" version = "0.5.0" -source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=732bec8e2166c33b9b851860363c81baa475203e#732bec8e2166c33b9b851860363c81baa475203e" dependencies = [ "anyhow", "cuda-bindings", @@ -173,7 +173,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" version = "0.5.0" -source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=732bec8e2166c33b9b851860363c81baa475203e#732bec8e2166c33b9b851860363c81baa475203e" dependencies = [ "proc-macro2", "quote", @@ -183,7 +183,7 @@ dependencies = [ [[package]] name = "cuda-device" version = "0.2.1" -source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=732bec8e2166c33b9b851860363c81baa475203e#732bec8e2166c33b9b851860363c81baa475203e" dependencies = [ "cuda-macros", ] @@ -191,7 +191,7 @@ dependencies = [ [[package]] name = "cuda-host" version = "0.2.1" -source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=732bec8e2166c33b9b851860363c81baa475203e#732bec8e2166c33b9b851860363c81baa475203e" dependencies = [ "cuda-artifact-finalizer", "cuda-core", @@ -208,7 +208,7 @@ dependencies = [ [[package]] name = "cuda-macros" version = "0.2.1" -source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=732bec8e2166c33b9b851860363c81baa475203e#732bec8e2166c33b9b851860363c81baa475203e" dependencies = [ "oxide-tile-export", "proc-macro2", @@ -222,7 +222,7 @@ dependencies = [ [[package]] name = "cuda-target-spec" version = "0.1.0" -source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=732bec8e2166c33b9b851860363c81baa475203e#732bec8e2166c33b9b851860363c81baa475203e" [[package]] name = "device_queries" @@ -379,7 +379,7 @@ dependencies = [ [[package]] name = "libnvvm-sys" version = "0.2.1" -source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=732bec8e2166c33b9b851860363c81baa475203e#732bec8e2166c33b9b851860363c81baa475203e" dependencies = [ "cuda-target-spec", "libloading", @@ -423,7 +423,7 @@ dependencies = [ [[package]] name = "nvjitlink-sys" version = "0.2.1" -source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=732bec8e2166c33b9b851860363c81baa475203e#732bec8e2166c33b9b851860363c81baa475203e" dependencies = [ "libloading", "thiserror", @@ -450,7 +450,7 @@ checksum = "9f7c3e4beb33f85d45ae3e3a1792185706c8e16d043238c593331cc7cd313b50" [[package]] name = "oxide-artifacts" version = "0.2.1" -source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=732bec8e2166c33b9b851860363c81baa475203e#732bec8e2166c33b9b851860363c81baa475203e" dependencies = [ "object", ] @@ -458,7 +458,7 @@ dependencies = [ [[package]] name = "oxide-tile-export" version = "0.2.1" -source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=732bec8e2166c33b9b851860363c81baa475203e#732bec8e2166c33b9b851860363c81baa475203e" dependencies = [ "serde", "serde_json", @@ -487,7 +487,7 @@ dependencies = [ [[package]] name = "ptx-parse" version = "0.1.0" -source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=732bec8e2166c33b9b851860363c81baa475203e#732bec8e2166c33b9b851860363c81baa475203e" [[package]] name = "quote" @@ -530,7 +530,7 @@ checksum = "d6f6ff9a378485b298a5286656da665ba74413d36db0979633275d2e708145d4" [[package]] name = "reserved-oxide-symbols" version = "0.2.1" -source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=732bec8e2166c33b9b851860363c81baa475203e#732bec8e2166c33b9b851860363c81baa475203e" [[package]] name = "rustc-hash" diff --git a/rust/cuda-coop/tests/device_queries/Cargo.toml b/rust/cuda-coop/tests/device_queries/Cargo.toml index 25e5d31be35b..2b149ba24726 100644 --- a/rust/cuda-coop/tests/device_queries/Cargo.toml +++ b/rust/cuda-coop/tests/device_queries/Cargo.toml @@ -11,6 +11,6 @@ publish = false [dependencies] cuda-coop = { path = "../.." } -cuda-device = { git = "https://github.com/tpn/cuda-rust.git", rev = "96caed46c776eb49226137eac849e795b83b7ce0" } -cuda-host = { git = "https://github.com/tpn/cuda-rust.git", rev = "96caed46c776eb49226137eac849e795b83b7ce0" } -cuda-core = { git = "https://github.com/tpn/cuda-rust.git", rev = "96caed46c776eb49226137eac849e795b83b7ce0" } +cuda-device = { git = "https://github.com/tpn/cuda-rust.git", rev = "732bec8e2166c33b9b851860363c81baa475203e" } +cuda-host = { git = "https://github.com/tpn/cuda-rust.git", rev = "732bec8e2166c33b9b851860363c81baa475203e" } +cuda-core = { git = "https://github.com/tpn/cuda-rust.git", rev = "732bec8e2166c33b9b851860363c81baa475203e" } diff --git a/rust/cuda-coop/tests/load_store/Cargo.lock b/rust/cuda-coop/tests/load_store/Cargo.lock index 514e255e8636..39bf1e33848f 100644 --- a/rust/cuda-coop/tests/load_store/Cargo.lock +++ b/rust/cuda-coop/tests/load_store/Cargo.lock @@ -135,7 +135,7 @@ dependencies = [ [[package]] name = "cuda-artifact-finalizer" version = "0.2.1" -source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=732bec8e2166c33b9b851860363c81baa475203e#732bec8e2166c33b9b851860363c81baa475203e" dependencies = [ "libnvvm-sys", "nvjitlink-sys", @@ -147,7 +147,7 @@ dependencies = [ [[package]] name = "cuda-async" version = "0.5.0" -source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=732bec8e2166c33b9b851860363c81baa475203e#732bec8e2166c33b9b851860363c81baa475203e" dependencies = [ "anyhow", "cuda-bindings", @@ -163,7 +163,7 @@ dependencies = [ [[package]] name = "cuda-bindings" version = "0.5.0" -source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=732bec8e2166c33b9b851860363c81baa475203e#732bec8e2166c33b9b851860363c81baa475203e" dependencies = [ "bindgen", "libloading", @@ -194,7 +194,7 @@ dependencies = [ [[package]] name = "cuda-core" version = "0.5.0" -source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=732bec8e2166c33b9b851860363c81baa475203e#732bec8e2166c33b9b851860363c81baa475203e" dependencies = [ "anyhow", "cuda-bindings", @@ -206,7 +206,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" version = "0.5.0" -source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=732bec8e2166c33b9b851860363c81baa475203e#732bec8e2166c33b9b851860363c81baa475203e" dependencies = [ "proc-macro2", "quote", @@ -216,7 +216,7 @@ dependencies = [ [[package]] name = "cuda-device" version = "0.2.1" -source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=732bec8e2166c33b9b851860363c81baa475203e#732bec8e2166c33b9b851860363c81baa475203e" dependencies = [ "cuda-macros", ] @@ -224,7 +224,7 @@ dependencies = [ [[package]] name = "cuda-host" version = "0.2.1" -source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=732bec8e2166c33b9b851860363c81baa475203e#732bec8e2166c33b9b851860363c81baa475203e" dependencies = [ "cuda-artifact-finalizer", "cuda-async", @@ -242,7 +242,7 @@ dependencies = [ [[package]] name = "cuda-macros" version = "0.2.1" -source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=732bec8e2166c33b9b851860363c81baa475203e#732bec8e2166c33b9b851860363c81baa475203e" dependencies = [ "oxide-tile-export", "proc-macro2", @@ -256,7 +256,7 @@ dependencies = [ [[package]] name = "cuda-target-spec" version = "0.1.0" -source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=732bec8e2166c33b9b851860363c81baa475203e#732bec8e2166c33b9b851860363c81baa475203e" [[package]] name = "dashmap" @@ -511,7 +511,7 @@ dependencies = [ [[package]] name = "libnvvm-sys" version = "0.2.1" -source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=732bec8e2166c33b9b851860363c81baa475203e#732bec8e2166c33b9b851860363c81baa475203e" dependencies = [ "cuda-target-spec", "libloading", @@ -564,7 +564,7 @@ dependencies = [ [[package]] name = "nvjitlink-sys" version = "0.2.1" -source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=732bec8e2166c33b9b851860363c81baa475203e#732bec8e2166c33b9b851860363c81baa475203e" dependencies = [ "libloading", "thiserror 2.0.21", @@ -591,7 +591,7 @@ checksum = "9f7c3e4beb33f85d45ae3e3a1792185706c8e16d043238c593331cc7cd313b50" [[package]] name = "oxide-artifacts" version = "0.2.1" -source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=732bec8e2166c33b9b851860363c81baa475203e#732bec8e2166c33b9b851860363c81baa475203e" dependencies = [ "object", ] @@ -599,7 +599,7 @@ dependencies = [ [[package]] name = "oxide-tile-export" version = "0.2.1" -source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=732bec8e2166c33b9b851860363c81baa475203e#732bec8e2166c33b9b851860363c81baa475203e" dependencies = [ "serde", "serde_json", @@ -647,7 +647,7 @@ dependencies = [ [[package]] name = "ptx-parse" version = "0.1.0" -source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=732bec8e2166c33b9b851860363c81baa475203e#732bec8e2166c33b9b851860363c81baa475203e" [[package]] name = "quote" @@ -699,7 +699,7 @@ checksum = "d6f6ff9a378485b298a5286656da665ba74413d36db0979633275d2e708145d4" [[package]] name = "reserved-oxide-symbols" version = "0.2.1" -source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=732bec8e2166c33b9b851860363c81baa475203e#732bec8e2166c33b9b851860363c81baa475203e" [[package]] name = "rustc-hash" diff --git a/rust/cuda-coop/tests/load_store/Cargo.toml b/rust/cuda-coop/tests/load_store/Cargo.toml index 626172bef415..52b9a20dfad2 100644 --- a/rust/cuda-coop/tests/load_store/Cargo.toml +++ b/rust/cuda-coop/tests/load_store/Cargo.toml @@ -11,7 +11,7 @@ publish = false [dependencies] cuda-coop = { path = "../.." } -cuda-device = { git = "https://github.com/tpn/cuda-rust.git", rev = "96caed46c776eb49226137eac849e795b83b7ce0" } -cuda-host = { git = "https://github.com/tpn/cuda-rust.git", rev = "96caed46c776eb49226137eac849e795b83b7ce0", features = ["async"] } -cuda-core = { git = "https://github.com/tpn/cuda-rust.git", rev = "96caed46c776eb49226137eac849e795b83b7ce0" } -cuda-async = { git = "https://github.com/tpn/cuda-rust.git", rev = "96caed46c776eb49226137eac849e795b83b7ce0" } +cuda-device = { git = "https://github.com/tpn/cuda-rust.git", rev = "732bec8e2166c33b9b851860363c81baa475203e" } +cuda-host = { git = "https://github.com/tpn/cuda-rust.git", rev = "732bec8e2166c33b9b851860363c81baa475203e", features = ["async"] } +cuda-core = { git = "https://github.com/tpn/cuda-rust.git", rev = "732bec8e2166c33b9b851860363c81baa475203e" } +cuda-async = { git = "https://github.com/tpn/cuda-rust.git", rev = "732bec8e2166c33b9b851860363c81baa475203e" } diff --git a/rust/cuda-coop/tests/reduce/.gitignore b/rust/cuda-coop/tests/reduce/.gitignore new file mode 100644 index 000000000000..74bffa8af4a4 --- /dev/null +++ b/rust/cuda-coop/tests/reduce/.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/reduce/Cargo.lock b/rust/cuda-coop/tests/reduce/Cargo.lock new file mode 100644 index 000000000000..8d47deaba2aa --- /dev/null +++ b/rust/cuda-coop/tests/reduce/Cargo.lock @@ -0,0 +1,1011 @@ +# 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-reduce-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=732bec8e2166c33b9b851860363c81baa475203e#732bec8e2166c33b9b851860363c81baa475203e" +dependencies = [ + "libnvvm-sys", + "nvjitlink-sys", + "serde", + "sha2", + "thiserror 2.0.21", +] + +[[package]] +name = "cuda-async" +version = "0.5.0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=732bec8e2166c33b9b851860363c81baa475203e#732bec8e2166c33b9b851860363c81baa475203e" +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.5.0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=732bec8e2166c33b9b851860363c81baa475203e#732bec8e2166c33b9b851860363c81baa475203e" +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.5.0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=732bec8e2166c33b9b851860363c81baa475203e#732bec8e2166c33b9b851860363c81baa475203e" +dependencies = [ + "anyhow", + "cuda-bindings", + "cuda-core-derive", + "half", + "oxide-artifacts", +] + +[[package]] +name = "cuda-core-derive" +version = "0.5.0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=732bec8e2166c33b9b851860363c81baa475203e#732bec8e2166c33b9b851860363c81baa475203e" +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=732bec8e2166c33b9b851860363c81baa475203e#732bec8e2166c33b9b851860363c81baa475203e" +dependencies = [ + "cuda-macros", +] + +[[package]] +name = "cuda-host" +version = "0.2.1" +source = "git+https://github.com/tpn/cuda-rust.git?rev=732bec8e2166c33b9b851860363c81baa475203e#732bec8e2166c33b9b851860363c81baa475203e" +dependencies = [ + "cuda-artifact-finalizer", + "cuda-async", + "cuda-core", + "cuda-device", + "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=732bec8e2166c33b9b851860363c81baa475203e#732bec8e2166c33b9b851860363c81baa475203e" +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=732bec8e2166c33b9b851860363c81baa475203e#732bec8e2166c33b9b851860363c81baa475203e" + +[[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=732bec8e2166c33b9b851860363c81baa475203e#732bec8e2166c33b9b851860363c81baa475203e" +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=732bec8e2166c33b9b851860363c81baa475203e#732bec8e2166c33b9b851860363c81baa475203e" +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=732bec8e2166c33b9b851860363c81baa475203e#732bec8e2166c33b9b851860363c81baa475203e" +dependencies = [ + "object", +] + +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +source = "git+https://github.com/tpn/cuda-rust.git?rev=732bec8e2166c33b9b851860363c81baa475203e#732bec8e2166c33b9b851860363c81baa475203e" +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=732bec8e2166c33b9b851860363c81baa475203e#732bec8e2166c33b9b851860363c81baa475203e" + +[[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=732bec8e2166c33b9b851860363c81baa475203e#732bec8e2166c33b9b851860363c81baa475203e" + +[[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/reduce/Cargo.toml b/rust/cuda-coop/tests/reduce/Cargo.toml new file mode 100644 index 000000000000..5a5eca34dba9 --- /dev/null +++ b/rust/cuda-coop/tests/reduce/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-reduce-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 = "732bec8e2166c33b9b851860363c81baa475203e" } +cuda-host = { git = "https://github.com/tpn/cuda-rust.git", rev = "732bec8e2166c33b9b851860363c81baa475203e", features = ["async"] } +cuda-core = { git = "https://github.com/tpn/cuda-rust.git", rev = "732bec8e2166c33b9b851860363c81baa475203e" } +cuda-async = { git = "https://github.com/tpn/cuda-rust.git", rev = "732bec8e2166c33b9b851860363c81baa475203e" } diff --git a/rust/cuda-coop/tests/reduce/src/main.rs b/rust/cuda-coop/tests/reduce/src/main.rs new file mode 100644 index 000000000000..9b32e2864fe7 --- /dev/null +++ b/rust/cuda-coop/tests/reduce/src/main.rs @@ -0,0 +1,527 @@ +// 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, thread}; +use cuda_host::cuda_module; + +#[cuda_module] +mod kernels { + use super::*; + + fn preserved( + values: &coop::ThreadData, + before: &coop::ThreadData, + ) -> bool { + let mut matches = true; + let mut i = 0; + while i < ITEMS_PER_THREAD { + matches &= values.get(i) == before.get(i); + i += 1; + } + matches && values.get(ITEMS_PER_THREAD).is_none() + } + + #[kernel] + #[launch_contract(domain = 1, block = (33, 1, 1), dynamic_shared = inferred)] + pub fn payload_sum( + input: &[T], + mut output: coop::DisjointSlice, + valid: usize, + zero: T, + invalid: bool, + ) { + let block = coop::this_block(); + let offset = block.rank_at::() as usize * 33 * ITEMS_PER_THREAD; + let mut values = coop::ThreadData::<_, ITEMS_PER_THREAD>::new(); + let options = coop::LoadOptions::default() + .offset(offset) + .valid_items(valid); + if invalid { + coop::load(block, input, &mut values, options); + } else { + coop::load(block, input, &mut values, options.oob_default(zero)); + } + let before = values; + let total = coop::sum(block, &values, Default::default()); + assert!(preserved(&values, &before)); + assert!(total.is_some() == (thread::threadIdx_x() == 0)); + if let Some(slot) = output.get_mut(thread::index_1d()) { + *slot = total.unwrap_or(zero); + } + } + + fn with_operator>( + block: coop::ThreadGroup, + values: &coop::ThreadData, + op: Op, + ) -> Option { + coop::reduce(block, values, coop::ReduceOptions::default().binary_op(op)) + } + + #[kernel] + #[launch_contract(domain = 3, block = (4, 4, 2), dynamic_shared = inferred)] + pub fn operators( + input: &[i32], + mut output: coop::DisjointSlice, + mode: u32, + ) { + let block = coop::this_block(); + let mut values = coop::ThreadData::<_, 4>::new(); + coop::load(block, input, &mut values, Default::default()); + let before = values; + let total = match mode { + 0 => coop::reduce(block, &mut values, Default::default()), + 1 => with_operator(block, &values, coop::Multiplies), + 2 => with_operator(block, &values, coop::Min), + 3 => with_operator(block, &values, coop::Max), + 4 => with_operator(block, &values, coop::BitAnd), + 5 => with_operator(block, &values, coop::BitOr), + 6 => with_operator(block, &values, coop::BitXor), + 7 => coop::sum( + block, + &values, + coop::SumOptions::default().algorithm(coop::RakingCommutativeOnly), + ), + 8 => coop::sum( + block, + &values, + coop::SumOptions::default().algorithm(coop::Raking), + ), + _ => coop::sum( + block, + &values, + coop::SumOptions::default().algorithm(coop::WarpReductions), + ), + }; + let leader = + thread::threadIdx_x() == 0 && thread::threadIdx_y() == 0 && thread::threadIdx_z() == 0; + assert!(preserved(&values, &before) && total.is_some() == leader); + if let Some(total) = total + && let Some(index) = thread::block_index() + && let Some(slot) = output.get_mut(index) + { + *slot = total; + } + } + + #[kernel] + #[launch_contract(domain = 3, block = (4, 4, 2), dynamic_shared = inferred)] + pub fn scalar_count( + mut output: coop::DisjointSlice, + valid: usize, + counted: bool, + ) { + let block = coop::this_block(); + let rank = thread::threadIdx_x() + 4 * (thread::threadIdx_y() + 4 * thread::threadIdx_z()); + let value = rank as i32 + 1; + let total = if counted { + coop::sum(block, value, coop::SumOptions::default().valid_items(valid)) + } else { + coop::reduce(block, value, Default::default()) + }; + assert!(total.is_some() == (rank == 0)); + if let Some(total) = total + && let Some(index) = thread::block_index() + && let Some(slot) = output.get_mut(index) + { + *slot = total; + } + } + + fn stored( + input: &[i32], + scratch: &mut coop::TempStorage<0, 0, AUTO, false>, + ) -> Option { + let block = coop::this_block(); + let mut values = coop::ThreadData::<_, 4>::new(); + coop::load(block, input, &mut values, Default::default()); + let before = values; + let mut matches = true; + let mut last = None; + let mut iteration = 0; + while iteration < 2 { + let total = coop::sum( + block, + &values, + coop::SumOptions::default().temp_storage(scratch), + ); + matches &= + preserved(&values, &before) && total.is_some() == (thread::threadIdx_x() == 0); + // Mutating a copied result cannot change the borrowed input. + let mut copied = total; + if let Some(value) = &mut copied { + *value += 1; + } + matches &= preserved(&values, &before); + if !AUTO { + thread::sync_threads(); + } + last = copied; + iteration += 1; + } + assert!(matches); + last + } + + #[kernel] + #[launch_contract(domain = 1, block = (32, 1, 1), dynamic_shared = inferred)] + pub fn storage(input: &[i32], mut output: coop::DisjointSlice, mode: u32) { + let total = match mode { + 0 => { + let block = coop::this_block(); + let mut values = coop::ThreadData::<_, 4>::new(); + coop::load(block, input, &mut values, Default::default()); + let before = values; + let first = coop::sum(block, &values, Default::default()); + let second = coop::sum(block, &values, Default::default()); + assert!(first == second && preserved(&values, &before)); + second + } + 1 => stored(input, &mut coop::TempStorage::new()), + _ => stored(input, &mut coop::TempStorage::new().auto_sync()), + }; + assert!(total.is_some() == (thread::threadIdx_x() == 0)); + if let Some(slot) = output.get_mut(thread::index_1d()) { + *slot = total.unwrap_or(-1); + } + } + + fn warp_payload, const ITEMS_PER_THREAD: usize>( + group: coop::ThreadGroup, + value: i32, + ) -> Option { + let mut values = coop::ThreadData::<_, ITEMS_PER_THREAD>::new(); + let mut i = 0; + while i < ITEMS_PER_THREAD { + values.write(i, value * (i as i32 + 1)); + i += 1; + } + let before = values; + let total = coop::sum(group, &values, Default::default()); + assert!(preserved(&values, &before)); + total + } + + #[kernel] + #[launch_contract(domain = 1, block = (64, 1, 1), dynamic_shared = inferred)] + pub fn logical(mut output: coop::DisjointSlice) { + let group = coop::this_warp().group_by::<8>(); + let tid = thread::threadIdx_x(); + let total = warp_payload::<_, ITEMS_PER_THREAD>(group, (tid / 8) as i32 + 1); + assert!(total.is_some() == (tid.is_multiple_of(8))); + if let Some(slot) = output.get_mut(thread::index_1d()) { + *slot = total.unwrap_or(-1); + } + } + + #[kernel] + #[launch_contract(domain = 1, block = (64, 1, 1), dynamic_shared = inferred)] + pub fn physical(mut output: coop::DisjointSlice) { + let tid = thread::threadIdx_x(); + let value = (tid / 32) as i32 + 1; + let total = coop::sum(coop::this_warp(), value, Default::default()); + assert!(total.is_some() == (tid.is_multiple_of(32))); + if let Some(slot) = output.get_mut(thread::index_1d()) { + *slot = total.unwrap_or(-1); + } + } + + #[kernel] + #[launch_contract(domain = 1, block = (64, 1, 1), dynamic_shared = inferred)] + pub fn prefix17(mut output: coop::DisjointSlice) { + let group = coop::this_warp().group_by_non_exhaustive::<17>(); + let tid = thread::threadIdx_x(); + let mut observed = -2; + if group.is_member() == 1 { + let total = warp_payload::<_, 1>(group, (tid / 32) as i32 + 1); + assert!(total.is_some() == (tid.is_multiple_of(32))); + observed = total.unwrap_or(-1); + } + if let Some(slot) = output.get_mut(thread::index_1d()) { + *slot = observed; + } + } + + #[kernel] + #[launch_contract(domain = 1, block = (64, 1, 1), dynamic_shared = inferred)] + pub fn prefix23(mut output: coop::DisjointSlice) { + let group = coop::this_warp().group_by_non_exhaustive::<23>(); + let tid = thread::threadIdx_x(); + let mut observed = -2; + if group.is_member() == 1 { + let total = warp_payload::<_, 4>(group, (tid / 32) as i32 + 1); + assert!(total.is_some() == (tid.is_multiple_of(32))); + observed = total.unwrap_or(-1); + } + if let Some(slot) = output.get_mut(thread::index_1d()) { + *slot = observed; + } + } + + #[kernel] + #[launch_contract(domain = 1, block = (64, 1, 1), dynamic_shared = inferred)] + pub fn prefix31(mut output: coop::DisjointSlice, valid: usize) { + let group = coop::this_warp().group_by_non_exhaustive::<31>(); + let tid = thread::threadIdx_x(); + let mut observed = -2; + if group.is_member() == 1 { + let value = (tid / 32 + 1) as i32 * (tid % 32 + 1) as i32; + let total = coop::sum(group, value, coop::SumOptions::default().valid_items(valid)); + assert!(total.is_some() == (tid.is_multiple_of(32))); + observed = total.unwrap_or(-1); + } + if let Some(slot) = output.get_mut(thread::index_1d()) { + *slot = observed; + } + } +} + +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, LaunchConfig3D}; + + init_device_contexts(0, 1)?; + let context = with_cuda_context(0, |context| context.clone())?; + let stream = context.default_stream(); + // SAFETY: the typed declarations and embedded kernels compile together. + let module = unsafe { kernels::load_async(0)? }; + let mode = std::env::args().nth(1); + if mode.as_deref() == Some("invalid-input") { + let plan = module.prepare_payload_sum::(LaunchConfig1D::auto_shared(1, 33))?; + let input = DeviceBuffer::from_host(&stream, &[1_i32; 132])?; + let output = DeviceBuffer::from_host(&stream, &[0_i32; 33])?; + module + .payload_sum_async_owned::(&plan, input, output, 131, 0, true) + .sync()?; + stream.synchronize()?; + panic!("uninitialized Reduce input unexpectedly completed"); + } + if let Some(count) = match mode.as_deref() { + Some("count-zero") => Some(0), + Some("count-high") => Some(33), + Some("count-max") => Some(usize::MAX), + _ => None, + } { + let plan = + module.prepare_scalar_count(LaunchConfig3D::auto_shared((1, 1, 1), (4, 4, 2)))?; + let output = DeviceBuffer::from_host(&stream, &[0_i32])?; + module + .scalar_count_async_owned(&plan, output, count, true) + .sync()?; + stream.synchronize()?; + panic!("invalid scalar Reduce count unexpectedly completed"); + } + let mut cases = 0; + macro_rules! typed { + ($ty:ty, $items:literal) => {{ + let tile: usize = 33 * $items; + let input: Vec<$ty> = (0..2 * tile) + .map(|i| { + (i.is_multiple_of($items) && (i / tile == 0 || (i / $items).is_multiple_of(2))) + as u8 as $ty + }) + .collect(); + let plan = + module.prepare_payload_sum::<$ty, $items>(LaunchConfig1D::auto_shared(2, 33))?; + for valid in [tile - 1, tile] { + let source = DeviceBuffer::from_host(&stream, &input)?; + let target = DeviceBuffer::from_host(&stream, &[0 as $ty; 66])?; + let (_, target) = module + .payload_sum_async_owned::<$ty, $items, _, _>( + &plan, source, target, valid, 0 as $ty, false, + ) + .sync()?; + let actual = target.to_host_vec(&stream)?; + for (block, chunk) in input.chunks(tile).enumerate() { + let mut total = 0 as $ty; + for value in &chunk[..valid] { + total += *value; + } + for lane in 0..33 { + assert_eq!( + actual[block * 33 + lane], + if lane == 0 { total } else { 0 as $ty }, + "{} N{} valid{} block{} lane{}", + stringify!($ty), + $items, + valid, + block, + lane + ); + } + } + cases += 1; + } + }}; + } + typed!(i8, 1); + typed!(i8, 4); + typed!(u8, 1); + typed!(u8, 4); + typed!(i16, 1); + typed!(i16, 4); + typed!(u16, 1); + typed!(u16, 4); + typed!(i32, 1); + typed!(i32, 4); + typed!(u32, 1); + typed!(u32, 4); + typed!(i64, 1); + typed!(i64, 4); + typed!(u64, 1); + typed!(u64, 4); + typed!(f32, 1); + typed!(f32, 4); + typed!(f64, 1); + typed!(f64, 4); + typed!(i32, 9); + // Empty input is legal for default-filled Load at offset zero; every + // payload slot then contains zero, unlike an invalid zero scalar count. + let plan = module.prepare_payload_sum::(LaunchConfig1D::auto_shared(1, 33))?; + let source = DeviceBuffer::from_host(&stream, &[] as &[i32])?; + let target = DeviceBuffer::from_host(&stream, &[-1_i32; 33])?; + let (_, target) = module + .payload_sum_async_owned::(&plan, source, target, 0, 0, false) + .sync()?; + assert_eq!(target.to_host_vec(&stream)?, [0; 33]); + cases += 1; + + let plan = module.prepare_operators(LaunchConfig3D::auto_shared((1, 1, 1), (4, 4, 2)))?; + for mode in 0..10 { + let input: Vec = (0..128) + .map(|i| { + if mode == 1 { + if i == 0 { 2 } else { 1 } + } else { + i % 7 + 1 + } + }) + .collect(); + let expected = input + .iter() + .copied() + .reduce(|a, b| match mode { + 1 => a * b, + 2 => a.min(b), + 3 => a.max(b), + 4 => a & b, + 5 => a | b, + 6 => a ^ b, + _ => a + b, + }) + .unwrap(); + let source = DeviceBuffer::from_host(&stream, &input)?; + let target = DeviceBuffer::from_host(&stream, &[-1_i32])?; + let (_, target) = module + .operators_async_owned(&plan, source, target, mode) + .sync()?; + assert_eq!( + target.to_host_vec(&stream)?, + [expected], + "operator mode {mode}" + ); + cases += 1; + } + let plan = module.prepare_scalar_count(LaunchConfig3D::auto_shared((1, 1, 1), (4, 4, 2)))?; + for (valid, counted) in [(1, true), (31, true), (32, true), (32, false)] { + let target = DeviceBuffer::from_host(&stream, &[-1_i32])?; + let target = module + .scalar_count_async_owned(&plan, target, valid, counted) + .sync()?; + assert_eq!( + target.to_host_vec(&stream)?, + [(valid * (valid + 1) / 2) as i32] + ); + cases += 1; + } + let plan = module.prepare_storage(LaunchConfig1D::auto_shared(1, 32))?; + for mode in 0..3 { + let source = DeviceBuffer::from_host(&stream, &[1_i32; 128])?; + let target = DeviceBuffer::from_host(&stream, &[-2_i32; 32])?; + let (_, target) = module + .storage_async_owned(&plan, source, target, mode) + .sync()?; + let expected: Vec<_> = (0..32) + .map(|lane| { + if lane == 0 { + if mode == 0 { 128 } else { 129 } + } else { + -1 + } + }) + .collect(); + assert_eq!( + target.to_host_vec(&stream)?, + expected, + "storage mode {mode}" + ); + cases += 1; + } + macro_rules! logical { + ($items:literal) => {{ + let plan = module.prepare_logical::<$items>(LaunchConfig1D::auto_shared(1, 64))?; + let target = DeviceBuffer::from_host(&stream, &[-3_i32; 64])?; + let target = module + .logical_async_owned::<$items, _>(&plan, target) + .sync()?; + let expected: Vec<_> = (0_i32..64) + .map(|tid| { + if tid % 8 == 0 { + 8 * (tid / 8 + 1) * $items * ($items + 1) / 2 + } else { + -1 + } + }) + .collect(); + assert_eq!(target.to_host_vec(&stream)?, expected); + cases += 1; + }}; + } + logical!(1); + logical!(4); + macro_rules! warp { + ($prepare:ident, $launch:ident, $width:literal, $items:literal) => {{ + let plan = module.$prepare(LaunchConfig1D::auto_shared(1, 64))?; + let target = DeviceBuffer::from_host(&stream, &[-3_i32; 64])?; + let target = module.$launch(&plan, target).sync()?; + let expected: Vec<_> = (0_i32..64) + .map(|tid| { + if tid % 32 >= $width { + -2 + } else if tid % 32 == 0 { + $width * (tid / 32 + 1) * $items * ($items + 1) / 2 + } else { + -1 + } + }) + .collect(); + assert_eq!(target.to_host_vec(&stream)?, expected, stringify!($launch)); + cases += 1; + }}; + } + warp!(prepare_physical, physical_async_owned, 32, 1); + warp!(prepare_prefix17, prefix17_async_owned, 17, 1); + warp!(prepare_prefix23, prefix23_async_owned, 23, 4); + let plan = module.prepare_prefix31(LaunchConfig1D::auto_shared(1, 64))?; + for valid in [1, 30, 31] { + let target = DeviceBuffer::from_host(&stream, &[-3_i32; 64])?; + let target = module.prefix31_async_owned(&plan, target, valid).sync()?; + let expected: Vec<_> = (0_i32..64) + .map(|tid| { + if tid % 32 == 31 { + -2 + } else if tid % 32 == 0 { + (tid / 32 + 1) * (valid * (valid + 1) / 2) as i32 + } else { + -1 + } + }) + .collect(); + assert_eq!(target.to_host_vec(&stream)?, expected); + cases += 1; + } + println!("PASS: {cases} public Reduce cases"); + Ok(()) +} diff --git a/rust/cuda-coop/tests/scan/Cargo.lock b/rust/cuda-coop/tests/scan/Cargo.lock index 0809eccb39af..0989d39428e1 100644 --- a/rust/cuda-coop/tests/scan/Cargo.lock +++ b/rust/cuda-coop/tests/scan/Cargo.lock @@ -135,7 +135,7 @@ dependencies = [ [[package]] name = "cuda-artifact-finalizer" version = "0.2.1" -source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=732bec8e2166c33b9b851860363c81baa475203e#732bec8e2166c33b9b851860363c81baa475203e" dependencies = [ "libnvvm-sys", "nvjitlink-sys", @@ -147,7 +147,7 @@ dependencies = [ [[package]] name = "cuda-async" version = "0.5.0" -source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=732bec8e2166c33b9b851860363c81baa475203e#732bec8e2166c33b9b851860363c81baa475203e" dependencies = [ "anyhow", "cuda-bindings", @@ -163,7 +163,7 @@ dependencies = [ [[package]] name = "cuda-bindings" version = "0.5.0" -source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=732bec8e2166c33b9b851860363c81baa475203e#732bec8e2166c33b9b851860363c81baa475203e" dependencies = [ "bindgen", "libloading", @@ -194,7 +194,7 @@ dependencies = [ [[package]] name = "cuda-core" version = "0.5.0" -source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=732bec8e2166c33b9b851860363c81baa475203e#732bec8e2166c33b9b851860363c81baa475203e" dependencies = [ "anyhow", "cuda-bindings", @@ -206,7 +206,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" version = "0.5.0" -source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=732bec8e2166c33b9b851860363c81baa475203e#732bec8e2166c33b9b851860363c81baa475203e" dependencies = [ "proc-macro2", "quote", @@ -216,7 +216,7 @@ dependencies = [ [[package]] name = "cuda-device" version = "0.2.1" -source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=732bec8e2166c33b9b851860363c81baa475203e#732bec8e2166c33b9b851860363c81baa475203e" dependencies = [ "cuda-macros", ] @@ -224,7 +224,7 @@ dependencies = [ [[package]] name = "cuda-host" version = "0.2.1" -source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=732bec8e2166c33b9b851860363c81baa475203e#732bec8e2166c33b9b851860363c81baa475203e" dependencies = [ "cuda-artifact-finalizer", "cuda-async", @@ -242,7 +242,7 @@ dependencies = [ [[package]] name = "cuda-macros" version = "0.2.1" -source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=732bec8e2166c33b9b851860363c81baa475203e#732bec8e2166c33b9b851860363c81baa475203e" dependencies = [ "oxide-tile-export", "proc-macro2", @@ -256,7 +256,7 @@ dependencies = [ [[package]] name = "cuda-target-spec" version = "0.1.0" -source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=732bec8e2166c33b9b851860363c81baa475203e#732bec8e2166c33b9b851860363c81baa475203e" [[package]] name = "dashmap" @@ -511,7 +511,7 @@ dependencies = [ [[package]] name = "libnvvm-sys" version = "0.2.1" -source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=732bec8e2166c33b9b851860363c81baa475203e#732bec8e2166c33b9b851860363c81baa475203e" dependencies = [ "cuda-target-spec", "libloading", @@ -564,7 +564,7 @@ dependencies = [ [[package]] name = "nvjitlink-sys" version = "0.2.1" -source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=732bec8e2166c33b9b851860363c81baa475203e#732bec8e2166c33b9b851860363c81baa475203e" dependencies = [ "libloading", "thiserror 2.0.21", @@ -591,7 +591,7 @@ checksum = "9f7c3e4beb33f85d45ae3e3a1792185706c8e16d043238c593331cc7cd313b50" [[package]] name = "oxide-artifacts" version = "0.2.1" -source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=732bec8e2166c33b9b851860363c81baa475203e#732bec8e2166c33b9b851860363c81baa475203e" dependencies = [ "object", ] @@ -599,7 +599,7 @@ dependencies = [ [[package]] name = "oxide-tile-export" version = "0.2.1" -source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=732bec8e2166c33b9b851860363c81baa475203e#732bec8e2166c33b9b851860363c81baa475203e" dependencies = [ "serde", "serde_json", @@ -647,7 +647,7 @@ dependencies = [ [[package]] name = "ptx-parse" version = "0.1.0" -source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=732bec8e2166c33b9b851860363c81baa475203e#732bec8e2166c33b9b851860363c81baa475203e" [[package]] name = "quote" @@ -699,7 +699,7 @@ checksum = "d6f6ff9a378485b298a5286656da665ba74413d36db0979633275d2e708145d4" [[package]] name = "reserved-oxide-symbols" version = "0.2.1" -source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=732bec8e2166c33b9b851860363c81baa475203e#732bec8e2166c33b9b851860363c81baa475203e" [[package]] name = "rustc-hash" diff --git a/rust/cuda-coop/tests/scan/Cargo.toml b/rust/cuda-coop/tests/scan/Cargo.toml index 05aab44b5af6..a4c015982eae 100644 --- a/rust/cuda-coop/tests/scan/Cargo.toml +++ b/rust/cuda-coop/tests/scan/Cargo.toml @@ -11,7 +11,7 @@ publish = false [dependencies] cuda-coop = { path = "../.." } -cuda-device = { git = "https://github.com/tpn/cuda-rust.git", rev = "96caed46c776eb49226137eac849e795b83b7ce0" } -cuda-host = { git = "https://github.com/tpn/cuda-rust.git", rev = "96caed46c776eb49226137eac849e795b83b7ce0", features = ["async"] } -cuda-core = { git = "https://github.com/tpn/cuda-rust.git", rev = "96caed46c776eb49226137eac849e795b83b7ce0" } -cuda-async = { git = "https://github.com/tpn/cuda-rust.git", rev = "96caed46c776eb49226137eac849e795b83b7ce0" } +cuda-device = { git = "https://github.com/tpn/cuda-rust.git", rev = "732bec8e2166c33b9b851860363c81baa475203e" } +cuda-host = { git = "https://github.com/tpn/cuda-rust.git", rev = "732bec8e2166c33b9b851860363c81baa475203e", features = ["async"] } +cuda-core = { git = "https://github.com/tpn/cuda-rust.git", rev = "732bec8e2166c33b9b851860363c81baa475203e" } +cuda-async = { git = "https://github.com/tpn/cuda-rust.git", rev = "732bec8e2166c33b9b851860363c81baa475203e" }