From 53da4923a2523814081c3a37f84a9082e34620b8 Mon Sep 17 00:00:00 2001 From: Trent Nelson Date: Sat, 24 Oct 2026 17:37:21 -0700 Subject: [PATCH 01/10] 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 02/10] 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 57d07306581cafb8fcd996ed477110fb0d3917d4 Mon Sep 17 00:00:00 2001 From: Trent Nelson Date: Sat, 24 Oct 2026 17:38:32 -0700 Subject: [PATCH 03/10] Guard batched warp reduction output copies by active lanes Avoid copying inactive batch output lanes while preserving active results and add focused coverage for narrow output batches. Signed-off-by: Trent Nelson --- .../warp_reduce_batched_wspro.cuh | 7 +- .../warp/catch2_test_warp_reduce_batched.cu | 79 +++++++++++++++++++ 2 files changed, 85 insertions(+), 1 deletion(-) diff --git a/cub/cub/warp/specializations/warp_reduce_batched_wspro.cuh b/cub/cub/warp/specializations/warp_reduce_batched_wspro.cuh index 4a36155469e1..889617648204 100644 --- a/cub/cub/warp/specializations/warp_reduce_batched_wspro.cuh +++ b/cub/cub/warp/specializations/warp_reduce_batched_wspro.cuh @@ -141,7 +141,12 @@ struct warp_reduce_batched_wspro _CCCL_PRAGMA_UNROLL_FULL() for (int i = 0; i < max_out_per_thread; ++i) { - outputs[i] = intermediate_outputs[i]; + const auto batch_idx = + ToBlocked ? logical_lane_id * max_out_per_thread + i : i * LogicalWarpThreads + logical_lane_id; + if (batch_idx < Batches) + { + outputs[i] = intermediate_outputs[i]; + } } } return true; diff --git a/cub/test/warp/catch2_test_warp_reduce_batched.cu b/cub/test/warp/catch2_test_warp_reduce_batched.cu index df744de2b2c7..632a85b4e495 100644 --- a/cub/test/warp/catch2_test_warp_reduce_batched.cu +++ b/cub/test/warp/catch2_test_warp_reduce_batched.cu @@ -1,6 +1,7 @@ // SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception +#include #include #include #include @@ -9,6 +10,8 @@ #include #include #include +#include +#include #include #include #include @@ -40,6 +43,82 @@ enum class WarpReduceBatchedMode ToBlocked }; +#if TEST_TYPES == 1 +struct assignment_counter_reference +{ + int* count; + + // Binding the source by reference observes the copy without reading an uninitialized scalar. + _CCCL_DEVICE_API void operator=(const int&) const noexcept + { + ++*count; + } +}; + +struct assignment_counter_accessor +{ + using offset_policy = assignment_counter_accessor; + using element_type = int; + using reference = assignment_counter_reference; + using data_handle_type = int*; + + [[nodiscard]] _CCCL_DEVICE_API reference access(data_handle_type counts, cuda::std::size_t index) const noexcept + { + return {counts + index}; + } + + [[nodiscard]] _CCCL_DEVICE_API data_handle_type offset(data_handle_type counts, cuda::std::size_t index) const noexcept + { + return counts + index; + } +}; + +__global__ void warp_reduce_batched_copy_count_kernel(int* counts, bool to_blocked) +{ + using warp_reduce_batched_t = cub::WarpReduceBatched; + __shared__ warp_reduce_batched_t::TempStorage temp_storage[block_size / warp_size]; + const int tid = static_cast(threadIdx.x); + const int logical_warp_id = tid / warp_size; + const cuda::std::array inputs{tid, tid + 1, tid + 2}; + int assignments = 0; + cuda::std::mdspan, cuda::std::layout_right, assignment_counter_accessor> outputs{ + &assignments}; + if (to_blocked) + { + warp_reduce_batched_t{temp_storage[logical_warp_id]}.ReduceToBlocked(inputs, outputs, cuda::std::plus<>{}); + } + else + { + warp_reduce_batched_t{temp_storage[logical_warp_id]}.ReduceToStriped(inputs, outputs, cuda::std::plus<>{}); + } + counts[tid] = assignments; +} + +CUB_TEST("WarpReduceBatched Redux copies only initialized intermediate outputs", "[warp][reduce][batched]", CUB_SMALL) +{ + int ptx_version = 0; + REQUIRE(cudaSuccess == cub::PtxVersion(ptx_version)); + if (ptx_version < 800) + { + SKIP("The Redux copy path requires SM80+; the fallback initializes every intermediate output."); + } + + const bool to_blocked = GENERATE(false, true); + CAPTURE(to_blocked); + c2h::device_vector counts(block_size); + warp_reduce_batched_copy_count_kernel<<<1, block_size>>>(thrust::raw_pointer_cast(counts.data()), to_blocked); + REQUIRE(cudaSuccess == cudaPeekAtLastError()); + REQUIRE(cudaSuccess == cudaDeviceSynchronize()); + + c2h::host_vector expected(block_size); + for (int tid = 0; tid < block_size; ++tid) + { + expected[tid] = tid % warp_size < 3 ? 1 : 0; + } + REQUIRE(expected == counts); +} +#endif // TEST_TYPES == 1 + template Date: Sat, 24 Oct 2026 17:40:52 -0700 Subject: [PATCH 04/10] 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" } From dc5456d9532c2033b5f08d2cf8d4064620473a1f Mon Sep 17 00:00:00 2001 From: Trent Nelson Date: Sat, 24 Oct 2026 17:41:46 -0700 Subject: [PATCH 05/10] Add Rust batched Reduce with exact provider shape identity Preserve the batch dimension and per-batch output validity while using the CUB output-copy guard. Include the provider shape correction. Signed-off-by: Trent Nelson --- docs/rust/index.rst | 9 +- docs/rust/payload.rst | 6 +- docs/rust/provider.rst | 9 + docs/rust/reduce_batched.rst | 135 +++ rust/.cargo/config.toml | 2 +- rust/cuda-coop-codegen/src/lib.rs | 39 +- rust/cuda-coop-codegen/src/oxide.rs | 214 +++- rust/cuda-coop-codegen/src/reduce.rs | 10 +- rust/cuda-coop-codegen/src/reduce_batched.rs | 208 ++++ rust/cuda-coop-codegen/src/scan.rs | 4 + rust/cuda-coop-codegen/tests/native.rs | 232 +++- rust/cuda-coop/build.rs | 17 +- .../examples/reduce_batched/.gitignore | 7 + .../examples/reduce_batched/Cargo.lock | 1019 +++++++++++++++++ .../examples/reduce_batched/Cargo.toml | 19 + .../examples/reduce_batched/helper/Cargo.toml | 11 + .../examples/reduce_batched/helper/src/lib.rs | 53 + .../examples/reduce_batched/src/main.rs | 113 ++ rust/cuda-coop/src/lib.rs | 14 +- rust/cuda-coop/src/reduce/boundary.rs | 17 +- rust/cuda-coop/src/reduce_batched.rs | 192 ++++ rust/cuda-coop/src/reduce_batched/boundary.rs | 80 ++ rust/cuda-coop/tests/check_load_store.py | 106 +- .../cuda-coop/tests/reduce_batched/.gitignore | 7 + .../cuda-coop/tests/reduce_batched/Cargo.lock | 1011 ++++++++++++++++ .../cuda-coop/tests/reduce_batched/Cargo.toml | 17 + rust/cuda-coop/tests/reduce_batched/README.md | 30 + .../tests/reduce_batched/src/main.rs | 418 +++++++ 28 files changed, 3959 insertions(+), 40 deletions(-) create mode 100644 docs/rust/reduce_batched.rst create mode 100644 rust/cuda-coop-codegen/src/reduce_batched.rs create mode 100644 rust/cuda-coop/examples/reduce_batched/.gitignore create mode 100644 rust/cuda-coop/examples/reduce_batched/Cargo.lock create mode 100644 rust/cuda-coop/examples/reduce_batched/Cargo.toml create mode 100644 rust/cuda-coop/examples/reduce_batched/helper/Cargo.toml create mode 100644 rust/cuda-coop/examples/reduce_batched/helper/src/lib.rs create mode 100644 rust/cuda-coop/examples/reduce_batched/src/main.rs create mode 100644 rust/cuda-coop/src/reduce_batched.rs create mode 100644 rust/cuda-coop/src/reduce_batched/boundary.rs create mode 100644 rust/cuda-coop/tests/reduce_batched/.gitignore create mode 100644 rust/cuda-coop/tests/reduce_batched/Cargo.lock create mode 100644 rust/cuda-coop/tests/reduce_batched/Cargo.toml create mode 100644 rust/cuda-coop/tests/reduce_batched/README.md create mode 100644 rust/cuda-coop/tests/reduce_batched/src/main.rs diff --git a/docs/rust/index.rst b/docs/rust/index.rst index 53926726b2d0..42e25b264a24 100644 --- a/docs/rust/index.rst +++ b/docs/rust/index.rst @@ -43,9 +43,11 @@ 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`, :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 +The public operations include :doc:`load_store`, :doc:`scan`, :doc:`reduce` +and :doc:`reduce_batched`. Scan returns fresh per-group prefixes; Reduce +returns a total only at each group's root. Batched reduction returns a fresh +payload whose extent and valid slots follow its Warp width and output layout. +Scan, Reduce and batched reduction preserve their input. Additional primitive families are under development. The crate is unpublished and its interface can change. @@ -116,6 +118,7 @@ Read next load_store scan reduce + reduce_batched provider Generate the API reference with ``cargo doc --workspace --no-deps`` from diff --git a/docs/rust/payload.rst b/docs/rust/payload.rst index 2991f1c19515..df414908189b 100644 --- a/docs/rust/payload.rst +++ b/docs/rust/payload.rst @@ -5,11 +5,15 @@ 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 +``ThreadDataLike`` interface. It is a ``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. +Ordinary payload consumers need no Rust feature gate. The calculated result +type of :doc:`reduce_batched` requires the pinned compiler's +``generic_const_exprs`` feature in the helper and consumer crates. + Construct and access a payload ------------------------------ diff --git a/docs/rust/provider.rst b/docs/rust/provider.rst index 344b94f5df9e..46b8f60ec784 100644 --- a/docs/rust/provider.rst +++ b/docs/rust/provider.rst @@ -86,6 +86,15 @@ 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. +``ReduceBatchedRequest`` has different input and output extents: each lane +supplies ``B`` batches and receives capacity ``ceil(B/W)`` at Warp width ``W``. +Its striped or blocked layout marks only actual batch results initialized. +The native request carries ``B`` and derives the output extent; the Rust API +uses the same calculation in its return type. See :doc:`reduce_batched` for +the unstable feature and generic-signature requirements. The provider +queries this operation's actual CUB storage, including a nonzero empty-type +layout, and retains ordinary group replication and reuse planning. + Compile a provider ------------------ diff --git a/docs/rust/reduce_batched.rst b/docs/rust/reduce_batched.rst new file mode 100644 index 000000000000..49f175a02298 --- /dev/null +++ b/docs/rust/reduce_batched.rst @@ -0,0 +1,135 @@ +.. Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. +.. SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +Warp batched reduction +====================== + +``reduce_batched`` reduces each payload slot independently across a Warp group. +If every lane supplies ``B`` items, there are ``B`` reductions: batch ``i`` +combines slot ``i`` from every member. The operation takes the group first, +preserves the borrowed input, and returns fresh ``ThreadData`` with the same +element type. Defaults select ``Sum`` and striped output ownership. + +Physical Warps and logical widths 1, 2, 4, 8, 16 and 32 are supported within +blocks containing complete physical warps. Each group reduces independently; +all its members participate. Block groups, non-power-of-two Warp prefixes and +scalar inputs are unsupported. Use :doc:`reduce` to combine every slot into +one root-only result. + +A complete program +------------------ + +The example uses eight-lane groups in three 64-thread blocks. It checks one, +four and ten input items per thread, including a short input tail and padding +between groups' output regions. The element type is inferred from the input. + +.. literalinclude:: ../../rust/cuda-coop/examples/reduce_batched/src/main.rs + :language: rust + :start-after: // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +The separate generic helper fills the missing input tail with zero before +summing. Every member calls both the reduction and Store, including lanes +whose result contains no initialized slots. + +.. literalinclude:: ../../rust/cuda-coop/examples/reduce_batched/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/reduce_batched + 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 +operations and plans their temporary storage. The consumer needs no C++ shim, +provider build script or manual scratch-byte calculation. + +Derived extent and validity +--------------------------- + +The result capacity per lane is ``K = batched_extent(B, W) = ceil(B / W)``, +where ``W`` is the group width. ``B`` is positive. The return type is +``ThreadData`` with default alignment; input alignment does not change +the result type. The caller supplies no separate output extent. Capacity and +validity are distinct: ``get(i)`` returns ``None`` when a slot has no batch. +Such slots remain unwritten, including when an entire lane has no result. + +For lane rank ``r`` and output slot ``i``, ownership is: + +.. list-table:: + :header-rows: 1 + :widths: 25 40 35 + + * - Output layout + - Batch index + - Validity + * - ``Striped`` (default) + - ``r + i * W`` + - Batch index is below ``B`` + * - ``Blocked`` + - ``r * K + i`` + - Batch index is below ``B`` + +For ``B = 10`` and ``W = 8``, each lane has capacity two. Striped lanes zero +and one receive batches ``(0, 8)`` and ``(1, 9)``; lanes two through seven +have only their first slot initialized. With blocked output, lanes zero +through four each receive two consecutive batches, and lanes five through +seven have no initialized result. + +Every input slot must be initialized. An unwritten input traps before its +numeric value is read. There is no valid-count or seed option. For a partial +input tile, fill missing items explicitly with a value suitable for the +operator, as the sum example does with zero. Invalid outputs are not filled +with zero or broadcast to other lanes. + +Calculated types on the pinned compiler +--------------------------------------- + +The calculated return signature currently requires +``#![feature(generic_const_exprs)]`` in every crate that instantiates it, +including a separate helper crate and its consumer. The example also uses +``#![allow(incomplete_features)]`` because this Rust feature is incomplete. +Generic helpers repeat the extent expression in the well-formedness bound +``[(); coop::batched_extent(ITEMS_PER_THREAD, 8)]:``, as shown in the executed +source above. Helpers generic over the group use +``::WIDTH`` in that expression. + +This is a limitation of the pinned compiler and calculated-result signature. +Ordinary payload, Load/Store, Scan and Reduce consumers do not need this +feature gate. The provider also checks CUB's signed-index limit for the padded +``W * K`` extent. + +Operators and safe Store +------------------------ + +``ReduceBatchedOptions::default()`` and ``Default::default()`` infer ``T``. +The consuming ``.binary_op(op)`` builder selects a builtin operator; +``.output_layout(coop::Blocked)`` selects blocked result ownership. +The operator value can remain generic under ``Op: ReduceOperator``. +All ten numeric types support ``Sum``, ``Multiplies``, ``Min`` and ``Max``; +integers also support ``BitAnd``, ``BitOr`` and ``BitXor``. Results retain the +input dtype. CUB may regroup floating-point arithmetic, so results need not +match a sequential host fold exactly. Custom operators are not yet supported. +There is no explicit algorithm or temporary-storage option. + +Store a striped result with the ``Striped`` Store algorithm; store a blocked +result with ``Direct``. In either case, ``valid_items(B)`` selects exactly the +initialized batch results. Keep Store collective: do not guard it with a +lane-local ``get(...).is_some()`` test. + +The output carrier is ``DisjointSlice>``. Each group's +region has ``W * K`` slots, of which the first ``B`` are written by the matching +Store. Padding remains unchanged. The example passes the block's output +offset, ``block_rank * block_threads * K``; Warp Store adds the subgroup's +origin itself. Do not add that subgroup origin a second time or substitute a +compact ``B``-slot stride for this ownership region. + +Scratch is implicit, queried from CUB, and planned per group with completion +before reuse. Result validity does not grant global-memory ownership or +complete unrelated Load/Store accesses. Retain the matching group's ``sync()`` +when a later global access depends on an earlier one. See :doc:`load_store` +for the output carrier and completion rules. diff --git a/rust/.cargo/config.toml b/rust/.cargo/config.toml index 0ab90f87e900..d59034c2bf7e 100644 --- a/rust/.cargo/config.toml +++ b/rust/.cargo/config.toml @@ -2,6 +2,6 @@ # SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception [env] -# cuda-device uses f16. +# cuda-coop uses generic_const_exprs; cuda-device uses f16. # Enable these features in processes Cargo launches from this tree. RUSTC_BOOTSTRAP = "1" diff --git a/rust/cuda-coop-codegen/src/lib.rs b/rust/cuda-coop-codegen/src/lib.rs index aa62e89e6b2b..630bd6ad603c 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, Scan and Reduce providers with +//! This private development crate compiles Load/Store, Scan, Reduce and batched 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. @@ -39,6 +39,9 @@ //! 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. +//! [`ReduceBatchedRequest`] reduces each payload position independently across +//! a power-of-two Warp. Its fresh per-lane output capacity is ceil(batches/width), +//! with striped or blocked distribution and validity only for real batches. //! [`Provider::compile_native`] batches these families in one native compilation. //! //! [`Provider::from_checkout`] snapshots the selected CCCL and CUDA headers; @@ -70,8 +73,10 @@ mod headers; mod reduce; +mod reduce_batched; mod scan; pub use reduce::{ReduceAlgorithm, ReduceForm, ReduceOperator, ReduceRecipe, ReduceRequest}; +pub use reduce_batched::{ReduceBatchedLayout, ReduceBatchedRecipe, ReduceBatchedRequest}; pub use scan::{ScanAlgorithm, ScanOperation, ScanOperator, ScanRecipe, ScanRequest}; #[allow(unsafe_code)] mod nvrtc; @@ -457,14 +462,26 @@ pub enum NativeRequest { Scan(ScanRequest), /// Numeric reduction with preserved input and a fresh root-only scalar. Reduce(ReduceRequest), + /// Batched Warp reduction with preserved input and partial fresh outputs. + ReduceBatched(ReduceBatchedRequest), } impl NativeRequest { + fn block_shape(self) -> [u32; 3] { + match self { + Self::LoadStore(request) => request.block_shape(), + Self::Scan(request) => request.block_shape(), + Self::Reduce(request) => request.block_shape(), + Self::ReduceBatched(request) => request.block_shape(), + } + } + fn source(self, symbol: &str) -> String { match self { Self::LoadStore(request) => source(request, symbol), Self::Scan(request) => request.source(symbol), Self::Reduce(request) => request.source(symbol), + Self::ReduceBatched(request) => request.source(symbol), } } @@ -473,13 +490,14 @@ impl NativeRequest { Self::LoadStore(request) => request.operation.token().into(), Self::Scan(request) => request.recipe.operation().token(), Self::Reduce(request) => request.recipe.token(), + Self::ReduceBatched(request) => request.recipe.token(), } } fn scratch_free(self) -> bool { match self { Self::LoadStore(request) => request.algorithm.scratch_free(), - Self::Scan(_) | Self::Reduce(_) => false, + Self::Scan(_) | Self::Reduce(_) | Self::ReduceBatched(_) => false, } } @@ -488,8 +506,8 @@ impl NativeRequest { Self::LoadStore(request) => { !request.operation.is_load() && !request.algorithm.scratch_free() } - // Scan and Reduce write fresh outputs, never rearranging inputs. - Self::Scan(_) | Self::Reduce(_) => false, + // These families write fresh outputs, never rearranging inputs. + Self::Scan(_) | Self::Reduce(_) | Self::ReduceBatched(_) => false, } } } @@ -508,8 +526,9 @@ pub struct Layout { /// # Device ABI obligations /// /// The memory-tile and payload-layout requirements below describe Load/Store. -/// Scan and Reduce use lane-private input and fresh output; their ABI and -/// validity contracts are documented by [`ScanRequest`] and [`ReduceRequest`]. +/// Scan, Reduce and batched Reduce use lane-private input and fresh output; +/// their ABI and validity contracts are documented by [`ScanRequest`], +/// [`ReduceRequest`] and [`ReduceBatchedRequest`]. /// /// 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. @@ -660,10 +679,10 @@ impl Provider { self.compile_native(&requests, target) } - /// Compile one ordered mixed Load/Store, Scan and Reduce batch in one NVRTC program. + /// Compile an ordered mixed batch of the supported families in one NVRTC program. /// /// Per-specialization identity is independent of unrelated batch members. - /// Scan/Reduce scratch describes one group and includes the actual CUB TempStorage, + /// Scan/Reduce/batched 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 @@ -727,6 +746,7 @@ impl Provider { include_bytes!("block_io.cuh"), include_bytes!("scan.rs"), include_bytes!("reduce.rs"), + include_bytes!("reduce_batched.rs"), ] { hash_field(&mut inputs, bytes); } @@ -748,6 +768,9 @@ impl Provider { _ => unreachable!("callback recipes are validated before native compilation"), }; let mut hash = inputs.clone(); + for axis in request.block_shape() { + hash_field(&mut hash, &axis.to_le_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); diff --git a/rust/cuda-coop-codegen/src/oxide.rs b/rust/cuda-coop-codegen/src/oxide.rs index ba64ee917b78..411300aca6d2 100644 --- a/rust/cuda-coop-codegen/src/oxide.rs +++ b/rust/cuda-coop-codegen/src/oxide.rs @@ -13,7 +13,8 @@ //! 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`], numeric Block/scalar Warp Scan through [`ScanRecipe`], and -//! root-result scalar/payload Reduce through [`ReduceRecipe`]. +//! root-result scalar/payload Reduce through [`ReduceRecipe`], and partial-output +//! Warp batched Reduce through [`ReduceBatchedRecipe`]. //! 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 @@ -22,8 +23,8 @@ use crate::{ Algorithm, Artifact, Dtype, Error, Layout, NativeRequest, Operation, Provider, ReduceAlgorithm, - ReduceForm, ReduceOperator, ReduceRecipe, ReduceRequest, Request, Result, ScanRecipe, - ScanRequest, + ReduceBatchedRecipe, ReduceBatchedRequest, ReduceForm, ReduceOperator, ReduceRecipe, + ReduceRequest, Request, Result, ScanRecipe, ScanRequest, }; use serde::{Deserialize, Serialize}; use std::io::{Read, Write}; @@ -190,6 +191,20 @@ struct SpecializationRequest { impl SpecializationRequest { fn native(&self) -> Result { + if self.operation.starts_with("reduce_batched_") { + let recipe = ReduceBatchedRecipe::parse(&self.operation, &self.algorithm)?; + let width = self + .warp_width + .ok_or_else(|| Error::message("batched Reduce requires a Warp scope"))?; + return ReduceBatchedRequest::warp( + self.dtype.dtype()?, + recipe, + self.block, + width, + self.items_per_thread, + ) + .map(NativeRequest::ReduceBatched); + } if self.operation.starts_with("reduce_") { let recipe = ReduceRecipe::parse(&self.operation, &self.algorithm)?; let dtype = self.dtype.dtype()?; @@ -315,6 +330,15 @@ struct BatchResponse { fn arguments(request: NativeRequest) -> Vec { let request = match request { + NativeRequest::ReduceBatched(request) => { + let dtype = Scalar::from(request.dtype); + return vec![ + Argument::ConstPointer(dtype), + Argument::MutPointer(dtype), + Argument::ConstPointer(Scalar::Unsigned(8)), + Argument::MutPointer(Scalar::Unsigned(8)), + ]; + } NativeRequest::Reduce(request) => { let dtype = Scalar::from(request.dtype); let mut arguments = vec![ @@ -527,6 +551,32 @@ pub fn reduce_registration(executable: &Path, recipe: ReduceRecipe) -> 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: &recipe.token(), + algorithm: "warp", + }) + .map_err(|error| Error::message(format!("cannot encode provider registration: {error}"))) +} + #[derive(Deserialize)] #[serde(deny_unknown_fields)] struct InputRequest { @@ -844,7 +894,163 @@ mod tests { } use super::*; - use crate::{ReduceAlgorithm, ReduceOperator, ScanAlgorithm, ScanOperation, ScanOperator}; + use crate::{ + ReduceAlgorithm, ReduceBatchedLayout, ReduceOperator, ScanAlgorithm, ScanOperation, + ScanOperator, + }; + + #[test] + fn batched_registration_preserves_derived_extent_and_pointer_abi() { + for (dtype, op, layout, width, batches, outputs) in [ + ( + Dtype::I8, + ReduceOperator::Sum, + ReduceBatchedLayout::Striped, + 32, + 1, + 1, + ), + ( + Dtype::U16, + ReduceOperator::BitAnd, + ReduceBatchedLayout::Blocked, + 8, + 10, + 2, + ), + ( + Dtype::F64, + ReduceOperator::Max, + ReduceBatchedLayout::Striped, + 1, + 9, + 9, + ), + ] { + let recipe = ReduceBatchedRecipe::new(op, layout); + let registration: serde_json::Value = serde_json::from_str( + &reduce_batched_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: batches, + warp_width: Some(width), + } + .native() + .unwrap(); + let NativeRequest::ReduceBatched(native) = request else { + panic!("batched recipe decoded as another family"); + }; + assert_eq!(native.outputs_per_thread(), outputs); + assert_eq!(native.recipe, recipe); + assert!(!request.scratch_free()); + assert!(!request.mutates_payload()); + assert_eq!( + arguments(request), + vec![ + Argument::ConstPointer(scalar), + Argument::MutPointer(scalar), + Argument::ConstPointer(Scalar::Unsigned(8)), + Argument::MutPointer(Scalar::Unsigned(8)), + ] + ); + assert_eq!(registration["provider"], "cuda_coop"); + assert_eq!(registration["version"], 1); + } + let recipe = ReduceBatchedRecipe::new(ReduceOperator::Sum, ReduceBatchedLayout::Striped); + assert!(reduce_batched_registration(Path::new("relative"), recipe).is_err()); + let last_padded = (i32::MAX as u32 / 32) * 32; + assert_eq!( + ReduceBatchedRequest::warp(Dtype::I32, recipe, [32, 1, 1], 32, last_padded) + .unwrap() + .outputs_per_thread(), + last_padded / 32, + ); + assert!( + ReduceBatchedRequest::warp(Dtype::I32, recipe, [32, 1, 1], 32, last_padded + 1,) + .is_err() + ); + } + + #[test] + fn rejects_invalid_batched_reduce_before_opening_toolkit() { + let base = serde_json::json!({"version":1,"target":"compute_120","cuda_root":"/absent", "requests":[{ + "operation":"reduce_batched_bit_xor_blocked_validity_bits_v1","algorithm":"warp", + "dtype":{"kind":"signed","bits":32},"block":[4,4,4],"items_per_thread":10,"warp_width":8 + }]}); + for (path, value, diagnostic) in [ + ( + "/requests/0/dtype", + serde_json::json!({"kind":"float","bits":32}), + "integer", + ), + ( + "/requests/0/algorithm", + serde_json::json!("raking"), + "Warp algorithm", + ), + ( + "/requests/0/operation", + serde_json::json!("reduce_batched_sum_payload_validity_bits_v1"), + "output layout", + ), + ( + "/requests/0/operation", + serde_json::json!("reduce_batched_unknown_striped_validity_bits_v1"), + "operator", + ), + ( + "/requests/0/items_per_thread", + serde_json::json!(0), + "positive batches", + ), + ( + "/requests/0/items_per_thread", + serde_json::json!(u32::MAX), + "padded extent", + ), + ( + "/requests/0/warp_width", + serde_json::json!(null), + "Warp scope", + ), + ( + "/requests/0/warp_width", + serde_json::json!(17), + "power-of-two", + ), + ( + "/requests/0/warp_width", + serde_json::json!(0), + "power-of-two", + ), + ( + "/requests/0/block", + serde_json::json!([48, 1, 1]), + "complete physical warps", + ), + ( + "/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}"); + } + } #[test] fn reduce_registration_preserves_scalar_count_and_payload_abi() { diff --git a/rust/cuda-coop-codegen/src/reduce.rs b/rust/cuda-coop-codegen/src/reduce.rs index 6c6c9dbcd415..1b81bb3c7b6d 100644 --- a/rust/cuda-coop-codegen/src/reduce.rs +++ b/rust/cuda-coop-codegen/src/reduce.rs @@ -25,7 +25,7 @@ pub enum ReduceOperator { } impl ReduceOperator { - fn token(self) -> &'static str { + pub(crate) fn token(self) -> &'static str { match self { Self::Sum => "sum", Self::Multiplies => "multiplies", @@ -37,7 +37,7 @@ impl ReduceOperator { } } - fn parse(token: &str) -> Result { + pub(crate) fn parse(token: &str) -> Result { match token { "sum" => Ok(Self::Sum), "multiplies" => Ok(Self::Multiplies), @@ -50,7 +50,7 @@ impl ReduceOperator { } } - fn functor(self) -> &'static str { + pub(crate) fn functor(self) -> &'static str { match self { Self::Sum => "cuda::std::plus", Self::Multiplies => "cuda::std::multiplies", @@ -284,6 +284,10 @@ impl ReduceRequest { }) } + pub(crate) fn block_shape(self) -> [u32; 3] { + self.block + } + pub(crate) fn source(self, symbol: &str) -> String { self.source_with_callback(symbol, None) } diff --git a/rust/cuda-coop-codegen/src/reduce_batched.rs b/rust/cuda-coop-codegen/src/reduce_batched.rs new file mode 100644 index 000000000000..09b3e99db31f --- /dev/null +++ b/rust/cuda-coop-codegen/src/reduce_batched.rs @@ -0,0 +1,208 @@ +// Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +//! Validated Warp batched reductions with partial, derived-extent outputs. + +use crate::{Dtype, Error, ReduceOperator, Result}; + +/// Distribution of batch results among the members of a Warp. +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub enum ReduceBatchedLayout { + /// Lane r receives batches r, r + width, and so on. + Striped, + /// Lane r receives consecutive batches starting at r * ceil(batches/width). + Blocked, +} + +/// Static batched-reduction operator and output layout. +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub struct ReduceBatchedRecipe { + operator: ReduceOperator, + layout: ReduceBatchedLayout, +} + +impl ReduceBatchedRecipe { + /// Select a builtin operator and the distribution of fresh output slots. + pub fn new(operator: ReduceOperator, layout: ReduceBatchedLayout) -> Self { + Self { operator, layout } + } + + pub(crate) fn token(self) -> String { + let layout = match self.layout { + ReduceBatchedLayout::Striped => "striped", + ReduceBatchedLayout::Blocked => "blocked", + }; + format!( + "reduce_batched_{}_{layout}_validity_bits_v1", + self.operator.token() + ) + } + + pub(crate) fn parse(operation: &str, algorithm: &str) -> Result { + if algorithm != "warp" { + return Err(Error::message("batched Reduce requires the Warp algorithm")); + } + let operation = operation + .strip_prefix("reduce_batched_") + .and_then(|s| s.strip_suffix("_validity_bits_v1")) + .ok_or_else(|| Error::message("unsupported batched Reduce operation"))?; + let (operator, layout) = operation + .rsplit_once('_') + .ok_or_else(|| Error::message("missing batched Reduce output layout"))?; + let layout = match layout { + "striped" => ReduceBatchedLayout::Striped, + "blocked" => ReduceBatchedLayout::Blocked, + _ => return Err(Error::message("unsupported batched Reduce output layout")), + }; + Ok(Self::new(ReduceOperator::parse(operator)?, layout)) + } +} + +/// Validated numeric batched reduction within a physical or logical Warp. +/// +/// Native calls use `(const T* input, T* output, u8* scratch, +/// const u8* input_bits, u8* output_bits)` with no scalar controls. Each lane +/// supplies `batches` initialized input slots and `ceil(batches/8)` packed, +/// low-bit-first validity bytes. Every input bit is checked before any numeric +/// read. Inputs and their bits are preserved. +/// +/// Output has exactly `ceil(batches/width)` slots and enough validity bytes for +/// that extent. Output values and bits are fresh, disjoint from input, live and +/// aligned. Only slots corresponding to real batches are written and marked +/// initialized; the other slots remain uninitialized, with zero validity bits. +/// All output validity bytes, including unused high bits, are reset. +/// +/// Every member of the aligned Warp participates. Blocks contain complete +/// physical warps, and supported widths are 1/2/4/8/16/32. Scratch describes +/// the actual CUB TempStorage for one group, including its nonzero empty-type +/// size. The caller owns its replication, alignment, lifetime and reuse +/// completion. This provider adds no completion barrier or result broadcast. +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub struct ReduceBatchedRequest { + pub(crate) recipe: ReduceBatchedRecipe, + pub(crate) dtype: Dtype, + block: [u32; 3], + width: u32, + batches: u32, +} + +impl ReduceBatchedRequest { + /// Validate one batched reduction in an aligned Warp with exact XYZ geometry. + /// + /// The input extent is `batches`; the output extent is derived from `width`. + /// Bitwise operators require integer elements. + /// + /// # Errors + /// Rejects zero/overflowing axes, more than 1024 threads, incomplete physical + /// warps, widths other than 1/2/4/8/16/32, zero batches, or a padded output + /// extent exceeding CUB's signed-int indexing range. + pub fn warp( + dtype: Dtype, + recipe: ReduceBatchedRecipe, + block: [u32; 3], + width: u32, + batches: u32, + ) -> Result { + let threads = block + .into_iter() + .try_fold(1_u32, u32::checked_mul) + .filter(|n| (1..=1024).contains(n)) + .ok_or_else(|| { + Error::message("batched Reduce requires 1..=1024 block threads with positive axes") + })?; + if threads % 32 != 0 { + return Err(Error::message( + "batched Reduce requires complete physical warps", + )); + } + if !width.is_power_of_two() || width > 32 { + return Err(Error::message( + "batched Reduce requires power-of-two width 1..=32", + )); + } + let outputs = batches / width + u32::from(batches % width != 0); + if batches == 0 || u64::from(outputs) * u64::from(width) > i32::MAX as u64 { + return Err(Error::message( + "batched Reduce requires positive batches and padded extent <= i32::MAX", + )); + } + if matches!(dtype, Dtype::F32 | Dtype::F64) + && matches!( + recipe.operator, + ReduceOperator::BitAnd | ReduceOperator::BitOr | ReduceOperator::BitXor + ) + { + return Err(Error::message( + "bitwise batched Reduce operators require integer elements", + )); + } + Ok(Self { + recipe, + dtype, + block, + width, + batches, + }) + } + + /// Derived per-lane output capacity; individual slots may be invalid. + pub fn outputs_per_thread(self) -> u32 { + self.batches / self.width + u32::from(self.batches % self.width != 0) + } + + pub(crate) fn block_shape(self) -> [u32; 3] { + self.block + } + + pub(crate) fn source(self, symbol: &str) -> String { + let (dtype, _) = self.dtype.native(); + let [x, y, z] = self.block; + let width = self.width; + let batches = self.batches; + let outputs = self.outputs_per_thread(); + let output_bytes = outputs / 8 + u32::from(outputs % 8 != 0); + let (method, batch_index) = match self.recipe.layout { + ReduceBatchedLayout::Striped => ("ReduceToStriped", format!("lane_rank + i * {width}")), + ReduceBatchedLayout::Blocked => { + ("ReduceToBlocked", format!("lane_rank * {outputs} + i")) + } + }; + let functor = self.recipe.operator.functor(); + format!( + r#" +#include +#include +#include +#include +namespace {symbol}_detail {{ +using value_type = {dtype}; +using primitive = cub::WarpReduceBatched; +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) {{ + using namespace {symbol}_detail; + const unsigned int lane_rank = cub::RowMajorTid({x}, {y}, {z}) % {width}; +#pragma unroll + for (int i = 0; i < {output_bytes}; ++i) output_bits[i] = 0; +#pragma unroll + for (int i = 0; i < {batches}; ++i) + if ((input_bits[i / 8] & (1u << (i % 8))) == 0) __trap(); + value_type local_input[{batches}]; +#pragma unroll + for (int i = 0; i < {batches}; ++i) local_input[i] = input[i]; + value_type local_output[{outputs}]; + primitive operation{{*reinterpret_cast(scratch)}}; + operation.{method}(local_input, local_output, {functor}{{}}); +#pragma unroll + for (int i = 0; i < {outputs}; ++i) {{ + if ({batch_index} < {batches}) {{ + output[i] = local_output[i]; + output_bits[i / 8] |= static_cast(1u << (i % 8)); + }} + }} +}} +"# + ) + } +} diff --git a/rust/cuda-coop-codegen/src/scan.rs b/rust/cuda-coop-codegen/src/scan.rs index 047187db6169..a4abfdcd687f 100644 --- a/rust/cuda-coop-codegen/src/scan.rs +++ b/rust/cuda-coop-codegen/src/scan.rs @@ -271,6 +271,10 @@ impl ScanRequest { }) } + pub(crate) fn block_shape(self) -> [u32; 3] { + self.block + } + pub(crate) fn source(self, symbol: &str) -> String { let (dtype, _) = self.dtype.native(); let [x, y, z] = self.block; diff --git a/rust/cuda-coop-codegen/tests/native.rs b/rust/cuda-coop-codegen/tests/native.rs index d8ebbd209a13..b7bdb29d4b73 100644 --- a/rust/cuda-coop-codegen/tests/native.rs +++ b/rust/cuda-coop-codegen/tests/native.rs @@ -4,12 +4,238 @@ //! Real native compiler coverage; ordinary host tests need no CUDA installation. use cuda_coop_codegen::{ - Algorithm, Dtype, Layout, NativeRequest, Operation, Provider, ReduceAlgorithm, ReduceForm, - ReduceOperator, ReduceRequest, Request, ScanAlgorithm, ScanOperation, ScanOperator, - ScanRequest, + Algorithm, Dtype, Layout, NativeRequest, Operation, Provider, ReduceAlgorithm, + ReduceBatchedLayout, ReduceBatchedRecipe, ReduceBatchedRequest, 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_warp_shape_identity() -> 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(); + // Two shapes have the same total threads. Warp bodies can otherwise be + // identical even though each root needs a distinct specialization record. + for block in [[64, 1, 1], [128, 1, 1], [8, 4, 2]] { + requests.extend([ + NativeRequest::LoadStore( + Request::new_typed_with_shape( + Dtype::I32, + Operation::LoadMaskedValidity, + Algorithm::Transpose, + block, + 4, + )? + .with_warp_width(32)?, + ), + NativeRequest::Scan(ScanRequest::warp( + Dtype::I32, + ScanOperation::ExclusiveSum, + block, + 32, + )?), + NativeRequest::Reduce(ReduceRequest::warp( + Dtype::I32, + ReduceOperator::Sum, + ReduceForm::Payload, + block, + 32, + 4, + )?), + NativeRequest::ReduceBatched(ReduceBatchedRequest::warp( + Dtype::I32, + ReduceBatchedRecipe::new(ReduceOperator::Sum, ReduceBatchedLayout::Striped), + block, + 32, + 4, + )?), + ]); + } + let artifact = provider.compile_native(&requests, &target)?; + assert_eq!(artifact.specializations.len(), 12); + assert!(!artifact.ltoir.is_empty()); + let mut identities = std::collections::HashSet::new(); + let mut symbols = std::collections::HashSet::new(); + let parameters = |symbol: &str| { + artifact + .source + .split_once(&format!("void {symbol}(")) + .expect("compiled native wrapper") + .1 + .split_once(") {") + .expect("native parameter list") + .0 + .split_whitespace() + .collect::>() + .join(" ") + }; + for (index, entry) in artifact.specializations.iter().enumerate() { + assert!(identities.insert(entry.identity)); + assert!(symbols.insert(&entry.symbol)); + let first = &artifact.specializations[index % 4]; + assert_eq!(entry.cub_layout, first.cub_layout); + assert_eq!(entry.scratch, first.scratch); + assert!(entry.scratch.size > 0); + assert!(!entry.scratch_free && !entry.mutates_payload); + assert_eq!(parameters(&entry.symbol), parameters(&first.symbol)); + eprintln!( + "{:?}: CUB {:?}, scratch {:?}", + requests[index], entry.cub_layout, entry.scratch + ); + } + requests.reverse(); + let reordered = provider.compile_native(&requests, &target)?; + assert_eq!( + reordered.specializations.len(), + artifact.specializations.len() + ); + assert!(!reordered.ltoir.is_empty()); + assert_ne!(reordered.identity, artifact.identity); + for (entry, previous) in reordered + .specializations + .iter() + .zip(artifact.specializations.iter().rev()) + { + assert_eq!(entry.identity, previous.identity); + assert_eq!(entry.symbol, previous.symbol); + assert_eq!(entry.cub_layout, previous.cub_layout); + assert_eq!(entry.scratch, previous.scratch); + assert_eq!(entry.scratch_free, previous.scratch_free); + assert_eq!(entry.mutates_payload, previous.mutates_payload); + } + Ok(()) +} + +#[test] +#[ignore = "requires CUDA_HOME and a Linux CUDA toolkit with NVRTC LTO support"] +fn native_mixed_reduce_batched_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 mut requests = Vec::new(); + for (width, batches, layout) in [ + (32, 1, ReduceBatchedLayout::Striped), + (8, 4, ReduceBatchedLayout::Striped), + (8, 10, ReduceBatchedLayout::Striped), + (8, 10, ReduceBatchedLayout::Blocked), + (1, 9, ReduceBatchedLayout::Blocked), + (8, 65, ReduceBatchedLayout::Striped), + (8, 65, ReduceBatchedLayout::Blocked), + ] { + requests.push(NativeRequest::ReduceBatched(ReduceBatchedRequest::warp( + Dtype::I32, + ReduceBatchedRecipe::new(ReduceOperator::Sum, layout), + [128, 1, 1], + width, + batches, + )?)); + } + // Spread types/operators/widths over cases rather than taking their product. + for (index, (dtype, op)) in [ + (Dtype::I8, ReduceOperator::Sum), + (Dtype::U8, ReduceOperator::Multiplies), + (Dtype::I16, ReduceOperator::Min), + (Dtype::U16, ReduceOperator::Max), + (Dtype::I32, ReduceOperator::BitAnd), + (Dtype::U32, ReduceOperator::BitOr), + (Dtype::I64, ReduceOperator::BitXor), + (Dtype::U64, ReduceOperator::Sum), + (Dtype::F32, ReduceOperator::Min), + (Dtype::F64, ReduceOperator::Max), + ] + .into_iter() + .enumerate() + { + requests.push(NativeRequest::ReduceBatched(ReduceBatchedRequest::warp( + dtype, + ReduceBatchedRecipe::new( + op, + if index % 2 == 0 { + ReduceBatchedLayout::Striped + } else { + ReduceBatchedLayout::Blocked + }, + ), + [4, 4, 4], + [1, 2, 4, 8, 16, 32][index % 6], + [1, 4, 10, 9, 65][index % 5], + )?)); + } + requests.push(NativeRequest::LoadStore(Request::new_typed( + Dtype::I32, + Operation::LoadFullValidity, + Algorithm::Direct, + 64, + 4, + )?)); + requests.push(NativeRequest::Scan(ScanRequest::warp( + Dtype::I32, + ScanOperation::ExclusiveSum, + [4, 4, 4], + 8, + )?)); + requests.push(NativeRequest::Reduce(ReduceRequest::warp( + Dtype::I32, + ReduceOperator::Sum, + ReduceForm::Payload, + [4, 4, 4], + 8, + 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 batched Reduce/Load/Scan/Reduce 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_reduce_recipes() -> Result<(), Box> { diff --git a/rust/cuda-coop/build.rs b/rust/cuda-coop/build.rs index 89a9817518b9..add35460896f 100644 --- a/rust/cuda-coop/build.rs +++ b/rust/cuda-coop/build.rs @@ -1,10 +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 Load, Store, Scan and Reduce. +//! Register the owning host provider for native Load, Store, Scan, Reduce and batched reduction. use cuda_coop_codegen::{ - Algorithm, Operation, ReduceAlgorithm, ReduceForm, ReduceOperator, ReduceRecipe, ScanAlgorithm, - ScanOperation, ScanOperator, ScanRecipe, oxide, + Algorithm, Operation, ReduceAlgorithm, ReduceBatchedLayout, ReduceBatchedRecipe, ReduceForm, + ReduceOperator, ReduceRecipe, ScanAlgorithm, ScanOperation, ScanOperator, ScanRecipe, oxide, }; fn main() { @@ -32,6 +32,17 @@ fn main() { (ReduceOperator::BitOr, "BIT_OR"), (ReduceOperator::BitXor, "BIT_XOR"), ] { + for (layout, suffix) in [ + (ReduceBatchedLayout::Striped, "STRIPED"), + (ReduceBatchedLayout::Blocked, "BLOCKED"), + ] { + let value = oxide::reduce_batched_registration( + &executable, + ReduceBatchedRecipe::new(operator, layout), + ) + .expect("cuda.coop batched reduce registration"); + println!("cargo::rustc-env=COOP_REDUCE_BATCHED_{name}_{suffix}={value}"); + } for (form, suffix) in [ (ReduceForm::Scalar, "SCALAR"), (ReduceForm::ScalarPrefix, "PREFIX"), diff --git a/rust/cuda-coop/examples/reduce_batched/.gitignore b/rust/cuda-coop/examples/reduce_batched/.gitignore new file mode 100644 index 000000000000..74bffa8af4a4 --- /dev/null +++ b/rust/cuda-coop/examples/reduce_batched/.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_batched/Cargo.lock b/rust/cuda-coop/examples/reduce_batched/Cargo.lock new file mode 100644 index 000000000000..4b2388701eea --- /dev/null +++ b/rust/cuda-coop/examples/reduce_batched/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-batched" +version = "0.0.0" +dependencies = [ + "coop-reduce-batched-helper", + "cuda-async", + "cuda-coop", + "cuda-core", + "cuda-device", + "cuda-host", +] + +[[package]] +name = "coop-reduce-batched-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_batched/Cargo.toml b/rust/cuda-coop/examples/reduce_batched/Cargo.toml new file mode 100644 index 000000000000..c92bfe327bb6 --- /dev/null +++ b/rust/cuda-coop/examples/reduce_batched/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-batched" +version = "0.0.0" +edition = "2024" +publish = false + +[workspace] +members = ["helper"] + +[dependencies] +cuda-coop = { path = "../.." } +coop-reduce-batched-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_batched/helper/Cargo.toml b/rust/cuda-coop/examples/reduce_batched/helper/Cargo.toml new file mode 100644 index 000000000000..f634c71dc915 --- /dev/null +++ b/rust/cuda-coop/examples/reduce_batched/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-batched-helper" +version = "0.0.0" +edition = "2024" +publish = false + +[dependencies] +cuda-coop = { path = "../../.." } diff --git a/rust/cuda-coop/examples/reduce_batched/helper/src/lib.rs b/rust/cuda-coop/examples/reduce_batched/helper/src/lib.rs new file mode 100644 index 000000000000..1592740e894f --- /dev/null +++ b/rust/cuda-coop/examples/reduce_batched/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 generic eight-lane batched sum with a calculated output payload. +#![no_std] +#![feature(generic_const_exprs)] +#![allow(incomplete_features)] +use cuda_coop as coop; + +/// Reduce each input slot across eight lanes and store its one striped result. +/// +/// `input_offset` and `output_offset` locate this block's input and padded output +/// regions. Warp Load/Store add each subgroup's origin themselves. `valid` is +/// the initialized source prefix for this subgroup; `zero` fills its short tail +/// before reduction. It must be the additive identity for `T`. +pub fn batches( + group: coop::ThreadGroup>, + input: &[T], + output: &mut coop::DisjointSlice< + '_, + T, + coop::CollectiveTiles<{ coop::batched_extent(ITEMS_PER_THREAD, 8) }>, + >, + input_offset: usize, + output_offset: usize, + valid: usize, + zero: T, +) where + [(); coop::batched_extent(ITEMS_PER_THREAD, 8)]:, +{ + let mut values = coop::ThreadData::<_, ITEMS_PER_THREAD>::new(); + coop::load( + group, + input, + &mut values, + coop::LoadOptions::default() + .offset(input_offset) + .valid_items(valid) + .oob_default(zero), + ); + let totals = coop::reduce_batched(group, &values, Default::default()); + // Every member participates, including lanes with no initialized result. + // Striped Store's count selects exactly the B owned batch results. + coop::store( + group, + output, + &totals, + coop::StoreOptions::default() + .algorithm(coop::Striped) + .offset(output_offset) + .valid_items(ITEMS_PER_THREAD), + ); +} diff --git a/rust/cuda-coop/examples/reduce_batched/src/main.rs b/rust/cuda-coop/examples/reduce_batched/src/main.rs new file mode 100644 index 000000000000..ec27a25dced9 --- /dev/null +++ b/rust/cuda-coop/examples/reduce_batched/src/main.rs @@ -0,0 +1,113 @@ +// Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +#![feature(generic_const_exprs)] +#![allow(incomplete_features)] +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 = (64, 1, 1), dynamic_shared = inferred)] + pub fn batched_sum( + input: &[i32], + mut output: coop::DisjointSlice< + i32, + coop::CollectiveTiles<{ coop::batched_extent(ITEMS_PER_THREAD, 8) }>, + >, + ) where + [(); coop::batched_extent(ITEMS_PER_THREAD, 8)]:, + { + let block = coop::this_block(); + let warp = coop::this_warp(); + let group = warp.group_by::<8>(); + let block_rank = block.rank_at::() as usize; + let threads = block.count() as usize; + let input_offset = block_rank * threads * ITEMS_PER_THREAD; + let output_offset = block_rank * threads * coop::batched_extent(ITEMS_PER_THREAD, 8); + let subgroup = + warp.rank_at::() as usize * 4 + group.rank_at::() as usize; + let begin = input_offset + subgroup * 8 * ITEMS_PER_THREAD; + let remaining = if begin <= input.len() { + input.len() - begin + } else { + 0 + }; + coop_reduce_batched_helper::batches::<_, ITEMS_PER_THREAD>( + group, + input, + &mut output, + input_offset, + output_offset, + remaining.min(8 * ITEMS_PER_THREAD), + 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: typed declarations and embedded kernels compile together. + let module = unsafe { kernels::load_async(0)? }; + macro_rules! run { + ($items:literal) => {{ + const THREADS: usize = 64; + const BLOCKS: usize = 3; + const WIDTH: usize = 8; + const K: usize = coop::batched_extent($items, WIDTH); + let input: Vec = (0..BLOCKS * THREADS) + .flat_map(|thread| { + (0..$items).map(move |batch| { + ((thread / THREADS + 1) * 1000 + (thread % THREADS) * 101 + batch * 7) + as i32 + }) + }) + .take(BLOCKS * THREADS * $items - 7) + .collect(); + let mut expected = vec![-1_i32; BLOCKS * THREADS * K]; + for group in 0..BLOCKS * THREADS / WIDTH { + for batch in 0..$items { + expected[group * WIDTH * K + batch] = (0..WIDTH) + .map(|lane| { + input + .get((group * WIDTH + lane) * $items + batch) + .copied() + .unwrap_or(0) + }) + .sum(); + } + } + let source = DeviceBuffer::from_host(&stream, &input)?; + let target = DeviceBuffer::from_host(&stream, &vec![-1_i32; expected.len()])?; + let plan = module.prepare_batched_sum::<$items>(LaunchConfig1D::auto_shared( + BLOCKS as u32, + THREADS as u32, + ))?; + let (_, target) = module + .batched_sum_async_owned::<$items, _, _>(&plan, source, target) + .sync()?; + let output = target.to_host_vec(&stream)?; + assert_eq!(output, expected); + println!( + "items/thread={}: {} batch sums and {} padding slots checked", + $items, + BLOCKS * THREADS / WIDTH * $items, + expected.len() - BLOCKS * THREADS / WIDTH * $items + ); + }}; + } + run!(1); + run!(4); + run!(10); + Ok(()) +} diff --git a/rust/cuda-coop/src/lib.rs b/rust/cuda-coop/src/lib.rs index 2e86271050c6..7364cdf928f0 100644 --- a/rust/cuda-coop/src/lib.rs +++ b/rust/cuda-coop/src/lib.rs @@ -3,7 +3,7 @@ //! Per-thread payloads and CUDA hierarchy queries for Rust `cuda.coop`. //! -//! This unpublished, feature-free `no_std` crate provides [`ThreadData`] and +//! This unpublished `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. @@ -49,15 +49,23 @@ //! 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. +//! [`reduce_batched`] independently reduces each payload slot across physical or +//! power-of-two logical Warp groups. Its fresh partial payload has a calculated +//! extent and explicit striped or blocked result ownership. The defining crate +//! and callers instantiating these calculated signatures enable unstable +//! `generic_const_exprs`; ordinary non-calculated consumers remain feature-free. //! Host payload tests need no GPU; //! executing hardware queries requires an Oxide-compiled CUDA kernel. #![no_std] +#![feature(generic_const_exprs)] +#![allow(incomplete_features)] mod alignment; mod group; mod load_store; mod reduce; +mod reduce_batched; mod scan; mod temp_storage; mod thread_data; @@ -72,6 +80,10 @@ pub use reduce::{ RakingCommutativeOnly, ReduceAlgorithm, ReduceGroup, ReduceOperator, ReduceOptions, ReduceValue, SumOptions, WarpReductions, reduce, sum, }; +pub use reduce_batched::{ + BatchedOutputLayout, Blocked, ReduceBatchedGroup, ReduceBatchedOptions, batched_extent, + reduce_batched, +}; 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/boundary.rs b/rust/cuda-coop/src/reduce/boundary.rs index 6c52b78e428a..e1fae6677c02 100644 --- a/rust/cuda-coop/src/reduce/boundary.rs +++ b/rust/cuda-coop/src/reduce/boundary.rs @@ -236,6 +236,15 @@ registrations!( (warp_bit_xor_payload, "COOP_REDUCE_WARP_BIT_XOR_PAYLOAD"), ); +struct Selection(core::marker::PhantomData<(Op, C)>); +impl Selection { + const KIND: u8 = if PAYLOAD { + 14 + Op::KIND + } else { + 7 * C::KIND + Op::KIND + }; +} + #[inline(always)] pub(super) fn reduce< T: Numeric, @@ -283,13 +292,7 @@ where // 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 - } - } { + match Selection::::KIND { 0 => call!( raking_commutative_only_sum_scalar, raking_sum_scalar, diff --git a/rust/cuda-coop/src/reduce_batched.rs b/rust/cuda-coop/src/reduce_batched.rs new file mode 100644 index 000000000000..d05b09b99b1d --- /dev/null +++ b/rust/cuda-coop/src/reduce_batched.rs @@ -0,0 +1,192 @@ +// 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, GroupScope, Numeric, ReduceOperator, Striped, Sum, SupportedAlignment, ThreadData, + ThreadGroup, ThreadsWithinWarp, Warp, +}; +use cuda_device::library; + +mod boundary; + +/// Capacity of each lane's batched-reduction result: ceiling division of +/// `batches` by `width`, without overflowing the addition in `(batches + width - 1)`. +/// +/// Both arguments must be positive. This is capacity, not the number of valid +/// slots in every lane. [`reduce_batched`] determines validity from output layout +/// and lane rank. Generic calculated signatures repeat this expression in a +/// `[(); batched_extent(B, W)]:` well-formedness bound. +pub const fn batched_extent(batches: usize, width: usize) -> usize { + assert!(batches > 0); + assert!(width > 0); + batches.div_ceil(width) +} + +/// Consecutive batched results per lane: lane `r`, slot `i` owns batch +/// `r * batched_extent(B, W) + i`. Slots with a batch index at least `B` are invalid. +#[derive(Clone, Copy, Debug, Default)] +pub struct Blocked; + +/// A sealed batched-result layout: [`Striped`] (default) or [`Blocked`]. +/// +/// Striped lane `r`, slot `i` owns batch `r + i * W`; Blocked owns +/// `r * batched_extent(B, W) + i`. This selection describes result ownership, +/// independently of a Load/Store algorithm. +pub trait BatchedOutputLayout: sealed::Layout {} +impl sealed::Layout for Striped { + const KIND: u8 = 0; +} +impl sealed::Layout for Blocked { + const KIND: u8 = 1; +} +impl BatchedOutputLayout for Striped {} +impl BatchedOutputLayout for Blocked {} + +/// A sealed group supporting batched reduction: physical [`Warp`] or logical +/// widths 1, 2, 4, 8, 16 and 32, with either exhaustive selection. +/// +/// The enclosing Block must contain complete physical warps. Block groups and +/// non-power-of-two prefixes are unsupported. All members of each selected +/// group must participate; sibling logical groups may choose different paths. +pub trait ReduceBatchedGroup: GroupScope + sealed::Group { + /// Number of lanes contributing to each batch. + /// + /// In helpers also bounded by [`crate::LoadStoreGroup`], spell this as + /// `::WIDTH` to distinguish the group capabilities. + const WIDTH: usize; +} +macro_rules! warp_group { + ($scope:ty, $width:literal $(, $e:ident)?) => { + impl<$(const $e: bool)?> ReduceBatchedGroup for $scope { const WIDTH: usize = $width; } + impl<$(const $e: bool)?> sealed::Group for $scope { + #[inline(always)] + unsafe fn call &'static str + Copy, P: Copy, const I: usize>( + registration: R, payload: P, + ) { + // SAFETY: the private boundary provides live payloads and the + // matching recipe; the checked marker proves participation. + unsafe { library::checked_warp_payload_call_v1::(registration, (), payload) } + } + } + }; +} +warp_group!(Warp, 32); +warp_group!(ThreadsWithinWarp<1, E>, 1, E); +warp_group!(ThreadsWithinWarp<2, E>, 2, E); +warp_group!(ThreadsWithinWarp<4, E>, 4, E); +warp_group!(ThreadsWithinWarp<8, E>, 8, E); +warp_group!(ThreadsWithinWarp<16, E>, 16, E); +warp_group!(ThreadsWithinWarp<32, E>, 32, E); + +/// Static operator and output layout for [`reduce_batched`]. +/// +/// `T` is inferred from the borrowed input. Canonical [`Default`] selects +/// [`Sum`] and [`Striped`]; builders consume options and preserve the other +/// selection. Built-ins use [`ReduceOperator`]'s numeric and integer bounds. +/// There is no count, seed, algorithm or explicit-storage option. +pub struct ReduceBatchedOptions { + binary_op: Op, + selection: PhantomData<(T, L)>, +} +impl Default for ReduceBatchedOptions { + fn default() -> Self { + Self { + binary_op: Sum, + selection: PhantomData, + } + } +} +impl ReduceBatchedOptions { + /// Select a built-in operator, retaining its actual generic value. + pub fn binary_op>(self, op: Op2) -> ReduceBatchedOptions { + ReduceBatchedOptions { + binary_op: op, + selection: PhantomData, + } + } + /// Select [`Striped`] or [`Blocked`] result ownership. + pub fn output_layout( + self, + _layout: L2, + ) -> ReduceBatchedOptions { + ReduceBatchedOptions { + binary_op: self.binary_op, + selection: PhantomData, + } + } +} + +/// Reduce each input slot independently across a Warp group into fresh partial +/// [`ThreadData`], preserving all input values and validity bits. +/// +/// Each lane supplies `ITEMS_PER_THREAD` initialized numeric slots (batches). +/// The result capacity is `ceil(ITEMS_PER_THREAD / G::WIDTH)`, with default +/// alignment. Each batch has one owning result slot according to the selected +/// [`BatchedOutputLayout`]; absent batches remain invalid and `get` returns None. +/// Input may be shared or mutably reborrowed; scalar input is unsupported. +/// +/// All selected members participate. Uninitialized input traps before native +/// numeric reads. Compiler-planned implicit Warp scratch is synchronized before +/// reuse; the operation does not imply global-memory completion. Use the matching +/// group's `sync()` when a later global access depends on preceding accesses. +/// +/// This crate and each crate instantiating the calculated result signature use +/// unstable `generic_const_exprs`. Ordinary payload/Load/Scan/Reduce consumers do +/// not need that feature. Generic helpers carry the explicit well-formedness +/// bound, as below; no caller-chosen output extent is needed. +/// +/// ```no_run +/// #![feature(generic_const_exprs)] +/// #![allow(incomplete_features)] +/// use cuda_coop as coop; +/// fn batches( +/// group: coop::ThreadGroup, values: &coop::ThreadData, +/// ) -> coop::ThreadData +/// where [(); coop::batched_extent(ITEMS_PER_THREAD, G::WIDTH)]:, +/// { +/// coop::reduce_batched(group, values, Default::default()) +/// } +/// ``` +#[inline(always)] +pub fn reduce_batched( + _group: ThreadGroup, + input: &ThreadData, + options: ReduceBatchedOptions, +) -> ThreadData::WIDTH) }> +where + T: Numeric, + G: ReduceBatchedGroup, + Op: ReduceOperator, + L: BatchedOutputLayout, + Alignment: SupportedAlignment, + [(); batched_extent(ITEMS_PER_THREAD, ::WIDTH)]:, +{ + boundary::reduce::(input, options) +} + +mod sealed { + pub trait Layout { + const KIND: u8; + } + pub trait Group { + unsafe fn call &'static str + Copy, P: Copy, const I: usize>( + registration: R, + payload: P, + ); + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn extent_handles_partial_and_maximum_batches() { + assert_eq!(batched_extent(1, 32), 1); + assert_eq!(batched_extent(10, 8), 2); + assert_eq!(batched_extent(usize::MAX, 1), usize::MAX); + assert_eq!(batched_extent(usize::MAX, 32), usize::MAX / 32 + 1); + } +} diff --git a/rust/cuda-coop/src/reduce_batched/boundary.rs b/rust/cuda-coop/src/reduce_batched/boundary.rs new file mode 100644 index 000000000000..a62f11657d88 --- /dev/null +++ b/rust/cuda-coop/src/reduce_batched/boundary.rs @@ -0,0 +1,80 @@ +// Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +use super::{BatchedOutputLayout, ReduceBatchedGroup, ReduceBatchedOptions, batched_extent}; +use crate::{Alignment, Numeric, ReduceOperator, SupportedAlignment, ThreadData}; +use cuda_device::library; + +macro_rules! registrations { + ($(($name:ident, $variable:literal)),+ $(,)?) => {$ ( + fn $name() -> &'static str { library::registration!($variable) } + )+}; +} +registrations!( + (sum_striped, "COOP_REDUCE_BATCHED_SUM_STRIPED"), + (multiplies_striped, "COOP_REDUCE_BATCHED_MULTIPLIES_STRIPED"), + (min_striped, "COOP_REDUCE_BATCHED_MIN_STRIPED"), + (max_striped, "COOP_REDUCE_BATCHED_MAX_STRIPED"), + (bit_and_striped, "COOP_REDUCE_BATCHED_BIT_AND_STRIPED"), + (bit_or_striped, "COOP_REDUCE_BATCHED_BIT_OR_STRIPED"), + (bit_xor_striped, "COOP_REDUCE_BATCHED_BIT_XOR_STRIPED"), + (sum_blocked, "COOP_REDUCE_BATCHED_SUM_BLOCKED"), + (multiplies_blocked, "COOP_REDUCE_BATCHED_MULTIPLIES_BLOCKED"), + (min_blocked, "COOP_REDUCE_BATCHED_MIN_BLOCKED"), + (max_blocked, "COOP_REDUCE_BATCHED_MAX_BLOCKED"), + (bit_and_blocked, "COOP_REDUCE_BATCHED_BIT_AND_BLOCKED"), + (bit_or_blocked, "COOP_REDUCE_BATCHED_BIT_OR_BLOCKED"), + (bit_xor_blocked, "COOP_REDUCE_BATCHED_BIT_XOR_BLOCKED"), +); + +struct Selection(core::marker::PhantomData<(T, Op, L)>); +impl, L: BatchedOutputLayout> Selection { + const KIND: u8 = 7 * L::KIND + Op::KIND; +} + +#[inline(always)] +pub(super) fn reduce( + input: &ThreadData, + options: ReduceBatchedOptions, +) -> ThreadData +where + T: Numeric, + G: ReduceBatchedGroup, + Op: ReduceOperator, + L: BatchedOutputLayout, + Alignment: SupportedAlignment, + [(); batched_extent(ITEMS, G::WIDTH)]:, +{ + 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 _operator = options.binary_op; + // SAFETY: input and fresh output are distinct live lane-local buffers with + // their actual extents and alignments. The typed recipe derives the same + // output extent from input extent/width, preflights every input bit before + // numeric reads, and preserves input. It clears output bits and copies only + // owned initialized results, then marks their bits. No invalid CUB output + // is read. Pointers cannot escape the synchronous call. The genuine marker + // checks participation and compiler-planned scratch allocation/reuse. + unsafe { + match Selection::::KIND { + 0 => G::call::(sum_striped, payload), + 1 => G::call::(multiplies_striped, payload), + 2 => G::call::(min_striped, payload), + 3 => G::call::(max_striped, payload), + 4 => G::call::(bit_and_striped, payload), + 5 => G::call::(bit_or_striped, payload), + 6 => G::call::(bit_xor_striped, payload), + 7 => G::call::(sum_blocked, payload), + 8 => G::call::(multiplies_blocked, payload), + 9 => G::call::(min_blocked, payload), + 10 => G::call::(max_blocked, payload), + 11 => G::call::(bit_and_blocked, payload), + 12 => G::call::(bit_or_blocked, payload), + 13 => G::call::(bit_xor_blocked, payload), + _ => panic!(), + } + } + output +} diff --git a/rust/cuda-coop/tests/check_load_store.py b/rust/cuda-coop/tests/check_load_store.py index 60c3d0d6643f..5794576d0cdc 100644 --- a/rust/cuda-coop/tests/check_load_store.py +++ b/rust/cuda-coop/tests/check_load_store.py @@ -3,9 +3,9 @@ """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 -can supply explicit Cargo patches without changing checked-in manifests. +Use the qualified installed Oxide CLI and its matching automatic backend, or +supply an explicit --backend for a labelled local compiler candidate. The optional +Cargo config supports local source integration without changing manifests. """ import argparse @@ -17,7 +17,11 @@ 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( + "--backend", + type=Path, + help="Explicit candidate backend; omitted uses the installed CLI cache", +) parser.add_argument("--output", required=True, type=Path) parser.add_argument("--cargo-config", type=Path) parser.add_argument( @@ -1067,6 +1071,84 @@ ), ) +# Batched reductions return a calculated partial payload. Result validity is +# lane-local and does not replace participation or global-access completion. +batched_group = "let group=coop::this_warp().group_by::<8>();" +batched_call = "let _result=coop::reduce_batched(group,&values,Default::default());" +cases.update( + batched_generic_defaults=( + load + "fn forward(" + "group:coop::ThreadGroup,input:&coop::ThreadData)" + "->coop::ThreadData " + "where [();coop::batched_extent(B,G::WIDTH)]: {" + "coop::reduce_batched(group,input,Default::default())}" + "fn selected,const B:usize>(" + "group:coop::ThreadGroup,input:&coop::ThreadData,op:Op)" + "->coop::ThreadData " + "where [();coop::batched_extent(B,G::WIDTH)]: {" + "coop::reduce_batched(group,input,coop::ReduceBatchedOptions::default().binary_op(op).output_layout(coop::Blocked))}" + "let _result=forward(coop::this_warp().group_by::<8>(),&mut values);" + "let _result=selected(coop::this_warp(),&values,coop::Min);", + None, + ), + batched_scalar=( + load + + "let _result=coop::reduce_batched(coop::this_warp(),&1_i32,Default::default());", + "mismatched types", + ), + batched_block=( + load + "let _result=coop::reduce_batched(block,&values,Default::default());", + "ReduceBatchedGroup", + ), + batched_prefix17=( + load + "let group=coop::this_warp().group_by_non_exhaustive::<17>();" + "if group.is_member()!=0 {" + batched_call + "}", + "ReduceBatchedGroup", + ), + batched_lane_divergence=( + load + batched_group + "if thread::threadIdx_x()&1==0 {" + batched_call + "}", + "lane-varying", + ), + batched_read_epoch=( + reload + batched_group + batched_call + store, + "requires a barrier covering both accesses", + ), + batched_read_completed=( + reload + batched_group + batched_call + sync + store, + None, + ), + batched_result_is_not_participation=( + load + + batched_group + + "let result=coop::reduce_batched(group,&values,Default::default());" + "if result.get(0).is_some() {" + batched_call + "}", + "lane-varying", + ), + batched_count=( + load + "let _result=coop::reduce_batched(coop::this_warp(),&values," + "coop::ReduceBatchedOptions::default().valid_items(valid));", + "no method named `valid_items`", + ), + batched_storage=( + load + + new_storage + + "let _result=coop::reduce_batched(coop::this_warp(),&values," + "coop::ReduceBatchedOptions::default().temp_storage(&mut scratch));", + "no method named `temp_storage`", + ), + batched_algorithm=( + load + "let _result=coop::reduce_batched(coop::this_warp(),&values," + "coop::ReduceBatchedOptions::default().algorithm(coop::Raking));", + "no method named `algorithm`", + ), + batched_float_bitwise=( + "values.write(0,0_i32);let mut floats=coop::ThreadData::<_,4>::new();" + "floats.write(0,1_f32);let _result=coop::reduce_batched(coop::this_warp(),&floats," + "coop::ReduceBatchedOptions::default().binary_op(coop::BitAnd));", + "ReduceOperator", + ), +) + if args.case: unknown = set(args.case).difference(cases) if unknown: @@ -1088,16 +1170,24 @@ 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()), ) +if args.backend: + env["CUDA_OXIDE_BACKEND"] = str(args.backend.resolve()) +else: + env.pop("CUDA_OXIDE_BACKEND", None) 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.startswith("batched_"): + case_prefix = case_prefix.replace( + "use cuda_coop as coop;", + "#![feature(generic_const_exprs)]\n#![allow(incomplete_features)]\nuse cuda_coop as coop;", + ) if name == "logical_width1_origin": case_prefix = case_prefix.replace( "valid:usize)", @@ -1137,7 +1227,11 @@ print(name, "PASS" if passed else "FAIL", flush=True) (output / "results.json").write_text( json.dumps( - {"command": command, "backend": env["CUDA_OXIDE_BACKEND"], "cases": results}, + { + "command": command, + "backend": env.get("CUDA_OXIDE_BACKEND"), + "cases": results, + }, indent=2, ) + "\n" diff --git a/rust/cuda-coop/tests/reduce_batched/.gitignore b/rust/cuda-coop/tests/reduce_batched/.gitignore new file mode 100644 index 000000000000..74bffa8af4a4 --- /dev/null +++ b/rust/cuda-coop/tests/reduce_batched/.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_batched/Cargo.lock b/rust/cuda-coop/tests/reduce_batched/Cargo.lock new file mode 100644 index 000000000000..99b88de73ef3 --- /dev/null +++ b/rust/cuda-coop/tests/reduce_batched/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-batched-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_batched/Cargo.toml b/rust/cuda-coop/tests/reduce_batched/Cargo.toml new file mode 100644 index 000000000000..1bd14cc0ec20 --- /dev/null +++ b/rust/cuda-coop/tests/reduce_batched/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-batched-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_batched/README.md b/rust/cuda-coop/tests/reduce_batched/README.md new file mode 100644 index 000000000000..97d6ce55bcf4 --- /dev/null +++ b/rust/cuda-coop/tests/reduce_batched/README.md @@ -0,0 +1,30 @@ +# Public batched Reduce fixture + +This standalone test uses the public `cuda_coop` API and the pinned Oxide +toolchain in `Cargo.toml`. It enables unstable `generic_const_exprs` for the +operation-derived result extent; callers do not supply an output size. + +The positive mode runs 14 cases and checks 1,792 per-lane reports. Six typed +shape rows exercise both Striped and Blocked results: physical W32/B1 and B4 +i32 Sum, W8/B10 i8 Min and f64 Max, W1/B9 u16 Sum, and W8/B65 u32 BitXor. +Logical rows use XYZ blocks. Two additional i8 cases select alternating W8 +siblings with differing group-uniform iteration counts. + +Every result slot is checked against a host-computed numerical and validity +oracle, including absent slots, extent/OOB access, unchanged input, and fresh +results after mutating an earlier result. A full-block Direct Store records +one observation per lane after all Warp work reconverges. The test does not +inspect private bitmap bytes or read absent numeric values. + +Build and run from this directory with the matching installed toolchain: + +```sh +cargo oxide build --arch sm_120 --materialize-cubin -- --release --locked +cargo oxide run --arch sm_120 --materialize-cubin +``` + +The built executable also accepts `invalid-input`. Run that mode in a separate +process: it leaves B4 slot 3 unwritten in physical lane 31, which owns no output. +An actual CUDA launch failure is expected; successful return or an unrelated +host error does not qualify the trap. Run positive sanitizer checks separately +from this deliberate failure. diff --git a/rust/cuda-coop/tests/reduce_batched/src/main.rs b/rust/cuda-coop/tests/reduce_batched/src/main.rs new file mode 100644 index 000000000000..ad316a023e35 --- /dev/null +++ b/rust/cuda-coop/tests/reduce_batched/src/main.rs @@ -0,0 +1,418 @@ +// Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +#![feature(generic_const_exprs)] +#![allow(incomplete_features)] + +use cuda_coop as coop; +use cuda_device::{kernel, launch_contract, thread}; +use cuda_host::cuda_module; + +const SKIPPED: u32 = 1 << 31; +const CHECKS: u32 = 15 << 16; + +#[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 < N { + matches &= values.get(i) == before.get(i); + i += 1; + } + matches && values.get(N).is_none() && values.get(usize::MAX).is_none() + } + + fn observe( + values: &coop::ThreadData, + expected: &[T], + offset: usize, + lane: usize, + width: usize, + batches: usize, + blocked: bool, + ) -> (u32, bool) { + let mut mask = 0; + let mut matches = values.len() == K; + let mut i = 0; + while i < K { + let batch = if blocked { + lane * K + i + } else { + lane + i * width + }; + let value = values.get(i); + matches &= value.is_some() == (batch < batches); + if let Some(value) = value { + mask |= 1 << i; + matches &= value == expected[offset + i]; + } + i += 1; + } + ( + mask, + matches && values.get(K).is_none() && values.get(usize::MAX).is_none(), + ) + } + + fn invoke( + group: coop::ThreadGroup, + input: &coop::ThreadData, + blocked: bool, + op: Op, + ) -> coop::ThreadData + where + T: coop::Numeric, + G: coop::ReduceBatchedGroup, + Op: coop::ReduceOperator, + [(); coop::batched_extent(ITEMS_PER_THREAD, G::WIDTH)]:, + { + let options = coop::ReduceBatchedOptions::default().binary_op(op); + if blocked { + coop::reduce_batched(group, input, options.output_layout(coop::Blocked)) + } else { + coop::reduce_batched(group, input, options) + } + } + + #[allow(clippy::too_many_arguments)] + fn exercise( + group: coop::ThreadGroup, + input: &[T], + expected: &[T], + tid: usize, + subgroup: usize, + blocked: bool, + selective: bool, + invalid: bool, + op: Op, + ) -> u32 + where + T: coop::Numeric + PartialEq, + G: coop::ReduceBatchedGroup, + Op: coop::ReduceOperator + Copy, + [(); coop::batched_extent(ITEMS_PER_THREAD, G::WIDTH)]:, + { + let mut values = coop::ThreadData::<_, ITEMS_PER_THREAD>::new(); + let mut i = 0; + while i < ITEMS_PER_THREAD { + // The isolated trap leaves one non-output-owner's input unwritten. + if !(invalid && tid == 31 && i == 3) { + values.write(i, input[tid * ITEMS_PER_THREAD + i]); + } + i += 1; + } + let before = values; + let mut word = SKIPPED; + if !selective || subgroup & 1 == 0 { + let capacity = coop::batched_extent(ITEMS_PER_THREAD, G::WIDTH); + let lane = tid % G::WIDTH; + let iterations = if selective { + 1 + ((subgroup / 2) & 1) + } else { + 2 + }; + let mut values_match = true; + let mut input_preserved = true; + let mut independent = true; + let mut masks_match = true; + let mut mask = 0; + let mut iteration = 0; + while iteration < iterations { + let mut first = invoke(group, &values, blocked, op); + let (first_mask, first_ok) = observe( + &first, + expected, + tid * capacity, + lane, + G::WIDTH, + ITEMS_PER_THREAD, + blocked, + ); + values_match &= first_ok; + input_preserved &= preserved(&values, &before); + // This also initializes the last slot on lanes where it was absent. + first.write(capacity - 1, input[tid * ITEMS_PER_THREAD]); + let mutated = first; + let second = invoke(group, &values, blocked, op); + let (second_mask, second_ok) = observe( + &second, + expected, + tid * capacity, + lane, + G::WIDTH, + ITEMS_PER_THREAD, + blocked, + ); + values_match &= second_ok; + masks_match &= first_mask == second_mask; + independent &= preserved(&first, &mutated); + input_preserved &= preserved(&values, &before); + mask = second_mask; + iteration += 1; + } + word = mask + | ((values_match as u32) << 16) + | ((input_preserved as u32) << 17) + | ((independent as u32) << 18) + | ((masks_match as u32) << 19); + } + word + } + + fn report(output: &mut coop::DisjointSlice>, word: u32) { + let block = coop::this_block(); + let mut value = coop::ThreadData::<_, 1>::new(); + value.write(0, word); + // All Warp branches have reconverged; Block Store owns the XYZ output. + coop::store( + block, + output, + &value, + coop::StoreOptions::default().offset(block.rank_at::() as usize * 64), + ); + } + + #[kernel] + #[launch_contract(domain = 1, block = (64, 1, 1), dynamic_shared = inferred)] + pub fn physical( + input: &[i32], + expected: &[i32], + mut output: coop::DisjointSlice>, + blocked: bool, + invalid: bool, + ) where + [(); coop::batched_extent(ITEMS_PER_THREAD, 32)]:, + { + let warp = coop::this_warp(); + let tid = (thread::blockIdx_x() * 64 + thread::threadIdx_x()) as usize; + let word = exercise::<_, _, _, ITEMS_PER_THREAD>( + warp, + input, + expected, + tid, + warp.rank_at::() as usize, + blocked, + false, + invalid, + coop::Sum, + ); + report(&mut output, word); + } + + #[kernel] + #[launch_contract(domain = 3, block = (4, 4, 4), dynamic_shared = inferred)] + pub fn logical( + input: &[T], + expected: &[T], + mut output: coop::DisjointSlice>, + blocked: bool, + selective: bool, + ) where + T: coop::Numeric + PartialEq, + Op: coop::ReduceOperator + Default + Copy, + coop::ThreadsWithinWarp: coop::ReduceBatchedGroup, + [(); coop::batched_extent( + ITEMS_PER_THREAD, + as coop::ReduceBatchedGroup>::WIDTH, + )]:, + { + let group = coop::this_warp().group_by::(); + let rank = thread::threadIdx_x() + 4 * (thread::threadIdx_y() + 4 * thread::threadIdx_z()); + let tid = (thread::blockIdx_x() * 64 + rank) as usize; + let word = exercise::<_, _, _, ITEMS_PER_THREAD>( + group, + input, + expected, + tid, + coop::this_warp().rank_at::() as usize * (32 / WIDTH) + + group.rank_at::() as usize, + blocked, + selective, + false, + Op::default(), + ); + report(&mut output, word); + } +} + +fn oracle( + input: &[T], + batches: usize, + width: usize, + blocked: bool, + selective: bool, + reduce: impl Fn(T, T) -> T, +) -> (Vec, Vec) { + let capacity = batches.div_ceil(width); + let mut expected = vec![input[0]; 128 * capacity]; + let mut words = vec![CHECKS; 128]; + for tid in 0..128 { + let lane = tid % width; + let group = tid / width; + if selective && group & 1 != 0 { + words[tid] = SKIPPED; + } + for i in 0..capacity { + let batch = if blocked { + lane * capacity + i + } else { + lane + i * width + }; + if batch < batches { + expected[tid * capacity + i] = (0..width) + .map(|lane| input[(group * width + lane) * batches + batch]) + .reduce(&reduce) + .unwrap(); + if words[tid] != SKIPPED { + words[tid] |= 1 << i; + } + } + } + } + (expected, words) +} + +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: typed declarations and embedded kernels are compiled together. + let module = unsafe { kernels::load_async(0)? }; + let invalid = match std::env::args().nth(1).as_deref() { + None => false, + Some("invalid-input") => true, + Some(other) => panic!("unknown test mode {other}"), + }; + if invalid { + let input = DeviceBuffer::from_host(&stream, &[1_i32; 128 * 4])?; + let expected = DeviceBuffer::from_host(&stream, &[32_i32; 128])?; + let output = DeviceBuffer::from_host(&stream, &[0_u32; 128])?; + let plan = module.prepare_physical::<4>(LaunchConfig1D::auto_shared(2, 64))?; + module + .physical_async_owned::<4, _, _, _>(&plan, input, expected, output, false, true) + .sync()?; + stream.synchronize()?; + panic!("unwritten batched input unexpectedly completed"); + } + let mut cases = 0; + macro_rules! physical { + ($items:literal) => {{ + let input: Vec = (0..128) + .flat_map(|tid| { + (0..$items).map(move |batch| { + (10000 * (tid / 64) + 1000 * (tid / 32) + 17 * batch + tid % 32) as i32 + }) + }) + .collect(); + for blocked in [false, true] { + let (expected, words) = oracle(&input, $items, 32, blocked, false, |a, b| a + b); + let source = DeviceBuffer::from_host(&stream, &input)?; + let expected = DeviceBuffer::from_host(&stream, &expected)?; + let target = DeviceBuffer::from_host(&stream, &[0_u32; 128])?; + let plan = module.prepare_physical::<$items>(LaunchConfig1D::auto_shared(2, 64))?; + let (_, _, target) = module + .physical_async_owned::<$items, _, _, _>( + &plan, source, expected, target, blocked, false, + ) + .sync()?; + assert_eq!( + target.to_host_vec(&stream)?, + words, + "physical B{} blocked{}", + $items, + blocked + ); + cases += 1; + } + }}; + } + physical!(1); + physical!(4); + macro_rules! logical { + ($ty:ty, $op:ty, $items:literal, $width:literal, $value:expr, $reduce:expr, $selective:expr) => {{ + let value = $value; + let input: Vec<$ty> = (0..128) + .flat_map(|tid| (0..$items).map(move |batch| value(tid, batch))) + .collect(); + for blocked in [false, true] { + let (expected, words) = + oracle(&input, $items, $width, blocked, $selective, $reduce); + let source = DeviceBuffer::from_host(&stream, &input)?; + let expected = DeviceBuffer::from_host(&stream, &expected)?; + let target = DeviceBuffer::from_host(&stream, &[0_u32; 128])?; + let plan = module.prepare_logical::<$ty, $op, $items, $width>( + LaunchConfig3D::auto_shared((2, 1, 1), (4, 4, 4)), + )?; + let (_, _, target) = module + .logical_async_owned::<$ty, $op, $items, $width, _, _, _>( + &plan, source, expected, target, blocked, $selective, + ) + .sync()?; + assert_eq!( + target.to_host_vec(&stream)?, + words, + "{} B{} W{} blocked{} selective{}", + stringify!($ty), + $items, + $width, + blocked, + $selective + ); + cases += 1; + } + }}; + } + for selective in [false, true] { + logical!( + i8, + coop::Min, + 10, + 8, + |tid: usize, batch: usize| -112_i8 + (6 * (tid / 8) + batch + tid % 8) as i8, + i8::min, + selective + ); + } + logical!( + f64, + coop::Max, + 10, + 8, + |tid: usize, batch: usize| (1000 * (tid / 64) + 100 * (tid / 8) + 7 * batch) as f64 + + (tid % 8) as f64 / 8.0, + f64::max, + false + ); + logical!( + u16, + coop::Sum, + 9, + 1, + |tid: usize, batch: usize| (tid * 17 + batch + 1) as u16, + |a, b| a + b, + false + ); + logical!( + u32, + coop::BitXor, + 65, + 8, + |tid: usize, batch: usize| 1_u32 << ((batch + 3 * (tid % 8) + tid / 8) % 31), + |a, b| a ^ b, + false + ); + println!( + "PASS: {cases} public batched Reduce cases, {} lane reports", + cases * 128 + ); + Ok(()) +} From 391c7a60061518965d66f1c972def30aea76e3a5 Mon Sep 17 00:00:00 2001 From: Trent Nelson Date: Sat, 24 Oct 2026 17:42:24 -0700 Subject: [PATCH 06/10] Add Rust cooperative Exchange and Shuffle Expose complete Block and Warp exchange operations and Block shifts with explicit layout, output validity and scratch completion. Signed-off-by: Trent Nelson --- docs/rust/exchange.rst | 114 ++ docs/rust/index.rst | 16 +- docs/rust/shuffle.rst | 90 ++ rust/cuda-coop-codegen/src/exchange.rs | 255 +++++ rust/cuda-coop-codegen/src/lib.rs | 42 +- rust/cuda-coop-codegen/src/oxide.rs | 203 +++- rust/cuda-coop-codegen/src/shuffle.rs | 179 +++ rust/cuda-coop-codegen/tests/native.rs | 236 +++- rust/cuda-coop/build.rs | 26 +- rust/cuda-coop/examples/exchange/.gitignore | 7 + rust/cuda-coop/examples/exchange/Cargo.lock | 1019 +++++++++++++++++ rust/cuda-coop/examples/exchange/Cargo.toml | 19 + .../examples/exchange/helper/Cargo.toml | 11 + .../examples/exchange/helper/src/lib.rs | 40 + rust/cuda-coop/examples/exchange/src/main.rs | 69 ++ rust/cuda-coop/examples/shuffle/.gitignore | 7 + rust/cuda-coop/examples/shuffle/Cargo.lock | 1019 +++++++++++++++++ rust/cuda-coop/examples/shuffle/Cargo.toml | 19 + .../examples/shuffle/helper/Cargo.toml | 11 + .../examples/shuffle/helper/src/lib.rs | 41 + rust/cuda-coop/examples/shuffle/src/main.rs | 99 ++ rust/cuda-coop/src/exchange.rs | 197 ++++ rust/cuda-coop/src/exchange/boundary.rs | 69 ++ rust/cuda-coop/src/lib.rs | 11 + rust/cuda-coop/src/shuffle.rs | 139 +++ rust/cuda-coop/src/shuffle/boundary.rs | 51 + rust/cuda-coop/tests/check_load_store.py | 169 +++ rust/cuda-coop/tests/exchange/.gitignore | 7 + rust/cuda-coop/tests/exchange/Cargo.lock | 1011 ++++++++++++++++ rust/cuda-coop/tests/exchange/Cargo.toml | 17 + rust/cuda-coop/tests/exchange/README.md | 29 + rust/cuda-coop/tests/exchange/src/main.rs | 401 +++++++ rust/cuda-coop/tests/shuffle/.gitignore | 7 + rust/cuda-coop/tests/shuffle/Cargo.lock | 1011 ++++++++++++++++ rust/cuda-coop/tests/shuffle/Cargo.toml | 17 + rust/cuda-coop/tests/shuffle/src/main.rs | 362 ++++++ 36 files changed, 6998 insertions(+), 22 deletions(-) create mode 100644 docs/rust/exchange.rst create mode 100644 docs/rust/shuffle.rst create mode 100644 rust/cuda-coop-codegen/src/exchange.rs create mode 100644 rust/cuda-coop-codegen/src/shuffle.rs create mode 100644 rust/cuda-coop/examples/exchange/.gitignore create mode 100644 rust/cuda-coop/examples/exchange/Cargo.lock create mode 100644 rust/cuda-coop/examples/exchange/Cargo.toml create mode 100644 rust/cuda-coop/examples/exchange/helper/Cargo.toml create mode 100644 rust/cuda-coop/examples/exchange/helper/src/lib.rs create mode 100644 rust/cuda-coop/examples/exchange/src/main.rs create mode 100644 rust/cuda-coop/examples/shuffle/.gitignore create mode 100644 rust/cuda-coop/examples/shuffle/Cargo.lock create mode 100644 rust/cuda-coop/examples/shuffle/Cargo.toml create mode 100644 rust/cuda-coop/examples/shuffle/helper/Cargo.toml create mode 100644 rust/cuda-coop/examples/shuffle/helper/src/lib.rs create mode 100644 rust/cuda-coop/examples/shuffle/src/main.rs create mode 100644 rust/cuda-coop/src/exchange.rs create mode 100644 rust/cuda-coop/src/exchange/boundary.rs create mode 100644 rust/cuda-coop/src/shuffle.rs create mode 100644 rust/cuda-coop/src/shuffle/boundary.rs create mode 100644 rust/cuda-coop/tests/exchange/.gitignore create mode 100644 rust/cuda-coop/tests/exchange/Cargo.lock create mode 100644 rust/cuda-coop/tests/exchange/Cargo.toml create mode 100644 rust/cuda-coop/tests/exchange/README.md create mode 100644 rust/cuda-coop/tests/exchange/src/main.rs create mode 100644 rust/cuda-coop/tests/shuffle/.gitignore create mode 100644 rust/cuda-coop/tests/shuffle/Cargo.lock create mode 100644 rust/cuda-coop/tests/shuffle/Cargo.toml create mode 100644 rust/cuda-coop/tests/shuffle/src/main.rs diff --git a/docs/rust/exchange.rst b/docs/rust/exchange.rst new file mode 100644 index 000000000000..8fa805cbef3b --- /dev/null +++ b/docs/rust/exchange.rst @@ -0,0 +1,114 @@ +.. Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. +.. SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +Exchange layouts +================ + +``exchange`` redistributes a group's per-thread values between striped and +blocked layouts. It preserves the input and returns a fresh ``ThreadData`` +with the same element type, item count and alignment. The default converts +striped input to blocked output: + +.. code-block:: rust + + let blocked = coop::exchange(group, &striped, Default::default()); + let striped_again = coop::exchange(group, &blocked, + coop::ExchangeOptions::default().mode(coop::BlockedToStriped)); + +Blocked order assigns consecutive items to each thread. Striped order assigns +consecutive items to consecutive threads. For four threads with two items each, +``StripedToBlocked`` produces: + +.. list-table:: + :header-rows: 1 + + * - Thread rank + - Striped input + - Blocked output + * - 0 + - ``[0, 4]`` + - ``[0, 1]`` + * - 1 + - ``[1, 5]`` + - ``[2, 3]`` + * - 2 + - ``[2, 6]`` + - ``[4, 5]`` + * - 3 + - ``[3, 7]`` + - ``[6, 7]`` + +``BlockedToStriped`` reverses this conversion. Exchange operates on the values +already held by the threads; Load and Store handle global-memory access. + +A complete program +------------------ + +This standalone Cargo example loads a tile in striped order, exchanges it to +blocked order, then uses Direct Store. Its reusable helper lives in a separate +crate. The element type is inferred, and ``ITEMS_PER_THREAD`` stays generic +through the host launch, kernel and helper. These same-extent signatures need +no ``generic_const_exprs`` feature in the consumer. + +.. literalinclude:: ../../rust/cuda-coop/examples/exchange/helper/src/lib.rs + :language: rust + :start-after: // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +The host specializes one and four items per thread across three blocks. A +filled Load initializes the short final tile, and Store writes only its valid +prefix. The resulting global output matches the input even though the +per-thread layout changes between Load and Store. + +.. literalinclude:: ../../rust/cuda-coop/examples/exchange/src/main.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/exchange + 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. ``auto_shared`` uses the compiler's +inferred storage plan; the consumer does not calculate scratch bytes. + +Supported groups and inputs +--------------------------- + +Exchange supports complete one-, two- and three-dimensional Blocks, physical +Warps and logical Warp groups of widths 1, 2, 4, 8, 16 and 32. Block thread rank +is x-fastest. Block size need not be a multiple of 32; Warp operations require +complete physical warps in the enclosing Block. Each Warp group converts its +own values independently. + +The payload has a positive item count and one of the ten ``Numeric`` element +types: ``i8``, ``u8``, ``i16``, ``u16``, ``i32``, ``u32``, ``i64``, ``u64``, +``f32`` or ``f64``. Pass a shared borrow, or an ordinary mutable reborrow. +Every input slot must be initialized, including a one-item payload. A masked +Load without a default can leave unwritten slots; Exchange traps before +reading them. Every returned slot is initialized, and changing the result +does not change the input. + +``ExchangeOptions::default()`` selects ``StripedToBlocked``. The consuming +``mode`` builder selects either direction at compile time. The API supports +these two payload conversions; scalar inputs, ranked scatter, algorithm and +explicit-storage options are unsupported. + +Participation and storage +------------------------- + +Every selected member must reach the same Exchange. Separate Warp groups may +take different paths, but members of one group must participate together. +Oxide checks participation and plans the actual CUB shared storage, including +independent regions for Warp groups and synchronization before reuse. + +Exchange does not promise completion of earlier global-memory accesses. Use +the group's ``sync()`` when subsequent conflicting accesses need completion. +See :doc:`groups` and :doc:`load_store` for the participation and memory rules. + +Generate the item reference with ``cargo doc --workspace --no-deps`` from +``rust/``. ``exchange``, ``ExchangeOptions``, ``ExchangeGroup`` and +``ExchangeMode`` document the full generic signatures and supported selectors. diff --git a/docs/rust/index.rst b/docs/rust/index.rst index 42e25b264a24..f68843d5f71a 100644 --- a/docs/rust/index.rst +++ b/docs/rust/index.rst @@ -43,13 +43,15 @@ 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`, :doc:`scan`, :doc:`reduce` -and :doc:`reduce_batched`. Scan returns fresh per-group prefixes; Reduce -returns a total only at each group's root. Batched reduction returns a fresh +The public operations include :doc:`load_store`, :doc:`scan`, :doc:`reduce`, +:doc:`reduce_batched`, :doc:`exchange` and :doc:`shuffle`. Scan returns fresh per-group prefixes; +Reduce returns a total only at each group's root. Batched reduction returns a fresh payload whose extent and valid slots follow its Warp width and output layout. -Scan, Reduce and batched reduction preserve their input. Additional primitive -families are under development. The crate is unpublished and its interface can -change. +Exchange converts between blocked and striped per-thread layouts. +Shuffle shifts blocked values by one item, leaving the exposed edge invalid. +Scan, Reduce, batched reduction, Exchange and Shuffle preserve their input. Additional +primitive families are under development. The crate is unpublished and its +interface can change. Build and run ------------- @@ -119,6 +121,8 @@ Read next scan reduce reduce_batched + exchange + shuffle provider Generate the API reference with ``cargo doc --workspace --no-deps`` from diff --git a/docs/rust/shuffle.rst b/docs/rust/shuffle.rst new file mode 100644 index 000000000000..f056f1120d96 --- /dev/null +++ b/docs/rust/shuffle.rst @@ -0,0 +1,90 @@ +.. Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. +.. SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +Shift a block's values +======================= + +``shuffle`` shifts a block's values by one item in blocked order. It preserves +the input and returns a fresh ``ThreadData`` with the same element type, item +count and alignment. The default direction, ``Down``, reads the next item; +``Up`` reads the previous item: + +.. code-block:: rust + + let next = coop::shuffle(block, &values, Default::default()); + let previous = coop::shuffle(block, &values, + coop::ShuffleOptions::default().mode(coop::Up)); + +The exposed edge remains invalid. For a block containing ``[10, 20, 30, 40]`` +in blocked order, Down produces ``[20, 30, 40, invalid]`` and Up produces +``[invalid, 10, 20, 30]``. Values cross thread boundaries; the shift neither +wraps around the block nor supplies a boundary value. + +A complete program +------------------ + +This helper loads a complete tile, shifts it, then initializes the exposed +edge through the safe payload API before storing. The input determines the +element type. ``ITEMS_PER_THREAD`` and the direction remain generic: + +.. literalinclude:: ../../rust/cuda-coop/examples/shuffle/helper/src/lib.rs + :language: rust + :start-after: // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +The host program specializes both directions for one and four items per +thread. Each of its three blocks shifts its own tile independently, with +zero supplied explicitly by the helper's caller for the missing edge: + +.. literalinclude:: ../../rust/cuda-coop/examples/shuffle/src/main.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/shuffle + 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 supplies the shared-storage +requirement to ``auto_shared``; the consumer does not calculate scratch bytes. + +Inputs and result validity +---------------------------- + +Shuffle supports complete one-, two- and three-dimensional Blocks, including +sizes that are not multiples of 32. Thread rank is x-fastest. Every thread +must participate with the same direction, element type and positive item count. +All ten ``Numeric`` types are supported. Ordinary same-extent helpers need +no ``generic_const_exprs`` feature. + +Every input slot must be initialized, including the one whose value does not +contribute to the result. A masked Load without a default can leave invalid +slots; Shuffle traps before reading them. The input remains unchanged. + +Down leaves the last item of the last thread invalid. Up leaves the first +item of the first thread invalid. ``get`` returns ``None`` at that edge, and +checked indexing or a Store that needs it traps. Use ``write`` to initialize +the slot before reading or storing it. A one-thread block with one item +returns an entirely invalid payload. + +``ShuffleOptions::default()`` selects Down. ``mode(Up)`` selects Up, and the +optional ``distance::<1>()`` spells the fixed distance explicitly. Other +distances fail compilation. The operation accepts payloads and complete +Blocks; it has no scalar, Warp, fill, count, algorithm or explicit-storage form. + +Participation and storage +------------------------- + +Oxide checks participation and plans the actual CUB shared storage, including +synchronization before reuse. Keep collective calls outside lane-specific +branches and reconverge after repairing the edge, as the helper does. +Shuffle does not promise completion of earlier global-memory accesses. Use +the block's ``sync()`` when dependent accesses need completion. + +See :doc:`payload` for checked access and :doc:`groups` for participation. +Generate the item reference with ``cargo doc --workspace --no-deps`` from +``rust/``; ``shuffle``, ``ShuffleOptions`` and ``ShuffleMode`` document the +full signatures and selectors. diff --git a/rust/cuda-coop-codegen/src/exchange.rs b/rust/cuda-coop-codegen/src/exchange.rs new file mode 100644 index 000000000000..3658faf62939 --- /dev/null +++ b/rust/cuda-coop-codegen/src/exchange.rs @@ -0,0 +1,255 @@ +// Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +//! Validated Block and Warp exchanges with preserved input and fresh output. + +use crate::{Dtype, Error, Result}; + +/// Direction of the redistribution of a group's complete tile. +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub enum ExchangeMode { + /// Convert striped input positions to blocked output positions. + StripedToBlocked, + /// Convert blocked input positions to striped output positions. + BlockedToStriped, +} + +/// Static group scope and direction for an out-of-place Exchange. +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub enum ExchangeRecipe { + /// Exchange across a complete Block, without warp time slicing. + Block(ExchangeMode), + /// Exchange within an aligned Warp using shared memory. + Warp(ExchangeMode), +} + +impl ExchangeRecipe { + fn mode(self) -> ExchangeMode { + match self { + Self::Block(mode) | Self::Warp(mode) => mode, + } + } + + pub(crate) fn token(self) -> &'static str { + match self.mode() { + ExchangeMode::StripedToBlocked => "exchange_striped_to_blocked_validity_bits_v1", + ExchangeMode::BlockedToStriped => "exchange_blocked_to_striped_validity_bits_v1", + } + } + + pub(crate) fn algorithm_token(self) -> &'static str { + match self { + Self::Block(_) => "block", + Self::Warp(_) => "warp", + } + } + + pub(crate) fn parse(operation: &str, algorithm: &str) -> Result { + let mode = match operation { + "exchange_striped_to_blocked_validity_bits_v1" => ExchangeMode::StripedToBlocked, + "exchange_blocked_to_striped_validity_bits_v1" => ExchangeMode::BlockedToStriped, + _ => return Err(Error::message("unsupported Exchange operation")), + }; + match algorithm { + "block" => Ok(Self::Block(mode)), + "warp" => Ok(Self::Warp(mode)), + _ => Err(Error::message("Exchange requires a Block or Warp recipe")), + } + } +} + +/// Validated numeric Exchange with a fresh output of the same per-lane extent. +/// +/// Native calls use `(const T* input, T* output, u8* scratch, +/// const u8* input_bits, u8* output_bits)` with no scalar controls. Each lane +/// supplies `items` initialized input slots and `ceil(items/8)` packed, +/// low-bit-first validity bytes. All input bits are checked before any numeric +/// read. Input values and validity bytes are preserved. +/// +/// Output contains `items` fresh slots and `ceil(items/8)` validity bytes, +/// disjoint from input and scratch. Every slot is written and marked initialized; +/// unused high validity bits are zero. All pointers must remain live and have +/// their natural alignment and the stated capacities throughout the call. +/// +/// Every member of the exact Block or aligned Warp participates in converged +/// control flow. Warp widths are 1/2/4/8/16/32 within complete physical warps. +/// The reported scratch is the actual CUB TempStorage for one group. The caller +/// owns its alignment, replication and lifetime, and synchronizes group members +/// before reuse. The provider adds no trailing completion barrier. +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub struct ExchangeRequest { + pub(crate) recipe: ExchangeRecipe, + pub(crate) dtype: Dtype, + block: [u32; 3], + items: u32, + width: Option, +} + +impl ExchangeRequest { + /// Validate a complete Block Exchange with exact XYZ geometry. + /// + /// # Errors + /// Rejects a Warp recipe, zero/overflowing axes, more than 1024 threads, + /// zero items, or a padded tile exceeding CUB's signed-int indexing range. + pub fn block( + dtype: Dtype, + recipe: ExchangeRecipe, + block: [u32; 3], + items: u32, + ) -> Result { + if !matches!(recipe, ExchangeRecipe::Block(_)) { + return Err(Error::message("Block Exchange requires a Block recipe")); + } + Self::new(dtype, recipe, block, items, None) + } + + /// Validate an aligned Warp Exchange within exact XYZ Block geometry. + /// + /// # Errors + /// Rejects a Block recipe, invalid Block geometry, incomplete physical warps, + /// widths other than 1/2/4/8/16/32, zero items, or a padded tile exceeding + /// CUB's signed-int indexing range. + pub fn warp( + dtype: Dtype, + recipe: ExchangeRecipe, + block: [u32; 3], + width: u32, + items: u32, + ) -> Result { + if !matches!(recipe, ExchangeRecipe::Warp(_)) { + return Err(Error::message("Warp Exchange requires a Warp recipe")); + } + Self::new(dtype, recipe, block, items, Some(width)) + } + + fn new( + dtype: Dtype, + recipe: ExchangeRecipe, + 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("Exchange requires 1..=1024 block threads with positive axes") + })?; + if let Some(width) = width { + if !width.is_power_of_two() || width > 32 { + return Err(Error::message( + "Warp Exchange requires power-of-two width 1..=32", + )); + } + if threads % 32 != 0 { + return Err(Error::message( + "Warp Exchange requires complete physical warps", + )); + } + } + // CUB adds one Warp sentinel and, for power-of-two extents > 4, + // bank-conflict padding. Its indexing and array bounds use signed int. + let tile = + u64::from(width.unwrap_or(threads)) * u64::from(items) + u64::from(width.is_some()); + let padding = if items > 4 && items.is_power_of_two() { + tile >> 5 + } else { + 0 + }; + if items == 0 || tile + padding > i32::MAX as u64 { + return Err(Error::message( + "Exchange requires positive items and padded tile <= i32::MAX", + )); + } + Ok(Self { + recipe, + dtype, + block, + items, + width, + }) + } + + pub(crate) fn block_shape(self) -> [u32; 3] { + self.block + } + + pub(crate) fn source(self, symbol: &str) -> String { + let (dtype, _) = self.dtype.native(); + let [x, y, z] = self.block; + let items = self.items; + let bytes = items / 8 + u32::from(items % 8 != 0); + let (header, primitive) = match self.width { + None => ( + "block/block_exchange.cuh", + format!("cub::BlockExchange"), + ), + Some(width) => ( + "warp/warp_exchange.cuh", + format!("cub::WarpExchange"), + ), + }; + let method = match self.recipe.mode() { + ExchangeMode::StripedToBlocked => "StripedToBlocked", + ExchangeMode::BlockedToStriped => "BlockedToStriped", + }; + format!( + r#" +#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) {{ + using namespace {symbol}_detail; +#pragma unroll + for (int i = 0; i < {bytes}; ++i) output_bits[i] = 0; +#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]; + value_type local_output[{items}]; + primitive operation{{*reinterpret_cast(scratch)}}; + operation.{method}(local_input, local_output); +#pragma unroll + for (int i = 0; i < {items}; ++i) {{ + output[i] = local_output[i]; + output_bits[i / 8] |= static_cast(1u << (i % 8)); + }} +}} +"# + ) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn scope_geometry_and_index_bounds() { + let block = ExchangeRecipe::Block(ExchangeMode::StripedToBlocked); + let warp = ExchangeRecipe::Warp(ExchangeMode::BlockedToStriped); + assert!(ExchangeRequest::block(Dtype::I8, block, [33, 1, 1], 8).is_ok()); + assert!(ExchangeRequest::block(Dtype::F64, block, [4, 3, 2], 4).is_ok()); + assert!(ExchangeRequest::warp(Dtype::U8, warp, [8, 4, 2], 1, 9).is_ok()); + assert!(ExchangeRequest::block(Dtype::I32, warp, [32, 1, 1], 1).is_err()); + assert!(ExchangeRequest::warp(Dtype::I32, block, [32, 1, 1], 32, 1).is_err()); + for shape in [[0, 1, 1], [1025, 1, 1], [u32::MAX, 2, 2]] { + assert!(ExchangeRequest::block(Dtype::I32, block, shape, 1).is_err()); + } + for width in [0, 3, 17, 31, 64] { + assert!(ExchangeRequest::warp(Dtype::I32, warp, [64, 1, 1], width, 1).is_err()); + } + assert!(ExchangeRequest::warp(Dtype::I32, warp, [33, 1, 1], 32, 1).is_err()); + assert!(ExchangeRequest::block(Dtype::I32, block, [1, 1, 1], 0).is_err()); + assert!(ExchangeRequest::block(Dtype::I32, block, [2, 1, 1], 1 << 30).is_err()); + // Warp adds a sentinel even when no bank-conflict padding is needed. + assert!(ExchangeRequest::warp(Dtype::I32, warp, [32, 1, 1], 1, i32::MAX as u32).is_err()); + } +} diff --git a/rust/cuda-coop-codegen/src/lib.rs b/rust/cuda-coop-codegen/src/lib.rs index 630bd6ad603c..d42f1a90d9fb 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 Load/Store, Scan, Reduce and batched Reduce providers with +//! This private development crate compiles Load/Store, Scan, Reduce, batched Reduce, Exchange and Shuffle 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. @@ -42,6 +42,11 @@ //! [`ReduceBatchedRequest`] reduces each payload position independently across //! a power-of-two Warp. Its fresh per-lane output capacity is ceil(batches/width), //! with striped or blocked distribution and validity only for real batches. +//! [`ExchangeRequest`] preserves complete numeric payloads and produces fresh +//! same-extent outputs redistributed between blocked and striped layouts. It +//! supports exact Blocks and power-of-two Warps using shared memory. +//! [`ShuffleRequest`] shifts a complete Block tile by one item, preserving input +//! and returning fresh output with exactly one invalid boundary slot. //! [`Provider::compile_native`] batches these families in one native compilation. //! //! [`Provider::from_checkout`] snapshots the selected CCCL and CUDA headers; @@ -71,13 +76,17 @@ #![deny(unsafe_code)] +mod exchange; mod headers; mod reduce; mod reduce_batched; mod scan; +mod shuffle; +pub use exchange::{ExchangeMode, ExchangeRecipe, ExchangeRequest}; pub use reduce::{ReduceAlgorithm, ReduceForm, ReduceOperator, ReduceRecipe, ReduceRequest}; pub use reduce_batched::{ReduceBatchedLayout, ReduceBatchedRecipe, ReduceBatchedRequest}; pub use scan::{ScanAlgorithm, ScanOperation, ScanOperator, ScanRecipe, ScanRequest}; +pub use shuffle::{ShuffleRecipe, ShuffleRequest}; #[allow(unsafe_code)] mod nvrtc; pub mod oxide; @@ -464,6 +473,10 @@ pub enum NativeRequest { Reduce(ReduceRequest), /// Batched Warp reduction with preserved input and partial fresh outputs. ReduceBatched(ReduceBatchedRequest), + /// Out-of-place Block or Warp Exchange with preserved input. + Exchange(ExchangeRequest), + /// One-item Block Shuffle with preserved input and a partial fresh output. + Shuffle(ShuffleRequest), } impl NativeRequest { @@ -473,6 +486,8 @@ impl NativeRequest { Self::Scan(request) => request.block_shape(), Self::Reduce(request) => request.block_shape(), Self::ReduceBatched(request) => request.block_shape(), + Self::Exchange(request) => request.block_shape(), + Self::Shuffle(request) => request.block_shape(), } } @@ -482,6 +497,8 @@ impl NativeRequest { Self::Scan(request) => request.source(symbol), Self::Reduce(request) => request.source(symbol), Self::ReduceBatched(request) => request.source(symbol), + Self::Exchange(request) => request.source(symbol), + Self::Shuffle(request) => request.source(symbol), } } @@ -491,13 +508,19 @@ impl NativeRequest { Self::Scan(request) => request.recipe.operation().token(), Self::Reduce(request) => request.recipe.token(), Self::ReduceBatched(request) => request.recipe.token(), + Self::Exchange(request) => request.recipe.token().into(), + Self::Shuffle(request) => request.recipe.token().into(), } } fn scratch_free(self) -> bool { match self { Self::LoadStore(request) => request.algorithm.scratch_free(), - Self::Scan(_) | Self::Reduce(_) | Self::ReduceBatched(_) => false, + Self::Scan(_) + | Self::Reduce(_) + | Self::ReduceBatched(_) + | Self::Exchange(_) + | Self::Shuffle(_) => false, } } @@ -507,7 +530,11 @@ impl NativeRequest { !request.operation.is_load() && !request.algorithm.scratch_free() } // These families write fresh outputs, never rearranging inputs. - Self::Scan(_) | Self::Reduce(_) | Self::ReduceBatched(_) => false, + Self::Scan(_) + | Self::Reduce(_) + | Self::ReduceBatched(_) + | Self::Exchange(_) + | Self::Shuffle(_) => false, } } } @@ -526,9 +553,9 @@ pub struct Layout { /// # Device ABI obligations /// /// The memory-tile and payload-layout requirements below describe Load/Store. -/// Scan, Reduce and batched Reduce use lane-private input and fresh output; +/// Scan, Reduce, batched Reduce, Exchange and Shuffle use lane-private input and fresh output; /// their ABI and validity contracts are documented by [`ScanRequest`], -/// [`ReduceRequest`] and [`ReduceBatchedRequest`]. +/// [`ReduceRequest`], [`ReduceBatchedRequest`], [`ExchangeRequest`] and [`ShuffleRequest`]. /// /// 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. @@ -682,7 +709,8 @@ impl Provider { /// Compile an ordered mixed batch of the supported families in one NVRTC program. /// /// Per-specialization identity is independent of unrelated batch members. - /// Scan/Reduce/batched Reduce scratch describes one group and includes the actual CUB TempStorage, + /// Scan/Reduce/batched Reduce/Exchange/Shuffle scratch describes one group and includes + /// the actual CUB TempStorage, /// even when its C++ empty type occupies one byte. The caller plans replication. /// /// # Errors @@ -747,6 +775,8 @@ impl Provider { include_bytes!("scan.rs"), include_bytes!("reduce.rs"), include_bytes!("reduce_batched.rs"), + include_bytes!("exchange.rs"), + include_bytes!("shuffle.rs"), ] { hash_field(&mut inputs, bytes); } diff --git a/rust/cuda-coop-codegen/src/oxide.rs b/rust/cuda-coop-codegen/src/oxide.rs index 411300aca6d2..e572f9acce69 100644 --- a/rust/cuda-coop-codegen/src/oxide.rs +++ b/rust/cuda-coop-codegen/src/oxide.rs @@ -14,7 +14,9 @@ //! The protocol supports raw i32 and typed validity Load/Store recipes documented //! by [`Operation`], numeric Block/scalar Warp Scan through [`ScanRecipe`], and //! root-result scalar/payload Reduce through [`ReduceRecipe`], and partial-output -//! Warp batched Reduce through [`ReduceBatchedRecipe`]. +//! Warp batched Reduce through [`ReduceBatchedRecipe`]. Out-of-place Block/Warp +//! Exchange uses [`ExchangeRecipe`] with no scalar controls; Block Shuffle uses +//! [`ShuffleRecipe`] for one-item shifts with partial fresh output. //! 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 @@ -22,9 +24,10 @@ //! checked callback marker; the provider does not select or execute Rust code. use crate::{ - Algorithm, Artifact, Dtype, Error, Layout, NativeRequest, Operation, Provider, ReduceAlgorithm, - ReduceBatchedRecipe, ReduceBatchedRequest, ReduceForm, ReduceOperator, ReduceRecipe, - ReduceRequest, Request, Result, ScanRecipe, ScanRequest, + Algorithm, Artifact, Dtype, Error, ExchangeRecipe, ExchangeRequest, Layout, NativeRequest, + Operation, Provider, ReduceAlgorithm, ReduceBatchedRecipe, ReduceBatchedRequest, ReduceForm, + ReduceOperator, ReduceRecipe, ReduceRequest, Request, Result, ScanRecipe, ScanRequest, + ShuffleRecipe, ShuffleRequest, }; use serde::{Deserialize, Serialize}; use std::io::{Read, Write}; @@ -191,6 +194,33 @@ struct SpecializationRequest { impl SpecializationRequest { fn native(&self) -> Result { + if self.operation.starts_with("shuffle_") { + let recipe = ShuffleRecipe::parse(&self.operation, &self.algorithm)?; + if self.warp_width.is_some() { + return Err(Error::message("Shuffle requires a Block scope")); + } + return ShuffleRequest::block( + self.dtype.dtype()?, + recipe, + self.block, + self.items_per_thread, + ) + .map(NativeRequest::Shuffle); + } + if self.operation.starts_with("exchange_") { + let recipe = ExchangeRecipe::parse(&self.operation, &self.algorithm)?; + let dtype = self.dtype.dtype()?; + let request = match (recipe, self.warp_width) { + (ExchangeRecipe::Block(_), None) => { + ExchangeRequest::block(dtype, recipe, self.block, self.items_per_thread)? + } + (ExchangeRecipe::Warp(_), Some(width)) => { + ExchangeRequest::warp(dtype, recipe, self.block, width, self.items_per_thread)? + } + _ => return Err(Error::message("Exchange recipe disagrees with Warp scope")), + }; + return Ok(NativeRequest::Exchange(request)); + } if self.operation.starts_with("reduce_batched_") { let recipe = ReduceBatchedRecipe::parse(&self.operation, &self.algorithm)?; let width = self @@ -330,6 +360,24 @@ struct BatchResponse { fn arguments(request: NativeRequest) -> Vec { let request = match request { + NativeRequest::Shuffle(request) => { + let dtype = Scalar::from(request.dtype); + return vec![ + Argument::ConstPointer(dtype), + Argument::MutPointer(dtype), + Argument::ConstPointer(Scalar::Unsigned(8)), + Argument::MutPointer(Scalar::Unsigned(8)), + ]; + } + NativeRequest::Exchange(request) => { + let dtype = Scalar::from(request.dtype); + return vec![ + Argument::ConstPointer(dtype), + Argument::MutPointer(dtype), + Argument::ConstPointer(Scalar::Unsigned(8)), + Argument::MutPointer(Scalar::Unsigned(8)), + ]; + } NativeRequest::ReduceBatched(request) => { let dtype = Scalar::from(request.dtype); return vec![ @@ -577,6 +625,52 @@ pub fn reduce_batched_registration( .map_err(|error| Error::message(format!("cannot encode provider registration: {error}"))) } +/// Encode a Block or Warp Exchange recipe for the existing provider executable. +/// +/// Dtype, exact Block geometry, optional Warp width and payload extent arrive +/// with specialization. Four typed payload pointers describe preserved input +/// and fresh same-extent output with packed validity; there are no scalar +/// controls. Scratch is inserted by the existing library-call ABI. +/// +/// # Errors +/// Rejects a nonabsolute executable or a path that cannot be encoded as JSON. +pub fn exchange_registration(executable: &Path, recipe: ExchangeRecipe) -> 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}"))) +} + +/// Encode a one-item Block Shuffle recipe for the existing provider executable. +/// +/// Dtype, exact Block geometry and payload extent arrive with specialization. +/// Four typed payload pointers describe preserved input and fresh same-extent +/// output with packed validity, including its one invalid boundary slot. There +/// are no scalar controls; scratch is inserted by the existing library-call ABI. +/// +/// # Errors +/// Rejects a nonabsolute executable or a path that cannot be encoded as JSON. +pub fn shuffle_registration(executable: &Path, recipe: ShuffleRecipe) -> 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: "block", + }) + .map_err(|error| Error::message(format!("cannot encode provider registration: {error}"))) +} + #[derive(Deserialize)] #[serde(deny_unknown_fields)] struct InputRequest { @@ -899,6 +993,107 @@ mod tests { ScanOperator, }; + #[test] + fn shuffle_registration_preserves_scope_extent_and_pointer_abi() { + for recipe in [ShuffleRecipe::Down, ShuffleRecipe::Up] { + let registration: serde_json::Value = serde_json::from_str( + &shuffle_registration(Path::new("/provider"), recipe).unwrap(), + ) + .unwrap(); + let mut request = SpecializationRequest { + operation: registration["operation"].as_str().unwrap().into(), + algorithm: registration["algorithm"].as_str().unwrap().into(), + dtype: Scalar::Float(64), + block: [4, 3, 2], + items_per_thread: 9, + warp_width: None, + }; + let native = request.native().unwrap(); + assert_eq!( + native, + NativeRequest::Shuffle( + ShuffleRequest::block(Dtype::F64, recipe, [4, 3, 2], 9).unwrap() + ) + ); + assert_eq!( + arguments(native), + vec![ + Argument::ConstPointer(Scalar::Float(64)), + Argument::MutPointer(Scalar::Float(64)), + Argument::ConstPointer(Scalar::Unsigned(8)), + Argument::MutPointer(Scalar::Unsigned(8)), + ] + ); + assert!(!native.scratch_free() && !native.mutates_payload()); + assert_eq!(registration["version"], 1); + assert_eq!(registration["provider"], "cuda_coop"); + request.warp_width = Some(32); + assert!(request.native().is_err()); + } + assert!(shuffle_registration(Path::new("relative"), ShuffleRecipe::Down).is_err()); + assert!(ShuffleRecipe::parse("shuffle_rotate", "block").is_err()); + assert!(ShuffleRecipe::parse(ShuffleRecipe::Down.token(), "warp").is_err()); + assert!(ShuffleRecipe::parse(ShuffleRecipe::Up.token(), "transpose").is_err()); + } + + #[test] + fn exchange_registration_preserves_scope_extent_and_pointer_abi() { + use crate::ExchangeMode::{BlockedToStriped, StripedToBlocked}; + for (recipe, width) in [ + (ExchangeRecipe::Block(StripedToBlocked), None), + (ExchangeRecipe::Block(BlockedToStriped), None), + (ExchangeRecipe::Warp(StripedToBlocked), Some(8)), + (ExchangeRecipe::Warp(BlockedToStriped), Some(8)), + ] { + let registration: serde_json::Value = serde_json::from_str( + &exchange_registration(Path::new("/provider"), recipe).unwrap(), + ) + .unwrap(); + let mut request = SpecializationRequest { + operation: registration["operation"].as_str().unwrap().into(), + algorithm: registration["algorithm"].as_str().unwrap().into(), + dtype: Scalar::Float(64), + block: [8, 4, 2], + items_per_thread: 9, + warp_width: width, + }; + let native = request.native().unwrap(); + let expected = match width { + None => ExchangeRequest::block(Dtype::F64, recipe, [8, 4, 2], 9).unwrap(), + Some(width) => { + ExchangeRequest::warp(Dtype::F64, recipe, [8, 4, 2], width, 9).unwrap() + } + }; + assert_eq!(native, NativeRequest::Exchange(expected)); + assert_eq!( + arguments(native), + vec![ + Argument::ConstPointer(Scalar::Float(64)), + Argument::MutPointer(Scalar::Float(64)), + Argument::ConstPointer(Scalar::Unsigned(8)), + Argument::MutPointer(Scalar::Unsigned(8)), + ] + ); + assert!(!native.scratch_free() && !native.mutates_payload()); + assert_eq!(registration["version"], 1); + assert_eq!(registration["provider"], "cuda_coop"); + request.warp_width = if width.is_some() { None } else { Some(8) }; + assert!(request.native().is_err()); + } + assert!( + exchange_registration( + Path::new("relative"), + ExchangeRecipe::Block(StripedToBlocked) + ) + .is_err() + ); + assert!(ExchangeRecipe::parse("exchange_ranked", "block").is_err()); + assert!( + ExchangeRecipe::parse(ExchangeRecipe::Block(StripedToBlocked).token(), "shuffle") + .is_err() + ); + } + #[test] fn batched_registration_preserves_derived_extent_and_pointer_abi() { for (dtype, op, layout, width, batches, outputs) in [ diff --git a/rust/cuda-coop-codegen/src/shuffle.rs b/rust/cuda-coop-codegen/src/shuffle.rs new file mode 100644 index 000000000000..e0ca6af12c16 --- /dev/null +++ b/rust/cuda-coop-codegen/src/shuffle.rs @@ -0,0 +1,179 @@ +// Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +//! Validated one-item Block shifts with preserved input and partial fresh output. + +use crate::{Dtype, Error, Result}; + +/// Direction of a one-item shift across a Block's blocked-order tile. +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub enum ShuffleRecipe { + /// Read each slot's successor; the last tile slot is invalid. + Down, + /// Read each slot's predecessor; the first tile slot is invalid. + Up, +} + +impl ShuffleRecipe { + pub(crate) fn token(self) -> &'static str { + match self { + Self::Down => "shuffle_down_validity_bits_v1", + Self::Up => "shuffle_up_validity_bits_v1", + } + } + + pub(crate) fn parse(operation: &str, algorithm: &str) -> Result { + if algorithm != "block" { + return Err(Error::message("Shuffle requires a Block recipe")); + } + match operation { + "shuffle_down_validity_bits_v1" => Ok(Self::Down), + "shuffle_up_validity_bits_v1" => Ok(Self::Up), + _ => Err(Error::message("unsupported Shuffle operation")), + } + } +} + +/// Validated numeric Block Shuffle with a fresh output of the same lane extent. +/// +/// Native calls use `(const T* input, T* output, u8* scratch, +/// const u8* input_bits, u8* output_bits)` with no scalar controls. Every lane +/// supplies `items` initialized input slots and `ceil(items/8)` packed, +/// low-bit-first validity bytes. Every input bit is checked before any numeric +/// read or collective call, including input slots discarded by the shift. +/// Input values and validity bytes are preserved. +/// +/// Output has `items` slots and `ceil(items/8)` validity bytes, disjoint from +/// input and scratch. For x-fastest Block rank `r`, item `i` and extent `N`, +/// Down reads the successor of blocked position `r*N+i`; Up reads its +/// predecessor. Down leaves the last Block rank's item `N-1` unwritten and +/// invalid; Up leaves rank zero's item zero unwritten and invalid. All other +/// slots are initialized. Unused high validity bits are zero. Block1/N1 has +/// no initialized output. There is no wraparound or boundary fill. +/// +/// Every member of the exact Block participates in converged control flow. +/// All pointers remain live with their natural alignment and stated capacities +/// throughout the call. Scratch has the reported actual CUB TempStorage layout; +/// the caller owns its alignment and lifetime and synchronizes Block members +/// before reuse. The provider adds no trailing completion barrier. +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub struct ShuffleRequest { + pub(crate) recipe: ShuffleRecipe, + pub(crate) dtype: Dtype, + block: [u32; 3], + items: u32, +} + +impl ShuffleRequest { + /// Validate a complete Block Shuffle with exact XYZ geometry and distance1. + /// + /// `dtype` selects one of the ten numeric types. `recipe` selects Down or + /// Up. `block` gives the three positive Block axes; `items` is each lane's + /// fixed input and output extent. Warp and scalar forms are not supported. + /// + /// # Errors + /// Rejects zero/overflowing axes, more than 1024 threads, zero items, or + /// an item extent above CUB's signed-int array indexing range. + pub fn block(dtype: Dtype, recipe: ShuffleRecipe, block: [u32; 3], items: u32) -> Result { + block + .into_iter() + .try_fold(1_u32, u32::checked_mul) + .filter(|n| (1..=1024).contains(n)) + .ok_or_else(|| { + Error::message("Shuffle requires 1..=1024 block threads with positive axes") + })?; + if items == 0 || items > i32::MAX as u32 { + return Err(Error::message("Shuffle requires items in 1..=i32::MAX")); + } + Ok(Self { + recipe, + dtype, + block, + items, + }) + } + + pub(crate) fn block_shape(self) -> [u32; 3] { + self.block + } + + pub(crate) fn source(self, symbol: &str) -> String { + let (dtype, _) = self.dtype.native(); + let [x, y, z] = self.block; + let threads = x * y * z; + let items = self.items; + let bytes = items / 8 + u32::from(items % 8 != 0); + let (method, valid) = match self.recipe { + ShuffleRecipe::Down => ( + "Down", + format!("thread_rank != {} || i != {}", threads - 1, items - 1), + ), + ShuffleRecipe::Up => ("Up", "thread_rank != 0 || i != 0".into()), + }; + format!( + r#" +#include +#include +namespace {symbol}_detail {{ +using value_type = {dtype}; +using primitive = cub::BlockShuffle; +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) {{ + using namespace {symbol}_detail; + const unsigned int thread_rank = cub::RowMajorTid({x}, {y}, {z}); +#pragma unroll + for (int i = 0; i < {bytes}; ++i) output_bits[i] = 0; +#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]; + value_type local_output[{items}]; + primitive operation{{*reinterpret_cast(scratch)}}; + operation.{method}(local_input, local_output); +#pragma unroll + for (int i = 0; i < {items}; ++i) {{ + if ({valid}) {{ + output[i] = local_output[i]; + output_bits[i / 8] |= static_cast(1u << (i % 8)); + }} + }} +}} +"# + ) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn geometry_and_local_index_bounds() { + for recipe in [ShuffleRecipe::Down, ShuffleRecipe::Up] { + for (shape, items) in [ + ([1, 1, 1], 1), + ([33, 1, 1], 9), + ([4, 3, 2], 4), + ([1024, 1, 1], i32::MAX as u32), + ] { + assert!(ShuffleRequest::block(Dtype::I8, recipe, shape, items).is_ok()); + } + for shape in [ + [0, 1, 1], + [1, 0, 1], + [1, 1, 0], + [1025, 1, 1], + [u32::MAX, 2, 2], + ] { + assert!(ShuffleRequest::block(Dtype::I32, recipe, shape, 1).is_err()); + } + for items in [0, i32::MAX as u32 + 1, u32::MAX] { + assert!(ShuffleRequest::block(Dtype::F64, recipe, [1, 1, 1], items).is_err()); + } + } + } +} diff --git a/rust/cuda-coop-codegen/tests/native.rs b/rust/cuda-coop-codegen/tests/native.rs index b7bdb29d4b73..ce3cbddfe0ea 100644 --- a/rust/cuda-coop-codegen/tests/native.rs +++ b/rust/cuda-coop-codegen/tests/native.rs @@ -4,12 +4,137 @@ //! Real native compiler coverage; ordinary host tests need no CUDA installation. use cuda_coop_codegen::{ - Algorithm, Dtype, Layout, NativeRequest, Operation, Provider, ReduceAlgorithm, - ReduceBatchedLayout, ReduceBatchedRecipe, ReduceBatchedRequest, ReduceForm, ReduceOperator, - ReduceRequest, Request, ScanAlgorithm, ScanOperation, ScanOperator, ScanRequest, + Algorithm, Dtype, ExchangeMode, ExchangeRecipe, ExchangeRequest, Layout, NativeRequest, + Operation, Provider, ReduceAlgorithm, ReduceBatchedLayout, ReduceBatchedRecipe, + ReduceBatchedRequest, ReduceForm, ReduceOperator, ReduceRequest, Request, ScanAlgorithm, + ScanOperation, ScanOperator, ScanRequest, ShuffleRecipe, ShuffleRequest, }; use std::{env, path::Path, time::Instant}; +#[test] +#[ignore = "requires CUDA_HOME and a Linux CUDA toolkit with NVRTC LTO support"] +fn native_mixed_shuffle_recipes() -> Result<(), Box> { + use ShuffleRecipe::{Down, Up}; + 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(); + // Cover all numeric types, both empty-result directions, cross-thread + // shifts, XYZ and a second validity byte without a Cartesian matrix. + for (dtype, block, items, recipe) in [ + (Dtype::I8, [1, 1, 1], 1, Down), + (Dtype::U8, [1, 1, 1], 1, Up), + (Dtype::I16, [33, 1, 1], 1, Down), + (Dtype::U16, [4, 3, 2], 4, Up), + (Dtype::I32, [1, 1, 1], 4, Down), + (Dtype::U32, [33, 1, 1], 9, Up), + (Dtype::I64, [64, 1, 1], 4, Down), + (Dtype::U64, [4, 3, 2], 1, Up), + (Dtype::F32, [128, 1, 1], 4, Down), + (Dtype::F64, [4, 3, 2], 9, Up), + (Dtype::I32, [33, 1, 1], 9, Down), + (Dtype::I32, [33, 1, 1], 9, Up), + ] { + requests.push(NativeRequest::Shuffle(ShuffleRequest::block( + dtype, recipe, block, items, + )?)); + } + requests.extend([ + NativeRequest::LoadStore(Request::new_typed( + Dtype::I32, + Operation::LoadFullValidity, + Algorithm::Direct, + 64, + 4, + )?), + NativeRequest::Scan(ScanRequest::warp( + Dtype::I32, + ScanOperation::ExclusiveSum, + [64, 1, 1], + 8, + )?), + NativeRequest::Reduce(ReduceRequest::warp( + Dtype::I32, + ReduceOperator::Sum, + ReduceForm::Payload, + [64, 1, 1], + 8, + 4, + )?), + NativeRequest::ReduceBatched(ReduceBatchedRequest::warp( + Dtype::I32, + ReduceBatchedRecipe::new(ReduceOperator::Sum, ReduceBatchedLayout::Striped), + [64, 1, 1], + 8, + 4, + )?), + NativeRequest::Exchange(ExchangeRequest::warp( + Dtype::I32, + ExchangeRecipe::Warp(ExchangeMode::StripedToBlocked), + [64, 1, 1], + 32, + 4, + )?), + ]); + let artifact = provider.compile_native(&requests, &target)?; + assert_eq!(artifact.specializations.len(), 17); + 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 + ); + } + // CUB scratch still exists for Block1/N1 even though the result is invalid. + assert_eq!( + artifact.specializations[0].scratch, + Layout { + size: 1, + alignment: 1 + } + ); + assert_eq!( + artifact.specializations[1].scratch, + Layout { + size: 1, + alignment: 1 + } + ); + assert_eq!( + artifact.specializations[9].scratch, + Layout { + size: 192, + alignment: 8 + } + ); + assert!( + provider + .compile_native(&[requests[0], requests[0]], &target) + .unwrap_err() + .to_string() + .contains("duplicate") + ); + eprintln!( + "{} mixed Shuffle/Load/Scan/Reduce/batched Reduce/Exchange 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_warp_shape_identity() -> Result<(), Box> { @@ -110,6 +235,111 @@ fn native_warp_shape_identity() -> Result<(), Box> { Ok(()) } +#[test] +#[ignore = "requires CUDA_HOME and a Linux CUDA toolkit with NVRTC LTO support"] +fn native_mixed_exchange_recipes() -> Result<(), Box> { + use ExchangeMode::{BlockedToStriped, StripedToBlocked}; + 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)?; + // Ten numeric types, both directions, all Warp widths, XYZ and padded + // extents are distributed over rows rather than taking their product. + let mut requests = Vec::new(); + for (dtype, block, width, items, mode) in [ + (Dtype::I8, [33, 1, 1], None, 1, StripedToBlocked), + (Dtype::U8, [64, 1, 1], Some(1), 9, BlockedToStriped), + (Dtype::I16, [4, 3, 2], None, 4, BlockedToStriped), + (Dtype::U16, [4, 4, 4], Some(2), 4, StripedToBlocked), + (Dtype::I32, [128, 1, 1], None, 4, StripedToBlocked), + (Dtype::U32, [64, 1, 1], Some(4), 8, BlockedToStriped), + (Dtype::I64, [1, 1, 1], None, 1, BlockedToStriped), + (Dtype::U64, [64, 1, 1], Some(8), 4, StripedToBlocked), + (Dtype::F32, [64, 1, 1], Some(16), 1, BlockedToStriped), + (Dtype::F64, [8, 4, 2], Some(32), 4, StripedToBlocked), + (Dtype::I32, [33, 1, 1], None, 8, BlockedToStriped), + (Dtype::I32, [64, 1, 1], Some(32), 4, StripedToBlocked), + (Dtype::I32, [128, 1, 1], Some(32), 4, StripedToBlocked), + ] { + requests.push(NativeRequest::Exchange(match width { + None => ExchangeRequest::block(dtype, ExchangeRecipe::Block(mode), block, items)?, + Some(width) => { + ExchangeRequest::warp(dtype, ExchangeRecipe::Warp(mode), block, width, items)? + } + })); + } + requests.push(NativeRequest::LoadStore(Request::new_typed( + Dtype::I32, + Operation::LoadFullValidity, + Algorithm::Direct, + 64, + 4, + )?)); + requests.push(NativeRequest::Scan(ScanRequest::warp( + Dtype::I32, + ScanOperation::ExclusiveSum, + [64, 1, 1], + 8, + )?)); + requests.push(NativeRequest::Reduce(ReduceRequest::warp( + Dtype::I32, + ReduceOperator::Sum, + ReduceForm::Payload, + [64, 1, 1], + 8, + 4, + )?)); + requests.push(NativeRequest::ReduceBatched(ReduceBatchedRequest::warp( + Dtype::I32, + ReduceBatchedRecipe::new(ReduceOperator::Sum, ReduceBatchedLayout::Striped), + [64, 1, 1], + 8, + 4, + )?)); + let artifact = provider.compile_native(&requests, &target)?; + assert_eq!(artifact.specializations.len(), 17); + 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 + ); + } + // Actual queried layouts include Warp sentinel/padding and Block padding. + assert_eq!( + artifact.specializations[5].scratch, + Layout { + size: 136, + alignment: 4 + } + ); + assert_eq!( + artifact.specializations[10].scratch, + Layout { + size: 1088, + alignment: 16 + } + ); + eprintln!( + "{} mixed Exchange/Load/Scan/Reduce/batched Reduce 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_reduce_batched_recipes() -> Result<(), Box> { diff --git a/rust/cuda-coop/build.rs b/rust/cuda-coop/build.rs index add35460896f..c1144e841372 100644 --- a/rust/cuda-coop/build.rs +++ b/rust/cuda-coop/build.rs @@ -1,10 +1,11 @@ // 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, Scan, Reduce and batched reduction. +//! Register the owning host provider for native cooperative operations. use cuda_coop_codegen::{ - Algorithm, Operation, ReduceAlgorithm, ReduceBatchedLayout, ReduceBatchedRecipe, ReduceForm, - ReduceOperator, ReduceRecipe, ScanAlgorithm, ScanOperation, ScanOperator, ScanRecipe, oxide, + Algorithm, ExchangeMode, ExchangeRecipe, Operation, ReduceAlgorithm, ReduceBatchedLayout, + ReduceBatchedRecipe, ReduceForm, ReduceOperator, ReduceRecipe, ScanAlgorithm, ScanOperation, + ScanOperator, ScanRecipe, ShuffleRecipe, oxide, }; fn main() { @@ -23,6 +24,25 @@ fn main() { let callable = oxide::reduce_callable_registration(&executable) .expect("cuda.coop callable Reduce registration"); println!("cargo::rustc-env=COOP_REDUCE_RAKING_CALLABLE_PAYLOAD={callable}"); + for (recipe, name) in [(ShuffleRecipe::Down, "DOWN"), (ShuffleRecipe::Up, "UP")] { + let value = oxide::shuffle_registration(&executable, recipe) + .expect("cuda.coop shuffle registration"); + println!("cargo::rustc-env=COOP_SHUFFLE_{name}={value}"); + } + for (mode, name) in [ + (ExchangeMode::StripedToBlocked, "STRIPED_TO_BLOCKED"), + (ExchangeMode::BlockedToStriped, "BLOCKED_TO_STRIPED"), + ] { + for (recipe, group) in [ + (ExchangeRecipe::Block(mode), "BLOCK"), + (ExchangeRecipe::Warp(mode), "WARP"), + ] { + let value = oxide::exchange_registration(&executable, recipe) + .expect("cuda.coop exchange registration"); + println!("cargo::rustc-env=COOP_EXCHANGE_{group}_{name}={value}"); + } + } + for (operator, name) in [ (ReduceOperator::Sum, "SUM"), (ReduceOperator::Multiplies, "MULTIPLIES"), diff --git a/rust/cuda-coop/examples/exchange/.gitignore b/rust/cuda-coop/examples/exchange/.gitignore new file mode 100644 index 000000000000..74bffa8af4a4 --- /dev/null +++ b/rust/cuda-coop/examples/exchange/.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/exchange/Cargo.lock b/rust/cuda-coop/examples/exchange/Cargo.lock new file mode 100644 index 000000000000..199801956139 --- /dev/null +++ b/rust/cuda-coop/examples/exchange/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-exchange" +version = "0.0.0" +dependencies = [ + "coop-exchange-helper", + "cuda-async", + "cuda-coop", + "cuda-core", + "cuda-device", + "cuda-host", +] + +[[package]] +name = "coop-exchange-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/exchange/Cargo.toml b/rust/cuda-coop/examples/exchange/Cargo.toml new file mode 100644 index 000000000000..ca37943a559b --- /dev/null +++ b/rust/cuda-coop/examples/exchange/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-exchange" +version = "0.0.0" +edition = "2024" +publish = false + +[workspace] +members = ["helper"] + +[dependencies] +cuda-coop = { path = "../.." } +coop-exchange-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/exchange/helper/Cargo.toml b/rust/cuda-coop/examples/exchange/helper/Cargo.toml new file mode 100644 index 000000000000..41e3e03a876a --- /dev/null +++ b/rust/cuda-coop/examples/exchange/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-exchange-helper" +version = "0.0.0" +edition = "2024" +publish = false + +[dependencies] +cuda-coop = { path = "../../.." } diff --git a/rust/cuda-coop/examples/exchange/helper/src/lib.rs b/rust/cuda-coop/examples/exchange/helper/src/lib.rs new file mode 100644 index 000000000000..8243057c6348 --- /dev/null +++ b/rust/cuda-coop/examples/exchange/helper/src/lib.rs @@ -0,0 +1,40 @@ +// Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +//! A reusable, element-type-preserving block layout conversion helper. +#![no_std] +use cuda_coop as coop; + +/// Copy a tile through striped input and blocked output layouts. +/// +/// `fill` initializes the missing tail before Exchange. Only the valid prefix +/// is stored, so any value of `T` may fill slots outside that prefix. +pub fn exchange_tile( + block: coop::ThreadGroup, + input: &[T], + output: &mut coop::DisjointSlice<'_, T, coop::CollectiveTiles>, + offset: usize, + valid: usize, + fill: T, +) { + let mut values = coop::ThreadData::<_, ITEMS_PER_THREAD>::new(); + coop::load( + block, + input, + &mut values, + coop::LoadOptions::default() + .algorithm(coop::Striped) + .offset(offset) + .valid_items(valid) + .oob_default(fill), + ); + let blocked = coop::exchange(block, &values, Default::default()); + coop::store( + block, + output, + &blocked, + coop::StoreOptions::default() + .offset(offset) + .valid_items(valid), + ); +} diff --git a/rust/cuda-coop/examples/exchange/src/main.rs b/rust/cuda-coop/examples/exchange/src/main.rs new file mode 100644 index 000000000000..dded76393d00 --- /dev/null +++ b/rust/cuda-coop/examples/exchange/src/main.rs @@ -0,0 +1,69 @@ +// Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +use cuda_coop as coop; +use cuda_device::{kernel, launch_contract}; +use cuda_host::cuda_module; + +#[cuda_module] +mod kernels { + use super::*; + + #[kernel] + #[launch_contract(domain = 1, block = (128, 1, 1), dynamic_shared = inferred)] + pub fn exchange_tile( + 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_exchange_helper::exchange_tile::<_, 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 as i32 - 31).collect(); + let expected = input.clone(); + let source = DeviceBuffer::from_host(&stream, &input)?; + let target = DeviceBuffer::from_host(&stream, &vec![-1_i32; input.len()])?; + let plan = + module.prepare_exchange_tile::<$items>(LaunchConfig1D::auto_shared(3, 128))?; + let (_, target) = module + .exchange_tile_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={}: {} values checked", $items, output.len()); + }}; + } + run!(1); + run!(4); + Ok(()) +} diff --git a/rust/cuda-coop/examples/shuffle/.gitignore b/rust/cuda-coop/examples/shuffle/.gitignore new file mode 100644 index 000000000000..74bffa8af4a4 --- /dev/null +++ b/rust/cuda-coop/examples/shuffle/.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/shuffle/Cargo.lock b/rust/cuda-coop/examples/shuffle/Cargo.lock new file mode 100644 index 000000000000..b309b992cd19 --- /dev/null +++ b/rust/cuda-coop/examples/shuffle/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-shuffle" +version = "0.0.0" +dependencies = [ + "coop-shuffle-helper", + "cuda-async", + "cuda-coop", + "cuda-core", + "cuda-device", + "cuda-host", +] + +[[package]] +name = "coop-shuffle-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/shuffle/Cargo.toml b/rust/cuda-coop/examples/shuffle/Cargo.toml new file mode 100644 index 000000000000..3b8d35695d11 --- /dev/null +++ b/rust/cuda-coop/examples/shuffle/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-shuffle" +version = "0.0.0" +edition = "2024" +publish = false + +[workspace] +members = ["helper"] + +[dependencies] +cuda-coop = { path = "../.." } +coop-shuffle-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/shuffle/helper/Cargo.toml b/rust/cuda-coop/examples/shuffle/helper/Cargo.toml new file mode 100644 index 000000000000..897a021b698c --- /dev/null +++ b/rust/cuda-coop/examples/shuffle/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-shuffle-helper" +version = "0.0.0" +edition = "2024" +publish = false + +[dependencies] +cuda-coop = { path = "../../.." } diff --git a/rust/cuda-coop/examples/shuffle/helper/src/lib.rs b/rust/cuda-coop/examples/shuffle/helper/src/lib.rs new file mode 100644 index 000000000000..4f71f9bfb8e7 --- /dev/null +++ b/rust/cuda-coop/examples/shuffle/helper/src/lib.rs @@ -0,0 +1,41 @@ +// Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +//! Shift a full block tile while retaining its element type and static extent. +#![no_std] +use cuda_coop as coop; + +/// Load a tile, shift it, initialize the exposed edge, and store the result. +/// +/// Every thread participates. `edge` supplies the one missing boundary value; +/// inspecting validity and writing that slot use the ordinary safe payload API. +pub fn shuffle_tile( + block: coop::ThreadGroup, + input: &[T], + output: &mut coop::DisjointSlice<'_, T, coop::CollectiveTiles>, + offset: usize, + edge: T, + options: coop::ShuffleOptions, +) { + let mut values = coop::ThreadData::<_, ITEMS_PER_THREAD>::new(); + coop::load( + block, + input, + &mut values, + coop::LoadOptions::default().offset(offset), + ); + let mut shifted = coop::shuffle(block, &values, options); + let mut item = 0; + while item < ITEMS_PER_THREAD { + if shifted.get(item).is_none() { + shifted.write(item, edge); + } + item += 1; + } + coop::store( + block, + output, + &shifted, + coop::StoreOptions::default().offset(offset), + ); +} diff --git a/rust/cuda-coop/examples/shuffle/src/main.rs b/rust/cuda-coop/examples/shuffle/src/main.rs new file mode 100644 index 000000000000..957afc4b2c5f --- /dev/null +++ b/rust/cuda-coop/examples/shuffle/src/main.rs @@ -0,0 +1,99 @@ +// 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 shuffle_tile( + input: &[i32], + mut output: coop::DisjointSlice>, + ) { + let block = coop::this_block(); + let offset = + block.rank_at::() as usize * block.count() as usize * ITEMS_PER_THREAD; + if UP { + coop_shuffle_helper::shuffle_tile::<_, _, ITEMS_PER_THREAD>( + block, + input, + &mut output, + offset, + 0_i32, + coop::ShuffleOptions::default() + .mode(coop::Up) + .distance::<1>(), + ); + } else { + coop_shuffle_helper::shuffle_tile::<_, _, ITEMS_PER_THREAD>( + block, + input, + &mut output, + offset, + 0_i32, + Default::default(), + ); + } + } +} + +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, $up:literal) => {{ + let tile = 128 * $items; + let input: Vec = (1..=3 * tile).map(|i| i as i32).collect(); + let expected: Vec = input + .chunks_exact(tile) + .flat_map(|values| { + if $up { + std::iter::once(0) + .chain(values[..tile - 1].iter().copied()) + .collect::>() + } else { + values[1..] + .iter() + .copied() + .chain(std::iter::once(0)) + .collect::>() + } + }) + .collect(); + let source = DeviceBuffer::from_host(&stream, &input)?; + let target = DeviceBuffer::from_host(&stream, &vec![-1_i32; input.len()])?; + let plan = + module.prepare_shuffle_tile::<$items, $up>(LaunchConfig1D::auto_shared(3, 128))?; + let (_, target) = module + .shuffle_tile_async_owned::<$items, $up, _, _>(&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={}, up={}: {} values checked", + $items, + $up, + output.len() + ); + }}; + } + run!(1, false); + run!(4, false); + run!(1, true); + run!(4, true); + Ok(()) +} diff --git a/rust/cuda-coop/src/exchange.rs b/rust/cuda-coop/src/exchange.rs new file mode 100644 index 000000000000..bc33e5bf6f30 --- /dev/null +++ b/rust/cuda-coop/src/exchange.rs @@ -0,0 +1,197 @@ +// 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, ThreadData, ThreadGroup, + ThreadsWithinWarp, Warp, +}; +use cuda_device::library; + +mod boundary; + +/// Convert striped input into blocked output (the default Exchange mode). +/// +/// For a group of `G` threads with `N` items each, output lane `r`, slot `i` +/// reads input lane `(r * N + i) % G`, slot `(r * N + i) / G`. +#[derive(Clone, Copy, Debug, Default)] +pub struct StripedToBlocked; + +/// Convert blocked input into striped output. +/// +/// For a group of `G` threads with `N` items each, output lane `r`, slot `i` +/// reads input lane `(i * G + r) / N`, slot `(i * G + r) % N`. +#[derive(Clone, Copy, Debug, Default)] +pub struct BlockedToStriped; + +/// A sealed compile-time Exchange direction: [`StripedToBlocked`] or +/// [`BlockedToStriped`]. Neither direction changes the input payload. +pub trait ExchangeMode: sealed::Mode {} +impl sealed::Mode for StripedToBlocked { + const KIND: u8 = 0; +} +impl sealed::Mode for BlockedToStriped { + const KIND: u8 = 1; +} +impl ExchangeMode for StripedToBlocked {} +impl ExchangeMode for BlockedToStriped {} + +/// A sealed group supporting Exchange: a complete [`Block`], physical [`Warp`], +/// or logical Warp of width 1, 2, 4, 8, 16 or 32, with either exhaustive selection. +/// +/// Every selected member must participate. Sibling Warp groups may choose +/// different paths. Warp operations require an enclosing Block containing +/// complete physical warps. Block supports all three launch dimensions and uses +/// x-fastest thread rank. Non-power-of-two prefixes and partial Block groups +/// are unsupported; a descriptor alone does not prove participation. +pub trait ExchangeGroup: GroupScope + sealed::Group {} +impl ExchangeGroup for Block {} +impl sealed::Group for Block { + const WARP: bool = false; + #[inline(always)] + unsafe fn call &'static str + Copy, P: Copy, const I: usize>( + registration: R, + payload: P, + ) { + // SAFETY: the private boundary supplies live typed payloads and the + // matching recipe. The marker checks participation and implicit scratch. + unsafe { + library::checked_block_payload_call_v1::( + registration, + (), + (), + payload, + ) + } + } +} +macro_rules! warp_group { + ($scope:ty, $width:literal $(, $e:ident)?) => { + impl<$(const $e: bool)?> ExchangeGroup for $scope {} + impl<$(const $e: bool)?> sealed::Group for $scope { + const WARP: bool = true; + #[inline(always)] + unsafe fn call &'static str + Copy, P: Copy, const I: usize>( + registration: R, payload: P, + ) { + // SAFETY: the private boundary supplies the matching recipe and + // live payloads; the genuine marker checks this canonical width. + unsafe { library::checked_warp_payload_call_v1::(registration, (), payload) } + } + } + }; +} +warp_group!(Warp, 32); +warp_group!(ThreadsWithinWarp<1, E>, 1, E); +warp_group!(ThreadsWithinWarp<2, E>, 2, E); +warp_group!(ThreadsWithinWarp<4, E>, 4, E); +warp_group!(ThreadsWithinWarp<8, E>, 8, E); +warp_group!(ThreadsWithinWarp<16, E>, 16, E); +warp_group!(ThreadsWithinWarp<32, E>, 32, E); + +/// Static direction for [`exchange`]. [`Default`] selects [`StripedToBlocked`]. +/// +/// `M` is an [`ExchangeMode`], selected by the consuming [`mode`](Self::mode) +/// builder. There are no count, algorithm, rank or explicit-storage options. +#[derive(Clone, Copy, Debug)] +pub struct ExchangeOptions { + mode: M, +} +impl Default for ExchangeOptions { + fn default() -> Self { + Self { + mode: StripedToBlocked, + } + } +} +impl ExchangeOptions { + /// Select the compile-time layout conversion. + /// + /// ``` + /// use cuda_coop as coop; + /// let _: coop::ExchangeOptions = + /// coop::ExchangeOptions::default().mode(coop::BlockedToStriped); + /// ``` + pub fn mode(self, mode: M2) -> ExchangeOptions { + ExchangeOptions { mode } + } +} + +/// Rearrange a group's per-thread values into a fresh payload of the same type, +/// extent and alignment, preserving every input value and validity bit. +/// +/// `group` selects an [`ExchangeGroup`]. `input` contains `ITEMS_PER_THREAD` +/// initialized slots per member, and `options` selects an [`ExchangeMode`]. +/// The default converts striped input into blocked output. Blocked order gives +/// lane `r`, slot `i` logical index `r * ITEMS_PER_THREAD + i`; striped order +/// gives it `i * group.count() + r`. Warp conversions act independently in each +/// canonical group. All members use the same element type and item count. +/// +/// `T` is inferred from input and supports all ten [`Numeric`] types. +/// `ITEMS_PER_THREAD` must be positive. `ALIGN` retains the input's supported +/// minimum alignment in the fresh result. Shared input and ordinary mutable +/// reborrows are accepted; scalars and owned-by-value payloads are unsupported. +/// Every returned slot is initialized, and changing the result cannot change +/// the input. The operation rearranges per-thread values; it does not load or store a +/// global-memory tile. +/// +/// Every selected member must participate. Uninitialized input traps before +/// native numeric reads. Oxide plans implicit shared scratch using the actual +/// CUB layout and synchronizes before its reuse. This operation does not promise +/// global-memory completion: use the group's `sync()` for dependent global +/// accesses. Unsupported groups or modes fail compilation. Host code can check +/// these signatures; execution requires an Oxide-compiled CUDA kernel. +/// +/// This same-extent signature needs no `generic_const_exprs` feature in callers. +/// +/// ```no_run +/// use cuda_coop as coop; +/// fn blocked( +/// group: coop::ThreadGroup, input: &coop::ThreadData, +/// ) -> coop::ThreadData { +/// coop::exchange(group, input, Default::default()) +/// } +/// ``` +/// +/// A scalar is not a one-item payload: +/// ```compile_fail,E0308 +/// use cuda_coop as coop; +/// fn invalid() { +/// coop::exchange(coop::this_block(), &3_i32, Default::default()); +/// } +/// ``` +/// +/// Non-power-of-two Warp prefixes do not support Exchange: +/// ```compile_fail,E0277 +/// use cuda_coop as coop; +/// fn invalid(values: &coop::ThreadData) { +/// coop::exchange(coop::this_warp().group_by_non_exhaustive::<17>(), values, +/// Default::default()); +/// } +/// ``` +#[inline(always)] +pub fn exchange( + _group: ThreadGroup, + input: &ThreadData, + options: ExchangeOptions, +) -> ThreadData +where + T: Numeric, + G: ExchangeGroup, + M: ExchangeMode, + Alignment: SupportedAlignment, +{ + boundary::exchange::(input, options) +} + +mod sealed { + pub trait Mode { + const KIND: u8; + } + pub trait Group { + const WARP: bool; + unsafe fn call &'static str + Copy, P: Copy, const I: usize>( + registration: R, + payload: P, + ); + } +} diff --git a/rust/cuda-coop/src/exchange/boundary.rs b/rust/cuda-coop/src/exchange/boundary.rs new file mode 100644 index 000000000000..63c8389e74d3 --- /dev/null +++ b/rust/cuda-coop/src/exchange/boundary.rs @@ -0,0 +1,69 @@ +// Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +use super::{ExchangeGroup, ExchangeMode, ExchangeOptions}; +use crate::{Alignment, Numeric, SupportedAlignment, ThreadData}; +use cuda_device::library; + +macro_rules! registrations { + ($(($name:ident, $variable:literal)),+ $(,)?) => {$ ( + fn $name() -> &'static str { library::registration!($variable) } + )+}; +} +registrations!( + ( + block_striped_to_blocked, + "COOP_EXCHANGE_BLOCK_STRIPED_TO_BLOCKED" + ), + ( + block_blocked_to_striped, + "COOP_EXCHANGE_BLOCK_BLOCKED_TO_STRIPED" + ), + ( + warp_striped_to_blocked, + "COOP_EXCHANGE_WARP_STRIPED_TO_BLOCKED" + ), + ( + warp_blocked_to_striped, + "COOP_EXCHANGE_WARP_BLOCKED_TO_STRIPED" + ), +); + +struct Selection(core::marker::PhantomData<(G, M)>); +impl Selection { + const KIND: u8 = 2 * G::WARP as u8 + M::KIND; +} + +#[inline(always)] +pub(super) fn exchange( + input: &ThreadData, + options: ExchangeOptions, +) -> ThreadData +where + T: Numeric, + G: ExchangeGroup, + M: ExchangeMode, + Alignment: SupportedAlignment, +{ + let mut output = ThreadData::::new(); + let (input_pointer, input_bits) = input.block_io_read_parts(); + let (output_pointer, output_bits) = output.block_io_parts(); + let payload = (input_pointer, output_pointer, input_bits, output_bits); + let _mode = options.mode; + // SAFETY: input and fresh output are distinct live per-thread buffers with + // the same numeric type and extent and their actual alignment. The typed + // provider preflights every input bit before numeric reads, preserves input, + // and initializes each output value before marking its validity bit. No + // pointer escapes the synchronous call. Genuine markers check participation + // and compiler-planned implicit scratch allocation and reuse. + unsafe { + match Selection::::KIND { + 0 => G::call::(block_striped_to_blocked, payload), + 1 => G::call::(block_blocked_to_striped, payload), + 2 => G::call::(warp_striped_to_blocked, payload), + 3 => G::call::(warp_blocked_to_striped, payload), + _ => panic!(), + } + } + output +} diff --git a/rust/cuda-coop/src/lib.rs b/rust/cuda-coop/src/lib.rs index 7364cdf928f0..da3458e2e44d 100644 --- a/rust/cuda-coop/src/lib.rs +++ b/rust/cuda-coop/src/lib.rs @@ -54,6 +54,11 @@ //! extent and explicit striped or blocked result ownership. The defining crate //! and callers instantiating these calculated signatures enable unstable //! `generic_const_exprs`; ordinary non-calculated consumers remain feature-free. +//! [`exchange`] converts between blocked and striped layouts on complete Blocks +//! and physical or power-of-two logical Warp groups. It preserves input and +//! returns a fresh payload with the same type, extent and alignment. +//! [`shuffle`] shifts a complete block's values by one item, preserving input +//! and returning the same payload shape with one invalid boundary slot. //! Host payload tests need no GPU; //! executing hardware queries requires an Oxide-compiled CUDA kernel. @@ -62,16 +67,21 @@ #![allow(incomplete_features)] mod alignment; +mod exchange; mod group; mod load_store; mod reduce; mod reduce_batched; mod scan; +mod shuffle; mod temp_storage; mod thread_data; pub use alignment::{Alignment, SupportedAlignment}; pub use cuda_device::{CollectiveTiles, DisjointSlice}; +pub use exchange::{ + BlockedToStriped, ExchangeGroup, ExchangeMode, ExchangeOptions, StripedToBlocked, exchange, +}; pub use load_store::{ Algorithm, Direct, LoadOptions, LoadSource, LoadStoreGroup, StoreOptions, StoreValue, Striped, Transpose, Vectorize, WarpTranspose, WarpTransposeTimesliced, load, store, @@ -90,6 +100,7 @@ pub use scan::{ ScanOptions, ScanValue, Sum, SumScanOptions, WarpScans, exclusive_scan, exclusive_sum, inclusive_scan, inclusive_sum, scan, }; +pub use shuffle::{Down, ShuffleMode, ShuffleOptions, Up, shuffle}; pub use temp_storage::{StorageReservation, TempStorage, TempStorageLike}; pub use thread_data::{Numeric, ThreadData, ThreadDataLike}; diff --git a/rust/cuda-coop/src/shuffle.rs b/rust/cuda-coop/src/shuffle.rs new file mode 100644 index 000000000000..b55ae77831b7 --- /dev/null +++ b/rust/cuda-coop/src/shuffle.rs @@ -0,0 +1,139 @@ +// Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +use crate::{Alignment, Block, Numeric, SupportedAlignment, ThreadData, ThreadGroup}; + +mod boundary; + +/// Read the next item in blocked order (the default Shuffle direction). +/// +/// The last item of the last thread has no successor and remains invalid. +#[derive(Clone, Copy, Debug, Default)] +pub struct Down; + +/// Read the previous item in blocked order. +/// +/// The first item of the first thread has no predecessor and remains invalid. +#[derive(Clone, Copy, Debug, Default)] +pub struct Up; + +/// A sealed compile-time Shuffle direction: [`Down`] or [`Up`]. +/// Neither direction changes the input or wraps around the block boundary. +pub trait ShuffleMode: sealed::Mode {} +impl sealed::Mode for Down { + const KIND: u8 = 0; +} +impl sealed::Mode for Up { + const KIND: u8 = 1; +} +impl ShuffleMode for Down {} +impl ShuffleMode for Up {} + +/// Static options for [`shuffle`]. [`Default`] selects [`Down`] at distance one. +/// +/// `M` is a [`ShuffleMode`], selected with [`mode`](Self::mode). The only +/// supported distance is one, optionally written as `.distance::<1>()`. +/// There are no count, fill, algorithm or explicit-storage options. +#[derive(Clone, Copy, Debug)] +pub struct ShuffleOptions { + mode: M, +} +impl Default for ShuffleOptions { + fn default() -> Self { + Self { mode: Down } + } +} +impl ShuffleOptions { + /// Select the compile-time direction, preserving the distance of one. + /// + /// ``` + /// use cuda_coop as coop; + /// let _: coop::ShuffleOptions = + /// coop::ShuffleOptions::default().mode(coop::Up).distance::<1>(); + /// ``` + pub fn mode(self, mode: M2) -> ShuffleOptions { + ShuffleOptions { mode } + } + + /// Specify the compile-time distance, which must be exactly one. + /// + /// `DISTANCE` is a const generic, not a runtime value. Instantiating this + /// builder with any other distance fails compilation. It may appear before + /// or after [`mode`](Self::mode). + /// + /// ```compile_fail,E0080 + /// use cuda_coop as coop; + /// let _ = coop::ShuffleOptions::default().distance::<2>(); + /// ``` + #[inline(always)] + pub fn distance(self) -> Self { + let () = Distance::::VALID; + self + } +} + +struct Distance; +impl Distance { + const VALID: () = assert!(DISTANCE == 1, "Shuffle supports only distance one"); +} + +/// Shift a complete block's values by one item into a fresh partial payload, +/// preserving every input value and validity bit. +/// +/// `group` must be the complete [`Block`]. `input` contains +/// `ITEMS_PER_THREAD` initialized items per thread in blocked order: thread +/// rank `r`, item `i` has index `r * ITEMS_PER_THREAD + i`. Rank is x-fastest +/// across the block's three launch dimensions. All threads must participate +/// with the same type, extent and direction. `options` defaults to [`Down`]. +/// +/// Down reads the next item and leaves the final block item invalid. [`Up`] +/// reads the previous item and leaves the first block item invalid. Every +/// other result slot is initialized. The boundary does not wrap or fill; +/// [`ThreadData::get`] returns `None` for that slot. A one-thread, one-item +/// block returns an entirely invalid payload. Use [`ThreadData::write`] to +/// initialize an exposed edge before a full [`crate::store`]. +/// +/// `T` is inferred from input and supports all ten [`Numeric`] types. +/// `ITEMS_PER_THREAD` must be positive; `ALIGN` preserves the input's supported +/// minimum alignment. Shared input and ordinary mutable reborrows are accepted. +/// The fresh result has the same `T`, extent and alignment, and changing it +/// cannot change the input. This signature needs no `generic_const_exprs` +/// feature in consumers. Scalars, Warp groups and partial Blocks are unsupported. +/// +/// Every input slot must be initialized, including the one that does not +/// contribute to a valid output; otherwise execution traps before native +/// numeric reads. Oxide plans implicit shared scratch from the actual CUB +/// layout and synchronizes before reuse. Shuffle does not promise global-memory +/// completion: use the block's `sync()` for dependent global accesses. +/// Unsupported groups, modes or distances fail compilation. Execution requires +/// an Oxide-compiled CUDA kernel; host code can check the signatures. +/// +/// ```no_run +/// use cuda_coop as coop; +/// fn next( +/// block: coop::ThreadGroup, +/// input: &coop::ThreadData, +/// ) -> coop::ThreadData +/// where coop::Alignment: coop::SupportedAlignment { +/// coop::shuffle(block, input, Default::default()) +/// } +/// ``` +#[inline(always)] +pub fn shuffle( + _group: ThreadGroup, + input: &ThreadData, + options: ShuffleOptions, +) -> ThreadData +where + T: Numeric, + M: ShuffleMode, + Alignment: SupportedAlignment, +{ + boundary::shuffle(input, options) +} + +mod sealed { + pub trait Mode { + const KIND: u8; + } +} diff --git a/rust/cuda-coop/src/shuffle/boundary.rs b/rust/cuda-coop/src/shuffle/boundary.rs new file mode 100644 index 000000000000..bcb39cd00d11 --- /dev/null +++ b/rust/cuda-coop/src/shuffle/boundary.rs @@ -0,0 +1,51 @@ +// Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +use super::{ShuffleMode, ShuffleOptions}; +use crate::{Alignment, Numeric, SupportedAlignment, ThreadData}; +use cuda_device::library; + +fn down() -> &'static str { + library::registration!("COOP_SHUFFLE_DOWN") +} +fn up() -> &'static str { + library::registration!("COOP_SHUFFLE_UP") +} + +#[inline(always)] +pub(super) fn shuffle( + input: &ThreadData, + options: ShuffleOptions, +) -> ThreadData +where + T: Numeric, + M: ShuffleMode, + Alignment: SupportedAlignment, +{ + let mut output = ThreadData::::new(); + let (input_pointer, input_bits) = input.block_io_read_parts(); + let (output_pointer, output_bits) = output.block_io_parts(); + let payload = (input_pointer, output_pointer, input_bits, output_bits); + let _mode = options.mode; + // SAFETY: the input and fresh output are distinct live per-thread buffers + // of the same numeric type, extent and alignment. The typed provider checks + // all input bits before reading numbers, preserves input, and only copies + // initialized CUB outputs before marking their bits. The exposed edge and + // padding bits remain invalid. Pointers do not escape the synchronous call. + // The genuine marker checks full-block participation and planned scratch. + unsafe { + match M::KIND { + 0 => library::checked_block_payload_call_v1::( + down, + (), + (), + payload, + ), + 1 => { + library::checked_block_payload_call_v1::(up, (), (), payload) + } + _ => panic!(), + } + } + output +} diff --git a/rust/cuda-coop/tests/check_load_store.py b/rust/cuda-coop/tests/check_load_store.py index 5794576d0cdc..e75be5e0a10a 100644 --- a/rust/cuda-coop/tests/check_load_store.py +++ b/rust/cuda-coop/tests/check_load_store.py @@ -1149,6 +1149,175 @@ ), ) +# Exchange preserves extent/alignment and input ownership. Fresh output does not +# establish participation or complete an unrelated global-memory access. +exchange = "values=coop::exchange(block,&values,Default::default());" +exchange_warp = "let warp=coop::this_warp().group_by_non_exhaustive::<8>();" +cases.update( + exchange_generic_defaults=( + load + + "fn forward(" + "group:coop::ThreadGroup,input:&coop::ThreadData)->coop::ThreadData " + "where coop::Alignment:coop::SupportedAlignment {" + "coop::exchange(group,input,Default::default())}" + "fn selected(" + "group:coop::ThreadGroup,input:&coop::ThreadData,options:coop::ExchangeOptions)" + "->coop::ThreadData where coop::Alignment:coop::SupportedAlignment {" + "coop::exchange(group,input,options)}" + "values=forward(block,&values);" + "values=selected(block,&mut values,coop::ExchangeOptions::default());" + "let mut one=coop::ThreadData::<_,1>::new();one.write(0,7_i32);" + "let _one=forward(block,&mut one);" + "let mut aligned=coop::ThreadData::<_,4,16>::new();" + "aligned.write(0,0_i32);aligned.write(1,1);aligned.write(2,2);aligned.write(3,3);" + "let aligned_result:coop::ThreadData<_,4,16>=selected(block,&mut aligned," + "coop::ExchangeOptions::default().mode(coop::BlockedToStriped));" + + store.replace("&values", "&aligned_result"), + None, + ), + exchange_warp_siblings=( + load + + "values=coop::exchange(coop::this_warp(),&values,Default::default());" + + exchange_warp + + "if warp.rank_at::()&1==0 {" + "values=coop::exchange(warp,&values,Default::default());" + "values=coop::exchange(warp,&values,coop::ExchangeOptions::default().mode(coop::BlockedToStriped));}" + + store, + None, + ), + exchange_block_partial_block=(load + exchange + store, None), + exchange_block_partition=( + load + + "let _result=coop::exchange(block.group_by::<2>(),&values,Default::default());", + "ExchangeGroup", + ), + exchange_prefix17=( + load + "let group=coop::this_warp().group_by_non_exhaustive::<17>();" + "if group.is_member()!=0 {let _result=coop::exchange(group,&values,Default::default());}", + "ExchangeGroup", + ), + exchange_bad_mode=( + load + "struct StripedToBlocked;" + "let _result=coop::exchange(block,&values,coop::ExchangeOptions::default().mode(StripedToBlocked));", + "ExchangeMode", + ), + exchange_block_lane_divergence=( + load + "if thread::threadIdx_x()&1==0 {" + exchange + "}", + "lane-varying", + ), + exchange_warp_lane_divergence=( + load + + exchange_warp + + "if thread::threadIdx_x()&1==0 {let _result=coop::exchange(warp,&values,Default::default());}", + "lane-varying", + ), + exchange_warp_partial_block=( + load + + "let _result=coop::exchange(coop::this_warp(),&values,Default::default());", + "complete physical warps", + ), + exchange_result_control=( + load + "let result=coop::exchange(block,&values,Default::default());" + "if result.get(0)==Some(0) {" + exchange + "}", + "lane-varying", + ), + exchange_read_epoch=( + reload + exchange + store, + "requires a barrier covering both accesses", + ), + exchange_read_completed=(reload + exchange + sync + store, None), +) + +# Shuffle preserves its input but returns one invalid edge. Neither result +# values nor validity prove participation or unrelated global-memory completion. +shuffle = "let result=coop::shuffle(block,&values,Default::default());" +cases.update( + shuffle_generic_defaults=( + load + "fn forward(" + "block:coop::ThreadGroup,input:&coop::ThreadData)->coop::ThreadData " + "where coop::Alignment:coop::SupportedAlignment {" + "coop::shuffle(block,input,Default::default())}" + "fn selected(" + "block:coop::ThreadGroup,input:&coop::ThreadData,options:coop::ShuffleOptions)" + "->coop::ThreadData where coop::Alignment:coop::SupportedAlignment {" + "coop::shuffle(block,input,options)}" + "let default_result=forward(block,&values);" + "let explicit_result=selected(block,&mut values,coop::ShuffleOptions::default().distance::<1>().mode(coop::Down));" + "let mut one=coop::ThreadData::<_,1>::new();one.write(0,7_i32);" + "let one_result=forward(block,&mut one);" + "let mut aligned=coop::ThreadData::<_,4,16>::new();" + "aligned.write(0,0_i32);aligned.write(1,1);aligned.write(2,2);aligned.write(3,3);" + "let mut aligned_result:coop::ThreadData<_,4,16>=selected(block,&mut aligned," + "coop::ShuffleOptions::default().mode(coop::Up).distance::<1>());" + "aligned_result.write(0,default_result.get(0).unwrap_or(0));" + "aligned_result.write(1,explicit_result.get(0).unwrap_or(0));" + "aligned_result.write(2,one_result.get(0).unwrap_or(0));" + + store.replace("&values", "&aligned_result"), + None, + ), + shuffle_block_partial_block=( + load + + shuffle + + "let mut repaired=result;if repaired.get(3).is_none(){repaired.write(3,0);}" + + store.replace("&values", "&repaired"), + None, + ), + shuffle_distance_zero=( + load + + "let result=coop::shuffle(block,&values,coop::ShuffleOptions::default().distance::<0>());" + + store, + "Shuffle supports only distance one", + ), + shuffle_distance_two=( + load + + "let result=coop::shuffle(block,&values,coop::ShuffleOptions::default().distance::<2>());" + + store, + "Shuffle supports only distance one", + ), + shuffle_warp_group=( + load + + "let result=coop::shuffle(coop::this_warp(),&values,Default::default());", + "mismatched types", + ), + shuffle_block_partition=( + load + + "let result=coop::shuffle(block.group_by::<2>(),&values,Default::default());", + "mismatched types", + ), + shuffle_scalar=( + load + "let result=coop::shuffle(block,&3_i32,Default::default());", + "mismatched types", + ), + shuffle_bad_mode=( + load + "struct Up;" + "let result=coop::shuffle(block,&values,coop::ShuffleOptions::default().mode(Up));", + "ShuffleMode", + ), + shuffle_lane_divergence=( + load + "if thread::threadIdx_x()&1==0 {" + shuffle + "}", + "lane-varying", + ), + shuffle_result_control=( + load + + "let first=coop::shuffle(block,&values,coop::ShuffleOptions::default().mode(coop::Up));" + "if first.get(0).is_some() {" + shuffle + "}", + "lane-varying", + ), + shuffle_read_epoch=( + reload.replace(".valid_items(valid)", ".valid_items(valid).oob_default(0)") + + shuffle + + store, + "requires a barrier covering both accesses", + ), + shuffle_read_completed=( + reload.replace(".valid_items(valid)", ".valid_items(valid).oob_default(0)") + + shuffle + + sync + + store, + None, + ), +) + if args.case: unknown = set(args.case).difference(cases) if unknown: diff --git a/rust/cuda-coop/tests/exchange/.gitignore b/rust/cuda-coop/tests/exchange/.gitignore new file mode 100644 index 000000000000..74bffa8af4a4 --- /dev/null +++ b/rust/cuda-coop/tests/exchange/.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/exchange/Cargo.lock b/rust/cuda-coop/tests/exchange/Cargo.lock new file mode 100644 index 000000000000..7a1d039e4c75 --- /dev/null +++ b/rust/cuda-coop/tests/exchange/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-exchange-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/exchange/Cargo.toml b/rust/cuda-coop/tests/exchange/Cargo.toml new file mode 100644 index 000000000000..733fe3955711 --- /dev/null +++ b/rust/cuda-coop/tests/exchange/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-exchange-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/exchange/README.md b/rust/cuda-coop/tests/exchange/README.md new file mode 100644 index 000000000000..4e4227a8c301 --- /dev/null +++ b/rust/cuda-coop/tests/exchange/README.md @@ -0,0 +1,29 @@ +# Public Exchange fixture + +This standalone test uses the public `cuda_coop` API and the pinned Oxide +toolchain in `Cargo.toml`. It needs no generic-const-expression feature. + +The positive mode runs 14 cases: both Exchange directions for Block33/i8/N1, +Block4×3×2/i16/N4, Block64/f64/N9, physical Warp32/f32/N4, +logical Warp8/i32/N8 and Warp1/u16/N9, plus both directions with selected W8 +siblings and different group-uniform iteration counts. Every launch uses two +blocks. The host independently computes each permutation. + +Every lane checks all output values and validity, extent and out-of-bounds +access, input preservation, inverse roundtrips, and a mutated fresh result +retained across later calls. A full-block Direct Store records one observation +per lane after Warp branches reconverge. No private validity bytes or raw +output pointers are used. + +Build and run from this directory with the matching installed toolchain: + +```sh +cargo oxide build --arch sm_120 --materialize-cubin -- --release --locked +cargo oxide run --arch sm_120 --materialize-cubin +``` + +The executable also accepts `invalid-input`. Run that mode in a separate +process: physical lane 31 leaves slot 3 unwritten before Exchange. An actual +CUDA launch failure is expected; successful return or an unrelated host error +does not qualify the trap. Positive sanitizer runs are separate from this +intentional failure. diff --git a/rust/cuda-coop/tests/exchange/src/main.rs b/rust/cuda-coop/tests/exchange/src/main.rs new file mode 100644 index 000000000000..c39b40c5ae9f --- /dev/null +++ b/rust/cuda-coop/tests/exchange/src/main.rs @@ -0,0 +1,401 @@ +// 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; + +const SKIPPED: u32 = 1 << 31; +const CHECKS: u32 = 15; + +#[cuda_module] +mod kernels { + use super::*; + + fn preserved( + values: &coop::ThreadData, + before: &coop::ThreadData, + ) -> bool { + let mut matches = values.len() == N; + let mut i = 0; + while i < N { + matches &= values.get(i).is_some() && values.get(i) == before.get(i); + i += 1; + } + matches && values.get(N).is_none() && values.get(usize::MAX).is_none() + } + + fn observe( + values: &coop::ThreadData, + expected: &[T], + offset: usize, + ) -> bool { + let mut matches = values.len() == N; + let mut i = 0; + while i < N { + matches &= values.get(i) == Some(expected[offset + i]); + i += 1; + } + matches && values.get(N).is_none() && values.get(usize::MAX).is_none() + } + + fn invoke( + group: coop::ThreadGroup, + input: &coop::ThreadData, + to_striped: bool, + ) -> coop::ThreadData { + if to_striped { + coop::exchange( + group, + input, + coop::ExchangeOptions::default().mode(coop::BlockedToStriped), + ) + } else { + coop::exchange(group, input, Default::default()) + } + } + + #[allow(clippy::too_many_arguments)] + fn exercise( + group: coop::ThreadGroup, + input: &[T], + expected: &[T], + tid: usize, + subgroup: usize, + to_striped: bool, + selective: bool, + invalid: bool, + ) -> u32 + where + T: coop::Numeric + PartialEq + Default, + G: coop::ExchangeGroup, + { + let mut values = coop::ThreadData::<_, ITEMS_PER_THREAD>::new(); + let mut i = 0; + while i < ITEMS_PER_THREAD { + if !(invalid && tid == 31 && i == 3) { + values.write(i, input[tid * ITEMS_PER_THREAD + i]); + } + i += 1; + } + let before = values; + let mut word = SKIPPED; + if !selective || subgroup & 1 == 0 { + let iterations = if selective { + 1 + ((subgroup / 2) & 1) + } else { + 2 + }; + let mut values_match = true; + let mut input_preserved = true; + let mut independent = true; + let mut roundtrip = true; + let mut iteration = 0; + while iteration < iterations { + let mut first = invoke(group, &values, to_striped); + values_match &= observe(&first, expected, tid * ITEMS_PER_THREAD); + input_preserved &= preserved(&values, &before); + let restored = invoke(group, &first, !to_striped); + input_preserved &= observe(&first, expected, tid * ITEMS_PER_THREAD); + roundtrip &= preserved(&restored, &before); + // All host inputs are nonzero. Keep this altered result live + // across subsequent exchanges to check fresh-result independence. + first.write(0, T::default()); + let mutated = first; + let second = invoke(group, &values, to_striped); + values_match &= observe(&second, expected, tid * ITEMS_PER_THREAD); + let restored = invoke(group, &second, !to_striped); + input_preserved &= observe(&second, expected, tid * ITEMS_PER_THREAD); + roundtrip &= preserved(&restored, &before); + independent &= preserved(&first, &mutated) && first.get(0) == Some(T::default()); + input_preserved &= preserved(&values, &before); + iteration += 1; + } + word = values_match as u32 + | ((input_preserved as u32) << 1) + | ((independent as u32) << 2) + | ((roundtrip as u32) << 3); + } + word + } + + fn report( + output: &mut coop::DisjointSlice>, + word: u32, + ) { + let block = coop::this_block(); + let mut value = coop::ThreadData::<_, 1>::new(); + value.write(0, word); + coop::store( + block, + output, + &value, + coop::StoreOptions::default() + .offset(block.rank_at::() as usize * BLOCK_THREADS), + ); + } + + #[kernel] + #[launch_contract(domain = 1, block = (33, 1, 1), dynamic_shared = inferred)] + pub fn block33( + input: &[i8], + expected: &[i8], + mut output: coop::DisjointSlice>, + to_striped: bool, + ) { + let tid = (thread::blockIdx_x() * 33 + thread::threadIdx_x()) as usize; + let word = exercise::<_, _, 1>( + coop::this_block(), + input, + expected, + tid, + 0, + to_striped, + false, + false, + ); + report::<33>(&mut output, word); + } + + #[kernel] + #[launch_contract(domain = 3, block = (4, 3, 2), dynamic_shared = inferred)] + pub fn block_xyz( + input: &[i16], + expected: &[i16], + mut output: coop::DisjointSlice>, + to_striped: bool, + ) { + let rank = thread::threadIdx_x() + 4 * (thread::threadIdx_y() + 3 * thread::threadIdx_z()); + let tid = (thread::blockIdx_x() * 24 + rank) as usize; + let word = exercise::<_, _, 4>( + coop::this_block(), + input, + expected, + tid, + 0, + to_striped, + false, + false, + ); + report::<24>(&mut output, word); + } + + #[kernel] + #[launch_contract(domain = 1, block = (64, 1, 1), dynamic_shared = inferred)] + pub fn block64( + input: &[f64], + expected: &[f64], + mut output: coop::DisjointSlice>, + to_striped: bool, + ) { + let tid = (thread::blockIdx_x() * 64 + thread::threadIdx_x()) as usize; + let word = exercise::<_, _, 9>( + coop::this_block(), + input, + expected, + tid, + 0, + to_striped, + false, + false, + ); + report::<64>(&mut output, word); + } + + #[kernel] + #[launch_contract(domain = 1, block = (64, 1, 1), dynamic_shared = inferred)] + pub fn physical( + input: &[f32], + expected: &[f32], + mut output: coop::DisjointSlice>, + to_striped: bool, + invalid: bool, + ) { + let tid = (thread::blockIdx_x() * 64 + thread::threadIdx_x()) as usize; + let word = exercise::<_, _, 4>( + coop::this_warp(), + input, + expected, + tid, + 0, + to_striped, + false, + invalid, + ); + report::<64>(&mut output, word); + } + + #[kernel] + #[launch_contract(domain = 3, block = (4, 4, 4), dynamic_shared = inferred)] + pub fn logical( + input: &[T], + expected: &[T], + mut output: coop::DisjointSlice>, + to_striped: bool, + selective: bool, + ) where + T: coop::Numeric + PartialEq + Default, + coop::ThreadsWithinWarp: coop::ExchangeGroup, + { + let group = coop::this_warp().group_by::(); + let rank = thread::threadIdx_x() + 4 * (thread::threadIdx_y() + 4 * thread::threadIdx_z()); + let tid = (thread::blockIdx_x() * 64 + rank) as usize; + let subgroup = coop::this_warp().rank_at::() as usize * (32 / WIDTH) + + group.rank_at::() as usize; + let word = exercise::<_, _, ITEMS_PER_THREAD>( + group, input, expected, tid, subgroup, to_striped, selective, false, + ); + // All selected Warp work has reconverged before this Block Store. + report::<64>(&mut output, word); + } +} + +fn oracle(input: &[T], items: usize, width: usize, to_striped: bool) -> Vec { + let mut expected = input.to_vec(); + for group in 0..input.len() / (width * items) { + for lane in 0..width { + for i in 0..items { + let logical = if to_striped { + i * width + lane + } else { + lane * items + i + }; + let (source_lane, source_slot) = if to_striped { + (logical / items, logical % items) + } else { + (logical % width, logical / width) + }; + expected[(group * width + lane) * items + i] = + input[(group * width + source_lane) * items + source_slot]; + } + } + } + expected +} + +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: typed declarations and embedded kernels are compiled together. + let module = unsafe { kernels::load_async(0)? }; + let invalid = match std::env::args().nth(1).as_deref() { + None => false, + Some("invalid-input") => true, + Some(other) => panic!("unknown test mode {other}"), + }; + if invalid { + let input = DeviceBuffer::from_host(&stream, &[1_f32; 128 * 4])?; + let expected = DeviceBuffer::from_host(&stream, &[1_f32; 128 * 4])?; + let output = DeviceBuffer::from_host(&stream, &[0_u32; 128])?; + let plan = module.prepare_physical(LaunchConfig1D::auto_shared(2, 64))?; + module + .physical_async_owned(&plan, input, expected, output, false, true) + .sync()?; + stream.synchronize()?; + panic!("unwritten Exchange input unexpectedly completed"); + } + let mut cases = 0; + let mut reports = 0; + macro_rules! fixed { + ($ty:ty, $items:literal, $threads:literal, $width:literal, $prepare:ident, $launch:ident, $config:expr $(, $extra:expr)?) => {{ + let input: Vec<$ty> = (1..=2 * $threads * $items).map(|i| i as $ty).collect(); + for to_striped in [false, true] { + let expected = oracle(&input, $items, $width, to_striped); + let source = DeviceBuffer::from_host(&stream, &input)?; + let expected = DeviceBuffer::from_host(&stream, &expected)?; + let target = DeviceBuffer::from_host(&stream, &[0_u32; 2 * $threads])?; + let plan = module.$prepare($config)?; + let (_, _, target) = module.$launch(&plan, source, expected, target, to_striped $(, $extra)?).sync()?; + assert_eq!(target.to_host_vec(&stream)?, vec![CHECKS; 2 * $threads], "{} striped{}", stringify!($prepare), to_striped); + cases += 1; + reports += 2 * $threads; + } + }}; + } + fixed!( + i8, + 1, + 33, + 33, + prepare_block33, + block33_async_owned, + LaunchConfig1D::auto_shared(2, 33) + ); + fixed!( + i16, + 4, + 24, + 24, + prepare_block_xyz, + block_xyz_async_owned, + LaunchConfig3D::auto_shared((2, 1, 1), (4, 3, 2)) + ); + fixed!( + f64, + 9, + 64, + 64, + prepare_block64, + block64_async_owned, + LaunchConfig1D::auto_shared(2, 64) + ); + fixed!( + f32, + 4, + 64, + 32, + prepare_physical, + physical_async_owned, + LaunchConfig1D::auto_shared(2, 64), + false + ); + macro_rules! logical { + ($ty:ty, $items:literal, $width:literal, $selective:expr) => {{ + let input: Vec<$ty> = (1..=128 * $items).map(|i| i as $ty).collect(); + for to_striped in [false, true] { + let expected = oracle(&input, $items, $width, to_striped); + let words: Vec = (0..128) + .map(|tid| { + if $selective && (tid / $width) & 1 != 0 { + SKIPPED + } else { + CHECKS + } + }) + .collect(); + let source = DeviceBuffer::from_host(&stream, &input)?; + let expected = DeviceBuffer::from_host(&stream, &expected)?; + let target = DeviceBuffer::from_host(&stream, &[0_u32; 128])?; + let plan = module.prepare_logical::<$ty, $items, $width>( + LaunchConfig3D::auto_shared((2, 1, 1), (4, 4, 4)), + )?; + let (_, _, target) = module + .logical_async_owned::<$ty, $items, $width, _, _, _>( + &plan, source, expected, target, to_striped, $selective, + ) + .sync()?; + assert_eq!( + target.to_host_vec(&stream)?, + words, + "{} N{} W{} striped{} selective{}", + stringify!($ty), + $items, + $width, + to_striped, + $selective + ); + cases += 1; + reports += 128; + } + }}; + } + logical!(i32, 8, 8, false); + logical!(u16, 9, 1, false); + logical!(i32, 4, 8, true); + println!("PASS: {cases} public Exchange cases, {reports} lane reports"); + Ok(()) +} diff --git a/rust/cuda-coop/tests/shuffle/.gitignore b/rust/cuda-coop/tests/shuffle/.gitignore new file mode 100644 index 000000000000..74bffa8af4a4 --- /dev/null +++ b/rust/cuda-coop/tests/shuffle/.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/shuffle/Cargo.lock b/rust/cuda-coop/tests/shuffle/Cargo.lock new file mode 100644 index 000000000000..1b34e139aed2 --- /dev/null +++ b/rust/cuda-coop/tests/shuffle/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-shuffle-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/shuffle/Cargo.toml b/rust/cuda-coop/tests/shuffle/Cargo.toml new file mode 100644 index 000000000000..5326fccc28e6 --- /dev/null +++ b/rust/cuda-coop/tests/shuffle/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-shuffle-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/shuffle/src/main.rs b/rust/cuda-coop/tests/shuffle/src/main.rs new file mode 100644 index 000000000000..3ad2684fac3a --- /dev/null +++ b/rust/cuda-coop/tests/shuffle/src/main.rs @@ -0,0 +1,362 @@ +// 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; + +const CHECKS: u32 = 7; + +#[cuda_module] +mod kernels { + use super::*; + + fn same( + left: &coop::ThreadData, + right: &coop::ThreadData, + ) -> bool + where + coop::Alignment: coop::SupportedAlignment, + { + let mut matches = left.len() == N; + let mut i = 0; + while i < N { + matches &= left.get(i) == right.get(i); + i += 1; + } + matches && left.get(N).is_none() && left.get(usize::MAX).is_none() + } + + fn shifted( + values: &coop::ThreadData, + expected: &[T], + tid: usize, + rank: usize, + threads: usize, + up: bool, + ) -> bool + where + coop::Alignment: coop::SupportedAlignment, + { + let mut matches = values.len() == N; + let mut i = 0; + while i < N { + let valid = if up { + rank != 0 || i != 0 + } else { + rank != threads - 1 || i != N - 1 + }; + matches &= if valid { + values.get(i) == Some(expected[tid * N + i]) + } else { + values.get(i).is_none() + }; + i += 1; + } + matches && values.get(N).is_none() && values.get(usize::MAX).is_none() + } + + fn invoke( + block: coop::ThreadGroup, + values: &coop::ThreadData, + up: bool, + ) -> coop::ThreadData + where + coop::Alignment: coop::SupportedAlignment, + { + if up { + coop::shuffle( + block, + values, + coop::ShuffleOptions::default().mode(coop::Up), + ) + } else { + coop::shuffle(block, values, Default::default()) + } + } + + #[allow(clippy::too_many_arguments)] + fn exercise( + input: &[T], + expected: &[T], + output: &mut coop::DisjointSlice<'_, T, coop::CollectiveTiles>, + reports: &mut coop::DisjointSlice<'_, u32, coop::CollectiveTiles<1>>, + rank: usize, + up: bool, + invalid: u32, + ) where + T: coop::Numeric + PartialEq + Default, + coop::Alignment: coop::SupportedAlignment, + { + let block = coop::this_block(); + let block_index = block.rank_at::() as usize; + let tid = block_index * B + rank; + let mut values = coop::ThreadData::<_, N, A>::new(); + let mut i = 0; + while i < N { + // These discarded values are still required by full-input preflight. + let missing = (invalid == 1 && rank == 0 && i == 0) + || (invalid == 2 && rank == B - 1 && i == N - 1); + if !missing { + values.write(i, input[tid * N + i]); + } + i += 1; + } + let before = values; + let mut first = invoke(block, &values, up); + let mut values_match = shifted(&first, expected, tid, rank, B, up); + let mut input_preserved = same(&values, &before); + // All input values are nonzero. Keep this changed result live across + // scratch reuse, including the entirely invalid Block1/N1 result. + first.write(0, T::default()); + let mutated = first; + let mut second = invoke(block, &values, up); + values_match &= shifted(&second, expected, tid, rank, B, up); + input_preserved &= same(&values, &before); + let independent = same(&first, &mutated) && first.get(0) == Some(T::default()); + let word = + values_match as u32 | ((input_preserved as u32) << 1) | ((independent as u32) << 2); + let unrepaired = second; + if invalid != 4 { + let mut i = 0; + while i < N { + if second.get(i).is_none() { + second.write(i, T::default()); + } + i += 1; + } + } + coop::store( + block, + output, + &second, + coop::StoreOptions::default().offset(block_index * B * N), + ); + let mut report = coop::ThreadData::<_, 1>::new(); + // If Index were to allow the invalid read, preserve its comparison in + // global memory and complete normally; only the host then reports the + // missing trap. A value-dependent assertion could falsely pass it. + let observed = if invalid == 3 { + word | (((unrepaired[0] == input[tid * N]) as u32) << 3) + } else { + word + }; + report.write(0, observed); + coop::store( + block, + reports, + &report, + coop::StoreOptions::default().offset(block_index * B), + ); + } + + #[kernel] + #[launch_contract(domain = 1, block = (33, 1, 1), dynamic_shared = inferred)] + pub fn block33( + input: &[i8], + expected: &[i8], + mut output: coop::DisjointSlice>, + mut reports: coop::DisjointSlice>, + up: bool, + invalid: u32, + ) { + exercise::<_, 1, 1, 33>( + input, + expected, + &mut output, + &mut reports, + thread::threadIdx_x() as usize, + up, + invalid, + ); + } + + #[kernel] + #[launch_contract(domain = 3, block = (4, 3, 2), dynamic_shared = inferred)] + pub fn block_xyz( + input: &[i32], + expected: &[i32], + mut output: coop::DisjointSlice>, + mut reports: coop::DisjointSlice>, + up: bool, + ) { + let rank = thread::threadIdx_x() + 4 * (thread::threadIdx_y() + 3 * thread::threadIdx_z()); + exercise::<_, 4, 16, 24>( + input, + expected, + &mut output, + &mut reports, + rank as usize, + up, + 0, + ); + } + + #[kernel] + #[launch_contract(domain = 1, block = (64, 1, 1), dynamic_shared = inferred)] + pub fn block64( + input: &[f64], + expected: &[f64], + mut output: coop::DisjointSlice>, + mut reports: coop::DisjointSlice>, + up: bool, + ) { + exercise::<_, 9, 1, 64>( + input, + expected, + &mut output, + &mut reports, + thread::threadIdx_x() as usize, + up, + 0, + ); + } + + #[kernel] + #[launch_contract(domain = 1, block = (1, 1, 1), dynamic_shared = inferred)] + pub fn single_item( + input: &[i32], + expected: &[i32], + mut output: coop::DisjointSlice>, + mut reports: coop::DisjointSlice>, + up: bool, + ) { + exercise::<_, 1, 1, 1>(input, expected, &mut output, &mut reports, 0, up, 0); + } + + #[kernel] + #[launch_contract(domain = 1, block = (1, 1, 1), dynamic_shared = inferred)] + pub fn single_thread( + input: &[u16], + expected: &[u16], + mut output: coop::DisjointSlice>, + mut reports: coop::DisjointSlice>, + up: bool, + ) { + exercise::<_, 4, 1, 1>(input, expected, &mut output, &mut reports, 0, up, 0); + } +} + +fn oracle(input: &[T], tile: usize, up: bool) -> Vec { + input + .chunks_exact(tile) + .flat_map(|values| { + if up { + std::iter::once(T::default()) + .chain(values[..tile - 1].iter().copied()) + .collect::>() + } else { + values[1..] + .iter() + .copied() + .chain(std::iter::once(T::default())) + .collect::>() + } + }) + .collect() +} + +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}; + + let invalid = match std::env::args().nth(1).as_deref() { + None => 0, + Some("missing-down-input") => 1, + Some("missing-up-input") => 2, + Some("invalid-edge-read") => 3, + Some("unrepaired-store") => 4, + Some(other) => panic!("unknown test mode {other}"), + }; + init_device_contexts(0, 1)?; + let context = with_cuda_context(0, |context| context.clone())?; + let stream = context.default_stream(); + // SAFETY: typed declarations and embedded kernels are compiled together. + let module = unsafe { kernels::load_async(0)? }; + if invalid != 0 { + let up = invalid == 2 || invalid == 3; + let input: Vec = (1..=33).collect(); + let expected = oracle(&input, 33, up); + let source = DeviceBuffer::from_host(&stream, &input)?; + let expected = DeviceBuffer::from_host(&stream, &expected)?; + let target = DeviceBuffer::from_host(&stream, &[-1_i8; 33])?; + let reports = DeviceBuffer::from_host(&stream, &[0_u32; 33])?; + let plan = module.prepare_block33(LaunchConfig1D::auto_shared(1, 33))?; + module + .block33_async_owned(&plan, source, expected, target, reports, up, invalid) + .sync()?; + stream.synchronize()?; + panic!("invalid Shuffle mode {invalid} unexpectedly completed"); + } + + let mut cases = 0; + let mut reports = 0; + let mut values = 0; + macro_rules! run { + ($ty:ty, $items:literal, $threads:literal, $prepare:ident, $launch:ident, $config:expr $(, $extra:expr)?) => {{ + let input: Vec<$ty> = (1..=3 * $threads * $items).map(|i| i as $ty).collect(); + for up in [false, true] { + let expected = oracle(&input, $threads * $items, up); + let source = DeviceBuffer::from_host(&stream, &input)?; + let expected_device = DeviceBuffer::from_host(&stream, &expected)?; + let target = DeviceBuffer::from_host(&stream, &vec![<$ty>::MAX; input.len()])?; + let words = DeviceBuffer::from_host(&stream, &[0_u32; 3 * $threads])?; + let plan = module.$prepare($config)?; + let (_, _, target, words) = module.$launch( + &plan, source, expected_device, target, words, up $(, $extra)? + ).sync()?; + assert_eq!(target.to_host_vec(&stream)?, expected, "{} up{}", stringify!($prepare), up); + assert_eq!(words.to_host_vec(&stream)?, vec![CHECKS; 3 * $threads], "{} up{}", stringify!($prepare), up); + cases += 1; + reports += 3 * $threads; + values += input.len(); + } + }}; + } + run!( + i8, + 1, + 33, + prepare_block33, + block33_async_owned, + LaunchConfig1D::auto_shared(3, 33), + 0 + ); + run!( + i32, + 4, + 24, + prepare_block_xyz, + block_xyz_async_owned, + LaunchConfig3D::auto_shared((3, 1, 1), (4, 3, 2)) + ); + run!( + f64, + 9, + 64, + prepare_block64, + block64_async_owned, + LaunchConfig1D::auto_shared(3, 64) + ); + run!( + i32, + 1, + 1, + prepare_single_item, + single_item_async_owned, + LaunchConfig1D::auto_shared(3, 1) + ); + run!( + u16, + 4, + 1, + prepare_single_thread, + single_thread_async_owned, + LaunchConfig1D::auto_shared(3, 1) + ); + println!( + "PASS: {cases} public Shuffle cases, {values} repaired values, {reports} lane reports" + ); + Ok(()) +} From 68077d5e3007ca6e5018b90f2a555a3ba2a7aa2b Mon Sep 17 00:00:00 2001 From: Trent Nelson Date: Sat, 24 Oct 2026 17:43:09 -0700 Subject: [PATCH 07/10] Add Rust Adjacent Difference and Discontinuity Provide neighboring-item transforms and head/tail flags with typed neighbor controls, payload validity and explicit scratch completion. Signed-off-by: Trent Nelson --- docs/rust/adjacent_difference.rst | 91 ++ docs/rust/discontinuity.rst | 90 ++ docs/rust/index.rst | 13 +- docs/rust/shuffle.rst | 2 +- .../src/adjacent_difference.rs | 267 +++++ rust/cuda-coop-codegen/src/discontinuity.rs | 293 +++++ rust/cuda-coop-codegen/src/lib.rs | 48 +- rust/cuda-coop-codegen/src/oxide.rs | 258 ++++- rust/cuda-coop-codegen/tests/native.rs | 253 +++- rust/cuda-coop/build.rs | 52 +- .../examples/adjacent_difference/.gitignore | 7 + .../examples/adjacent_difference/Cargo.lock | 1019 ++++++++++++++++ .../examples/adjacent_difference/Cargo.toml | 19 + .../adjacent_difference/helper/Cargo.toml | 11 + .../adjacent_difference/helper/src/lib.rs | 56 + .../examples/adjacent_difference/src/main.rs | 101 ++ .../examples/discontinuity/.gitignore | 7 + .../examples/discontinuity/Cargo.lock | 1019 ++++++++++++++++ .../examples/discontinuity/Cargo.toml | 19 + .../examples/discontinuity/helper/Cargo.toml | 11 + .../examples/discontinuity/helper/src/lib.rs | 39 + .../examples/discontinuity/src/main.rs | 128 +++ rust/cuda-coop/src/adjacent_difference.rs | 264 +++++ .../src/adjacent_difference/boundary.rs | 97 ++ rust/cuda-coop/src/discontinuity.rs | 276 +++++ rust/cuda-coop/src/discontinuity/boundary.rs | 132 +++ rust/cuda-coop/src/lib.rs | 14 + rust/cuda-coop/src/shuffle.rs | 2 +- .../tests/adjacent_difference/.gitignore | 7 + .../tests/adjacent_difference/Cargo.lock | 1011 ++++++++++++++++ .../tests/adjacent_difference/Cargo.toml | 17 + .../tests/adjacent_difference/src/main.rs | 600 ++++++++++ rust/cuda-coop/tests/check_load_store.py | 238 ++++ rust/cuda-coop/tests/discontinuity/.gitignore | 7 + rust/cuda-coop/tests/discontinuity/Cargo.lock | 1011 ++++++++++++++++ rust/cuda-coop/tests/discontinuity/Cargo.toml | 17 + .../cuda-coop/tests/discontinuity/src/main.rs | 1023 +++++++++++++++++ 37 files changed, 8497 insertions(+), 22 deletions(-) create mode 100644 docs/rust/adjacent_difference.rst create mode 100644 docs/rust/discontinuity.rst create mode 100644 rust/cuda-coop-codegen/src/adjacent_difference.rs create mode 100644 rust/cuda-coop-codegen/src/discontinuity.rs create mode 100644 rust/cuda-coop/examples/adjacent_difference/.gitignore create mode 100644 rust/cuda-coop/examples/adjacent_difference/Cargo.lock create mode 100644 rust/cuda-coop/examples/adjacent_difference/Cargo.toml create mode 100644 rust/cuda-coop/examples/adjacent_difference/helper/Cargo.toml create mode 100644 rust/cuda-coop/examples/adjacent_difference/helper/src/lib.rs create mode 100644 rust/cuda-coop/examples/adjacent_difference/src/main.rs create mode 100644 rust/cuda-coop/examples/discontinuity/.gitignore create mode 100644 rust/cuda-coop/examples/discontinuity/Cargo.lock create mode 100644 rust/cuda-coop/examples/discontinuity/Cargo.toml create mode 100644 rust/cuda-coop/examples/discontinuity/helper/Cargo.toml create mode 100644 rust/cuda-coop/examples/discontinuity/helper/src/lib.rs create mode 100644 rust/cuda-coop/examples/discontinuity/src/main.rs create mode 100644 rust/cuda-coop/src/adjacent_difference.rs create mode 100644 rust/cuda-coop/src/adjacent_difference/boundary.rs create mode 100644 rust/cuda-coop/src/discontinuity.rs create mode 100644 rust/cuda-coop/src/discontinuity/boundary.rs create mode 100644 rust/cuda-coop/tests/adjacent_difference/.gitignore create mode 100644 rust/cuda-coop/tests/adjacent_difference/Cargo.lock create mode 100644 rust/cuda-coop/tests/adjacent_difference/Cargo.toml create mode 100644 rust/cuda-coop/tests/adjacent_difference/src/main.rs create mode 100644 rust/cuda-coop/tests/discontinuity/.gitignore create mode 100644 rust/cuda-coop/tests/discontinuity/Cargo.lock create mode 100644 rust/cuda-coop/tests/discontinuity/Cargo.toml create mode 100644 rust/cuda-coop/tests/discontinuity/src/main.rs diff --git a/docs/rust/adjacent_difference.rst b/docs/rust/adjacent_difference.rst new file mode 100644 index 000000000000..4521c47664f9 --- /dev/null +++ b/docs/rust/adjacent_difference.rst @@ -0,0 +1,91 @@ +.. Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. +.. SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +Differences between neighboring values +======================================= + +``adjacent_difference`` subtracts neighboring values in a block's blocked +order. It preserves the input and returns a fresh ``ThreadData`` with the +same element type, item count and alignment. The default ``Left`` direction +subtracts the previous item; ``Right`` subtracts the next item: + +.. code-block:: rust + + let left = coop::adjacent_difference(block, &values, Default::default()); + let right = coop::adjacent_difference(block, &values, + coop::AdjacentDifferenceOptions::default().direction(coop::Right)); + +For ``[10, 20, 40, 70]``, Left produces ``[10, 10, 20, 30]`` and Right +produces ``[-10, -20, -30, 70]``. Neighbors cross thread boundaries. Without +an outside neighbor, the first Left item or last Right item is copied. + +A complete program +------------------ + +The helper keeps ``ITEMS_PER_THREAD`` generic and infers the element type +from the input slice. Load initializes the payload, Adjacent Difference +returns fresh values, and Store writes the result: + +.. literalinclude:: ../../rust/cuda-coop/examples/adjacent_difference/helper/src/lib.rs + :language: rust + :start-after: // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +The host program specializes one and four items per thread: + +.. literalinclude:: ../../rust/cuda-coop/examples/adjacent_difference/src/main.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/adjacent_difference + 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. ``auto_shared`` uses the compiler's +shared-storage plan, including the actual CUB temporary-storage requirement. + +Counts and outside neighbors +----------------------------- + +``valid_items(count)`` limits the subtraction to the first ``count`` items +of the whole block tile. Every later item is copied unchanged. With count +two, the example above produces ``[10, 10, 40, 70]`` for Left and +``[-10, 20, 40, 70]`` for Right. Count zero returns a copy of the whole input. +Counts must be block-uniform and lie between zero and the tile capacity. + +Left accepts ``tile_predecessor_item(value)`` with or without a count. +Right accepts ``tile_successor_item(value)`` only when the count is omitted. +A supplied count and successor are rejected even when the count equals the +full tile size. Outside neighbors have exactly the input element type and +must be block-uniform. Changing direction preserves the option fields; +it does not reinterpret a predecessor as a successor. + +Every input slot must be initialized, including the copied suffix and every +slot at count zero. Fill missing Load items or initialize them with ``write`` +before calling. Every result slot is initialized; changing a result cannot +change the input. Arithmetic uses native subtraction in the input type. + +Groups and temporary storage +---------------------------- + +The operation accepts complete one-, two- and three-dimensional Blocks and +all ten ``Numeric`` types. Rank is x-fastest; item count is positive and the +whole tile must fit CUB's signed 32-bit indexing range. All threads participate +with the same direction, type and extent. Scalars and Warp groups are not +supported. Same-extent generic helpers need no ``generic_const_exprs`` feature. + +Omitted storage uses managed synchronization before scratch reuse. Borrow an +explicit descriptor with ``temp_storage(&mut scratch)`` to use the policies +described in :doc:`load_store`: manual reuse requires block completion, +``auto_sync()`` supplies a trailing barrier, and exclusive storage has its +own allocation. Adjacent Difference does not itself promise completion of +unrelated global-memory accesses; use ``cuda_device::thread::sync_threads()`` +when those accesses need block completion. + +See :doc:`payload` for checked access and :doc:`groups` for participation. +Generate the API reference with ``cargo doc --workspace --no-deps`` from +``rust/`` for the full signatures and legal option combinations. diff --git a/docs/rust/discontinuity.rst b/docs/rust/discontinuity.rst new file mode 100644 index 000000000000..88575aa69727 --- /dev/null +++ b/docs/rust/discontinuity.rst @@ -0,0 +1,90 @@ +.. Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. +.. SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +Flags at changes between values +=============================== + +``discontinuity`` marks changes between neighboring items in blocked order. +It preserves the input and returns fresh ``i32`` flags: one for unequal +neighbors, zero for equal neighbors. ``Heads`` is the default and compares +each item with its predecessor. ``Tails`` compares with its successor. +``HeadsAndTails`` returns two independent payloads, heads first: + +.. code-block:: rust + + let heads = coop::discontinuity(block, &values, Default::default()); + let (heads, tails) = coop::discontinuity(block, &values, + coop::DiscontinuityOptions::default().mode(coop::HeadsAndTails)); + +For ``[2, 2, 5, 5]``, heads are ``[1, 0, 1, 0]`` and tails are +``[0, 1, 0, 1]``. Neighbors cross thread boundaries. Without an outside +neighbor, the first head and last tail are one. + +A complete program +------------------ + +The helper infers the input type and keeps ``ITEMS_PER_THREAD`` generic. +It loads a complete tile and returns both flag payloads and the preserved +input. The flag type is ``i32`` even when the input is ``f64``: + +.. literalinclude:: ../../rust/cuda-coop/examples/discontinuity/helper/src/lib.rs + :language: rust + :start-after: // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +The host program uses a ``(4, 3, 2)`` block, specializes one and four items +per thread, and checks both flag arrays and the preserved input: + +.. literalinclude:: ../../rust/cuda-coop/examples/discontinuity/src/main.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/discontinuity + 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. ``auto_shared`` uses the compiler's +plan for the actual CUB temporary-storage requirement. + +Outside neighbors and initialized input +--------------------------------------- + +Heads accepts ``tile_predecessor_item(value)``; Tails accepts +``tile_successor_item(value)``. The combined mode accepts either, both or +neither. Outside neighbors have exactly the input type and are uniform +across the block. A supplied neighbor makes its edge flag a normal +comparison: predecessor two changes the example's first head to zero; +successor five changes its last tail to zero. Builders preserve both fields +when changing mode. + +Every input slot must be initialized. There is no count option: all values, +including any fill supplied to Load, participate in comparisons. Floating +inequality treats NaNs as unequal and positive and negative zero as equal. +Every output slot is initialized. Both outputs retain the input's item +count and minimum alignment, with the natural alignment required by ``i32``. +Changing either result cannot change the input or the other result. + +Groups and temporary storage +---------------------------- + +Complete one-, two- and three-dimensional Blocks and all ten ``Numeric`` +input types are supported. Rank is x-fastest, item count is positive, and +the complete tile must fit CUB's signed 32-bit indexing range. Every thread +participates with the same mode, type and extent. Generic helpers can return +one payload or the pair without a ``generic_const_exprs`` feature. + +Omitted storage uses managed synchronization before scratch reuse. Use +``temp_storage(&mut scratch)`` for an explicit descriptor: manual reuse +requires block completion, ``auto_sync()`` supplies a trailing barrier, +and exclusive storage keeps its existing allocation policy. For unrelated +global accesses that need block completion, use +``cuda_device::thread::sync_threads()``. + +See :doc:`payload`, :doc:`groups` and :doc:`load_store` for checked access, +participation and storage. Generate the API reference with +``cargo doc --workspace --no-deps`` from ``rust/`` for the full signatures +and mode-dependent result types. diff --git a/docs/rust/index.rst b/docs/rust/index.rst index f68843d5f71a..9f10b3ff44ef 100644 --- a/docs/rust/index.rst +++ b/docs/rust/index.rst @@ -44,14 +44,17 @@ 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`, :doc:`scan`, :doc:`reduce`, -:doc:`reduce_batched`, :doc:`exchange` and :doc:`shuffle`. Scan returns fresh per-group prefixes; +:doc:`reduce_batched`, :doc:`exchange`, :doc:`shuffle`, +:doc:`adjacent_difference` and :doc:`discontinuity`. +Scan returns fresh per-group prefixes; Reduce returns a total only at each group's root. Batched reduction returns a fresh payload whose extent and valid slots follow its Warp width and output layout. Exchange converts between blocked and striped per-thread layouts. Shuffle shifts blocked values by one item, leaving the exposed edge invalid. -Scan, Reduce, batched reduction, Exchange and Shuffle preserve their input. Additional -primitive families are under development. The crate is unpublished and its -interface can change. +Discontinuity returns head flags, tail flags or both as fresh ``i32`` payloads. +Scan, Reduce, batched reduction, Exchange, Shuffle, Adjacent Difference and +Discontinuity preserve their input. Additional primitive families are under +development. The crate is unpublished and its interface can change. Build and run ------------- @@ -123,6 +126,8 @@ Read next reduce_batched exchange shuffle + adjacent_difference + discontinuity provider Generate the API reference with ``cargo doc --workspace --no-deps`` from diff --git a/docs/rust/shuffle.rst b/docs/rust/shuffle.rst index f056f1120d96..a0325fb9622d 100644 --- a/docs/rust/shuffle.rst +++ b/docs/rust/shuffle.rst @@ -82,7 +82,7 @@ Oxide checks participation and plans the actual CUB shared storage, including synchronization before reuse. Keep collective calls outside lane-specific branches and reconverge after repairing the edge, as the helper does. Shuffle does not promise completion of earlier global-memory accesses. Use -the block's ``sync()`` when dependent accesses need completion. +``cuda_device::thread::sync_threads()`` when dependent accesses need block completion. See :doc:`payload` for checked access and :doc:`groups` for participation. Generate the item reference with ``cargo doc --workspace --no-deps`` from diff --git a/rust/cuda-coop-codegen/src/adjacent_difference.rs b/rust/cuda-coop-codegen/src/adjacent_difference.rs new file mode 100644 index 000000000000..b3724dce49a8 --- /dev/null +++ b/rust/cuda-coop-codegen/src/adjacent_difference.rs @@ -0,0 +1,267 @@ +// Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +//! Validated numeric Block adjacent differences with preserved input. + +use crate::{Dtype, Error, Result}; + +/// Direction, tile extent and optional boundary value for adjacent subtraction. +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub enum AdjacentDifferenceRecipe { + /// Subtract each predecessor, copying the first tile item unchanged. + LeftFull, + /// Subtract each predecessor, including a supplied tile predecessor. + LeftFullPredecessor, + /// Subtract within a valid prefix, copying its first item and the suffix. + LeftPartial, + /// Subtract within a valid prefix using a supplied tile predecessor. + LeftPartialPredecessor, + /// Subtract each successor, copying the last tile item unchanged. + RightFull, + /// Subtract each successor, including a supplied tile successor. + RightFullSuccessor, + /// Subtract within a valid prefix, copying its last item and the suffix. + RightPartial, +} + +impl AdjacentDifferenceRecipe { + pub(crate) fn token(self) -> &'static str { + match self { + Self::LeftFull => "adjacent_difference_left_full_validity_bits_v1", + Self::LeftFullPredecessor => { + "adjacent_difference_left_full_predecessor_validity_bits_v1" + } + Self::LeftPartial => "adjacent_difference_left_partial_validity_bits_v1", + Self::LeftPartialPredecessor => { + "adjacent_difference_left_partial_predecessor_validity_bits_v1" + } + Self::RightFull => "adjacent_difference_right_full_validity_bits_v1", + Self::RightFullSuccessor => "adjacent_difference_right_full_successor_validity_bits_v1", + Self::RightPartial => "adjacent_difference_right_partial_validity_bits_v1", + } + } + + pub(crate) fn parse(operation: &str, algorithm: &str) -> Result { + if algorithm != "block" { + return Err(Error::message( + "Adjacent Difference requires a Block recipe", + )); + } + match operation { + "adjacent_difference_left_full_validity_bits_v1" => Ok(Self::LeftFull), + "adjacent_difference_left_full_predecessor_validity_bits_v1" => { + Ok(Self::LeftFullPredecessor) + } + "adjacent_difference_left_partial_validity_bits_v1" => Ok(Self::LeftPartial), + "adjacent_difference_left_partial_predecessor_validity_bits_v1" => { + Ok(Self::LeftPartialPredecessor) + } + "adjacent_difference_right_full_validity_bits_v1" => Ok(Self::RightFull), + "adjacent_difference_right_full_successor_validity_bits_v1" => { + Ok(Self::RightFullSuccessor) + } + "adjacent_difference_right_partial_validity_bits_v1" => Ok(Self::RightPartial), + _ => Err(Error::message("unsupported Adjacent Difference operation")), + } + } + + pub(crate) fn counted(self) -> bool { + matches!( + self, + Self::LeftPartial | Self::LeftPartialPredecessor | Self::RightPartial + ) + } + + pub(crate) fn has_neighbor(self) -> bool { + matches!( + self, + Self::LeftFullPredecessor | Self::LeftPartialPredecessor | Self::RightFullSuccessor + ) + } +} + +/// Validated Block adjacent subtraction with fresh, fully initialized output. +/// +/// Native calls use `(const T* input, T* output, u8* scratch, +/// const u8* input_bits, u8* output_bits)`, followed by an `i32` valid-item +/// count for partial recipes and then a value of exactly `T` for recipes with +/// a tile predecessor or successor. Controls are uniform across the Block. +/// Counts lie in `0..=block_threads*items`; invalid counts trap before the call. +/// A partial right recipe has no successor parameter. +/// +/// Each lane supplies `items` initialized values and `ceil(items/8)` packed, +/// low-bit-first validity bytes. Every input bit is checked before any numeric +/// read or collective call, even for a zero count. Input values and validity +/// are preserved. Output has the same capacities, is disjoint from input and +/// scratch, and receives every value and validity bit; unused high bits are zero. +/// +/// For x-fastest Block rank `r` and item `i`, position `r*items+i` uses native +/// `cuda::std::minus` with its left or right neighbor. Without an explicit +/// neighbor, the corresponding tile boundary is copied unchanged. Partial +/// recipes copy every slot beyond the valid prefix and its unsupplied boundary. +/// A zero count copies the complete input; it still executes the CUB collective. +/// +/// Every member of the exact Block participates in converged control flow. +/// Pointers remain live with their natural alignment and stated capacities. +/// Scratch uses the reported actual CUB TempStorage, including both endpoint +/// planes. The caller owns its alignment and lifetime and synchronizes Block +/// members before reuse. The provider adds no trailing completion barrier. +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub struct AdjacentDifferenceRequest { + pub(crate) recipe: AdjacentDifferenceRecipe, + pub(crate) dtype: Dtype, + block: [u32; 3], + items: u32, +} + +impl AdjacentDifferenceRequest { + /// Validate adjacent subtraction for an exact complete Block. + /// + /// `dtype` selects one of the ten numeric types; `recipe` selects one of + /// the seven boundary/count combinations. `block` gives positive XYZ axes + /// and `items` is the fixed per-lane input and output extent. + /// + /// # Errors + /// Rejects zero/overflowing axes, more than 1024 threads, zero items, or + /// a complete tile exceeding CUB's signed-int indexing range. + pub fn block( + dtype: Dtype, + recipe: AdjacentDifferenceRecipe, + block: [u32; 3], + items: u32, + ) -> Result { + let threads = block + .into_iter() + .try_fold(1_u32, u32::checked_mul) + .filter(|n| (1..=1024).contains(n)) + .ok_or_else(|| { + Error::message( + "Adjacent Difference requires 1..=1024 block threads with positive axes", + ) + })?; + if items == 0 || items > i32::MAX as u32 / threads { + return Err(Error::message( + "Adjacent Difference requires a positive tile no larger than i32::MAX", + )); + } + Ok(Self { + recipe, + dtype, + block, + items, + }) + } + + pub(crate) fn block_shape(self) -> [u32; 3] { + self.block + } + + pub(crate) fn source(self, symbol: &str) -> String { + let (dtype, _) = self.dtype.native(); + let [x, y, z] = self.block; + let items = self.items; + let tile = x * y * z * items; + let bytes = items / 8 + u32::from(items % 8 != 0); + let mut parameters = String::new(); + let mut controls = String::new(); + let mut guard = String::new(); + if self.recipe.counted() { + parameters.push_str(", int valid_items"); + controls.push_str(", valid_items"); + guard = format!("if (valid_items < 0 || valid_items > {tile}) __trap();"); + } + if self.recipe.has_neighbor() { + parameters.push_str(&format!(", {dtype} neighbor")); + controls.push_str(", neighbor"); + } + let method = + match self.recipe { + AdjacentDifferenceRecipe::LeftFull + | AdjacentDifferenceRecipe::LeftFullPredecessor => "SubtractLeft", + AdjacentDifferenceRecipe::LeftPartial + | AdjacentDifferenceRecipe::LeftPartialPredecessor => "SubtractLeftPartialTile", + AdjacentDifferenceRecipe::RightFull + | AdjacentDifferenceRecipe::RightFullSuccessor => "SubtractRight", + AdjacentDifferenceRecipe::RightPartial => "SubtractRightPartialTile", + }; + format!( + r#" +#include +#include +namespace {symbol}_detail {{ +using value_type = {dtype}; +using primitive = cub::BlockAdjacentDifference; +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{parameters}) {{ + using namespace {symbol}_detail; + {guard} +#pragma unroll + for (int i = 0; i < {bytes}; ++i) output_bits[i] = 0; +#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]; + value_type local_output[{items}]; + primitive operation{{*reinterpret_cast(scratch)}}; + operation.{method}(local_input, local_output, ::cuda::std::minus{{}}{controls}); +#pragma unroll + for (int i = 0; i < {items}; ++i) {{ + output[i] = local_output[i]; + output_bits[i / 8] |= static_cast(1u << (i % 8)); + }} +}} +"# + ) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn geometry_and_complete_tile_index_bounds() { + for (shape, items) in [ + ([1, 1, 1], i32::MAX as u32), + ([33, 1, 1], 9), + ([4, 3, 2], 4), + ([1024, 1, 1], i32::MAX as u32 / 1024), + ] { + assert!( + AdjacentDifferenceRequest::block( + Dtype::I8, + AdjacentDifferenceRecipe::LeftFull, + shape, + items + ) + .is_ok() + ); + } + for (shape, items) in [ + ([0, 1, 1], 1), + ([1, 0, 1], 1), + ([1, 1, 0], 1), + ([1025, 1, 1], 1), + ([u32::MAX, 2, 2], 1), + ([1, 1, 1], 0), + ([1, 1, 1], i32::MAX as u32 + 1), + ([1, 1, 1], u32::MAX), + ([2, 1, 1], i32::MAX as u32), + ([1024, 1, 1], i32::MAX as u32 / 1024 + 1), + ] { + assert!( + AdjacentDifferenceRequest::block( + Dtype::F64, + AdjacentDifferenceRecipe::RightPartial, + shape, + items + ) + .is_err() + ); + } + } +} diff --git a/rust/cuda-coop-codegen/src/discontinuity.rs b/rust/cuda-coop-codegen/src/discontinuity.rs new file mode 100644 index 000000000000..f88f446aa06d --- /dev/null +++ b/rust/cuda-coop-codegen/src/discontinuity.rs @@ -0,0 +1,293 @@ +// Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +//! Validated Block discontinuity flags with preserved numeric input. + +use crate::{Dtype, Error, Result}; + +/// Head/tail results and optional neighbors outside a complete Block tile. +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub enum DiscontinuityRecipe { + /// Head flags, with the first tile flag set to one. + Heads, + /// Head flags, comparing the first item with a supplied predecessor. + HeadsPredecessor, + /// Tail flags, with the last tile flag set to one. + Tails, + /// Tail flags, comparing the last item with a supplied successor. + TailsSuccessor, + /// Separate head and tail flags, with both tile boundaries set to one. + HeadsAndTails, + /// Both results with a supplied predecessor and an implicit final boundary. + HeadsAndTailsPredecessor, + /// Both results with an implicit first boundary and a supplied successor. + HeadsAndTailsSuccessor, + /// Both results with supplied predecessor and successor values. + HeadsAndTailsBoth, +} + +impl DiscontinuityRecipe { + pub(crate) fn token(self) -> &'static str { + match self { + Self::Heads => "discontinuity_heads_validity_bits_v1", + Self::HeadsPredecessor => "discontinuity_heads_predecessor_validity_bits_v1", + Self::Tails => "discontinuity_tails_validity_bits_v1", + Self::TailsSuccessor => "discontinuity_tails_successor_validity_bits_v1", + Self::HeadsAndTails => "discontinuity_heads_and_tails_validity_bits_v1", + Self::HeadsAndTailsPredecessor => { + "discontinuity_heads_and_tails_predecessor_validity_bits_v1" + } + Self::HeadsAndTailsSuccessor => { + "discontinuity_heads_and_tails_successor_validity_bits_v1" + } + Self::HeadsAndTailsBoth => "discontinuity_heads_and_tails_both_validity_bits_v1", + } + } + + pub(crate) fn parse(operation: &str, algorithm: &str) -> Result { + if algorithm != "block" { + return Err(Error::message("Discontinuity requires a Block recipe")); + } + match operation { + "discontinuity_heads_validity_bits_v1" => Ok(Self::Heads), + "discontinuity_heads_predecessor_validity_bits_v1" => Ok(Self::HeadsPredecessor), + "discontinuity_tails_validity_bits_v1" => Ok(Self::Tails), + "discontinuity_tails_successor_validity_bits_v1" => Ok(Self::TailsSuccessor), + "discontinuity_heads_and_tails_validity_bits_v1" => Ok(Self::HeadsAndTails), + "discontinuity_heads_and_tails_predecessor_validity_bits_v1" => { + Ok(Self::HeadsAndTailsPredecessor) + } + "discontinuity_heads_and_tails_successor_validity_bits_v1" => { + Ok(Self::HeadsAndTailsSuccessor) + } + "discontinuity_heads_and_tails_both_validity_bits_v1" => Ok(Self::HeadsAndTailsBoth), + _ => Err(Error::message("unsupported Discontinuity operation")), + } + } + + pub(crate) fn paired(self) -> bool { + matches!( + self, + Self::HeadsAndTails + | Self::HeadsAndTailsPredecessor + | Self::HeadsAndTailsSuccessor + | Self::HeadsAndTailsBoth + ) + } + + pub(crate) fn predecessor(self) -> bool { + matches!( + self, + Self::HeadsPredecessor | Self::HeadsAndTailsPredecessor | Self::HeadsAndTailsBoth + ) + } + + pub(crate) fn successor(self) -> bool { + matches!( + self, + Self::TailsSuccessor | Self::HeadsAndTailsSuccessor | Self::HeadsAndTailsBoth + ) + } +} + +/// Validated Block discontinuity with one or two fresh, initialized i32 outputs. +/// +/// Single-result calls use `(const T* input, i32* output, u8* scratch, +/// const u8* input_bits, u8* output_bits)`. Paired calls use +/// `(const T* input, i32* heads, u8* scratch, i32* tails, +/// const u8* input_bits, u8* heads_bits, u8* tails_bits)`. +/// A supplied predecessor and/or successor follows these pointers, in that +/// order, as values of exactly `T`, uniform across the Block. No count exists. +/// +/// Every lane supplies `items` initialized numeric values and `ceil(items/8)` +/// low-bit-first validity bytes. All input bits are checked before any numeric +/// read or collective call. Each output has `items` i32 slots and its own +/// equally sized bitmap. The provider initializes every flag, marks every slot +/// valid and clears unused high bits. Input values and validity are preserved. +/// All input, output, bitmap and scratch buffers must be disjoint, live and +/// naturally aligned, with the stated capacities. +/// +/// Flags follow blocked order with x-fastest Block rank. A head compares an +/// item with its predecessor; a tail compares it with its successor. Native +/// `cuda::std::not_equal_to` produces flags zero or one, including its usual +/// floating-point NaN and signed-zero behavior. A tile boundary without a +/// supplied neighbor is flagged one. Paired results are ordered heads, tails. +/// +/// Every member of the exact Block participates in converged control flow. +/// Scratch is the reported actual CUB TempStorage, including both endpoint +/// planes. The caller provides its alignment/lifetime and synchronizes Block +/// members before reuse. The provider adds no trailing completion barrier. +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub struct DiscontinuityRequest { + pub(crate) recipe: DiscontinuityRecipe, + pub(crate) dtype: Dtype, + block: [u32; 3], + items: u32, +} + +impl DiscontinuityRequest { + /// Validate a complete Block's single or paired discontinuity results. + /// + /// `dtype` selects one of the ten numeric input types; every output is i32. + /// `recipe` selects one of eight mode/neighbor combinations. `block` gives + /// positive XYZ axes; `items` is the per-lane extent of input and each output. + /// + /// # Errors + /// Rejects zero/overflowing axes, more than 1024 threads, zero items, or + /// a complete tile exceeding CUB's signed-int indexing range. + pub fn block( + dtype: Dtype, + recipe: DiscontinuityRecipe, + block: [u32; 3], + items: u32, + ) -> Result { + let threads = block + .into_iter() + .try_fold(1_u32, u32::checked_mul) + .filter(|n| (1..=1024).contains(n)) + .ok_or_else(|| { + Error::message("Discontinuity requires 1..=1024 block threads with positive axes") + })?; + if items == 0 || items > i32::MAX as u32 / threads { + return Err(Error::message( + "Discontinuity requires a positive tile no larger than i32::MAX", + )); + } + Ok(Self { + recipe, + dtype, + block, + items, + }) + } + + pub(crate) fn block_shape(self) -> [u32; 3] { + self.block + } + + pub(crate) fn source(self, symbol: &str) -> String { + let (dtype, _) = self.dtype.native(); + let [x, y, z] = self.block; + let items = self.items; + let bytes = items / 8 + u32::from(items % 8 != 0); + let mut controls = String::new(); + if self.recipe.predecessor() { + controls.push_str(&format!(", {dtype} predecessor")); + } + if self.recipe.successor() { + controls.push_str(&format!(", {dtype} successor")); + } + let (tail_parameter, tail_bits_parameter, clear_tails, declare_tails, copy_tails) = if self + .recipe + .paired() + { + ( + ", int* tails", + ", unsigned char* tails_bits", + "tails_bits[i] = 0;", + format!("int local_tails[{items}];"), + "tails[i] = local_tails[i]; tails_bits[i / 8] |= static_cast(1u << (i % 8));", + ) + } else { + ("", "", "", String::new(), "") + }; + let invocation = match self.recipe { + DiscontinuityRecipe::Heads => "FlagHeads(local_output, local_input, flag_op)", + DiscontinuityRecipe::HeadsPredecessor => { + "FlagHeads(local_output, local_input, flag_op, predecessor)" + } + DiscontinuityRecipe::Tails => "FlagTails(local_output, local_input, flag_op)", + DiscontinuityRecipe::TailsSuccessor => { + "FlagTails(local_output, local_input, flag_op, successor)" + } + DiscontinuityRecipe::HeadsAndTails => { + "FlagHeadsAndTails(local_output, local_tails, local_input, flag_op)" + } + DiscontinuityRecipe::HeadsAndTailsPredecessor => { + "FlagHeadsAndTails(local_output, predecessor, local_tails, local_input, flag_op)" + } + DiscontinuityRecipe::HeadsAndTailsSuccessor => { + "FlagHeadsAndTails(local_output, local_tails, successor, local_input, flag_op)" + } + DiscontinuityRecipe::HeadsAndTailsBoth => { + "FlagHeadsAndTails(local_output, predecessor, local_tails, successor, local_input, flag_op)" + } + }; + format!( + r#" +#include +#include +namespace {symbol}_detail {{ +using value_type = {dtype}; +using primitive = cub::BlockDiscontinuity; +using scratch_storage = primitive::TempStorage; +}} +extern "C" __device__ void {symbol}(const {dtype}* input, int* output, + unsigned char* scratch{tail_parameter}, const unsigned char* input_bits, + unsigned char* output_bits{tail_bits_parameter}{controls}) {{ + using namespace {symbol}_detail; +#pragma unroll + for (int i = 0; i < {bytes}; ++i) {{ output_bits[i] = 0; {clear_tails} }} +#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]; + int local_output[{items}]; + {declare_tails} + const auto flag_op = ::cuda::std::not_equal_to{{}}; + primitive operation{{*reinterpret_cast(scratch)}}; + operation.{invocation}; +#pragma unroll + for (int i = 0; i < {items}; ++i) {{ + output[i] = local_output[i]; + output_bits[i / 8] |= static_cast(1u << (i % 8)); + {copy_tails} + }} +}} +"# + ) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn geometry_and_complete_tile_index_bounds() { + for (shape, items) in [ + ([1, 1, 1], i32::MAX as u32), + ([33, 1, 1], 9), + ([4, 3, 2], 4), + ([1024, 1, 1], i32::MAX as u32 / 1024), + ] { + assert!( + DiscontinuityRequest::block(Dtype::I8, DiscontinuityRecipe::Heads, shape, items) + .is_ok() + ); + } + for (shape, items) in [ + ([0, 1, 1], 1), + ([1, 0, 1], 1), + ([1, 1, 0], 1), + ([1025, 1, 1], 1), + ([u32::MAX, 2, 2], 1), + ([1, 1, 1], 0), + ([1, 1, 1], i32::MAX as u32 + 1), + ([2, 1, 1], i32::MAX as u32), + ([1024, 1, 1], i32::MAX as u32 / 1024 + 1), + ] { + assert!( + DiscontinuityRequest::block( + Dtype::F64, + DiscontinuityRecipe::HeadsAndTailsBoth, + shape, + items, + ) + .is_err() + ); + } + } +} diff --git a/rust/cuda-coop-codegen/src/lib.rs b/rust/cuda-coop-codegen/src/lib.rs index d42f1a90d9fb..b88e65b594c0 100644 --- a/rust/cuda-coop-codegen/src/lib.rs +++ b/rust/cuda-coop-codegen/src/lib.rs @@ -3,8 +3,9 @@ //! Native, host-only CUB specialization for Rust cooperative algorithms. //! -//! This private development crate compiles Load/Store, Scan, Reduce, batched Reduce, Exchange and Shuffle with -//! NVRTC. It returns linkable LTO-IR and the exact C++ scratch layout from the +//! This private development crate compiles Load/Store, Scan, Reduce, batched +//! Reduce, Exchange, Shuffle, Adjacent Difference and Discontinuity 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. //! @@ -47,6 +48,11 @@ //! supports exact Blocks and power-of-two Warps using shared memory. //! [`ShuffleRequest`] shifts a complete Block tile by one item, preserving input //! and returning fresh output with exactly one invalid boundary slot. +//! [`AdjacentDifferenceRequest`] subtracts left or right neighbors in a Block +//! tile or prefix, preserving input and copying unchanged boundary/suffix items +//! into a fully initialized fresh output. +//! [`DiscontinuityRequest`] compares neighboring Block items, preserving input +//! and returning fresh i32 head flags, tail flags or both, with complete validity. //! [`Provider::compile_native`] batches these families in one native compilation. //! //! [`Provider::from_checkout`] snapshots the selected CCCL and CUDA headers; @@ -76,12 +82,16 @@ #![deny(unsafe_code)] +mod adjacent_difference; +mod discontinuity; mod exchange; mod headers; mod reduce; mod reduce_batched; mod scan; mod shuffle; +pub use adjacent_difference::{AdjacentDifferenceRecipe, AdjacentDifferenceRequest}; +pub use discontinuity::{DiscontinuityRecipe, DiscontinuityRequest}; pub use exchange::{ExchangeMode, ExchangeRecipe, ExchangeRequest}; pub use reduce::{ReduceAlgorithm, ReduceForm, ReduceOperator, ReduceRecipe, ReduceRequest}; pub use reduce_batched::{ReduceBatchedLayout, ReduceBatchedRecipe, ReduceBatchedRequest}; @@ -477,6 +487,10 @@ pub enum NativeRequest { Exchange(ExchangeRequest), /// One-item Block Shuffle with preserved input and a partial fresh output. Shuffle(ShuffleRequest), + /// Block adjacent subtraction with preserved input and a complete fresh output. + AdjacentDifference(AdjacentDifferenceRequest), + /// Block head/tail flags with preserved input and one or two fresh i32 outputs. + Discontinuity(DiscontinuityRequest), } impl NativeRequest { @@ -488,6 +502,8 @@ impl NativeRequest { Self::ReduceBatched(request) => request.block_shape(), Self::Exchange(request) => request.block_shape(), Self::Shuffle(request) => request.block_shape(), + Self::AdjacentDifference(request) => request.block_shape(), + Self::Discontinuity(request) => request.block_shape(), } } @@ -499,6 +515,8 @@ impl NativeRequest { Self::ReduceBatched(request) => request.source(symbol), Self::Exchange(request) => request.source(symbol), Self::Shuffle(request) => request.source(symbol), + Self::AdjacentDifference(request) => request.source(symbol), + Self::Discontinuity(request) => request.source(symbol), } } @@ -510,6 +528,8 @@ impl NativeRequest { Self::ReduceBatched(request) => request.recipe.token(), Self::Exchange(request) => request.recipe.token().into(), Self::Shuffle(request) => request.recipe.token().into(), + Self::AdjacentDifference(request) => request.recipe.token().into(), + Self::Discontinuity(request) => request.recipe.token().into(), } } @@ -520,7 +540,9 @@ impl NativeRequest { | Self::Reduce(_) | Self::ReduceBatched(_) | Self::Exchange(_) - | Self::Shuffle(_) => false, + | Self::Shuffle(_) + | Self::AdjacentDifference(_) + | Self::Discontinuity(_) => false, } } @@ -534,7 +556,9 @@ impl NativeRequest { | Self::Reduce(_) | Self::ReduceBatched(_) | Self::Exchange(_) - | Self::Shuffle(_) => false, + | Self::Shuffle(_) + | Self::AdjacentDifference(_) + | Self::Discontinuity(_) => false, } } } @@ -553,9 +577,11 @@ pub struct Layout { /// # Device ABI obligations /// /// The memory-tile and payload-layout requirements below describe Load/Store. -/// Scan, Reduce, batched Reduce, Exchange and Shuffle use lane-private input and fresh output; -/// their ABI and validity contracts are documented by [`ScanRequest`], -/// [`ReduceRequest`], [`ReduceBatchedRequest`], [`ExchangeRequest`] and [`ShuffleRequest`]. +/// Scan, Reduce, batched Reduce, Exchange, Shuffle, Adjacent Difference and +/// Discontinuity use lane-private input and fresh output. Their ABI and validity contracts are +/// documented by [`ScanRequest`], [`ReduceRequest`], [`ReduceBatchedRequest`], +/// [`ExchangeRequest`], [`ShuffleRequest`], [`AdjacentDifferenceRequest`] and +/// [`DiscontinuityRequest`]. /// /// 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. @@ -709,9 +735,9 @@ impl Provider { /// Compile an ordered mixed batch of the supported families in one NVRTC program. /// /// Per-specialization identity is independent of unrelated batch members. - /// Scan/Reduce/batched Reduce/Exchange/Shuffle scratch describes one group and includes - /// the actual CUB TempStorage, - /// even when its C++ empty type occupies one byte. The caller plans replication. + /// Scan/Reduce/batched Reduce/Exchange/Shuffle/Adjacent Difference/Discontinuity + /// 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`] @@ -777,6 +803,8 @@ impl Provider { include_bytes!("reduce_batched.rs"), include_bytes!("exchange.rs"), include_bytes!("shuffle.rs"), + include_bytes!("adjacent_difference.rs"), + include_bytes!("discontinuity.rs"), ] { hash_field(&mut inputs, bytes); } diff --git a/rust/cuda-coop-codegen/src/oxide.rs b/rust/cuda-coop-codegen/src/oxide.rs index e572f9acce69..9975c44ea084 100644 --- a/rust/cuda-coop-codegen/src/oxide.rs +++ b/rust/cuda-coop-codegen/src/oxide.rs @@ -17,6 +17,10 @@ //! Warp batched Reduce through [`ReduceBatchedRecipe`]. Out-of-place Block/Warp //! Exchange uses [`ExchangeRecipe`] with no scalar controls; Block Shuffle uses //! [`ShuffleRecipe`] for one-item shifts with partial fresh output. +//! [`AdjacentDifferenceRecipe`] selects Block neighbor subtraction with typed +//! optional count and predecessor/successor controls. +//! [`DiscontinuityRecipe`] selects Block head/tail flags with one or two i32 +//! outputs and optional neighbors of the input type. //! 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 @@ -24,7 +28,8 @@ //! checked callback marker; the provider does not select or execute Rust code. use crate::{ - Algorithm, Artifact, Dtype, Error, ExchangeRecipe, ExchangeRequest, Layout, NativeRequest, + AdjacentDifferenceRecipe, AdjacentDifferenceRequest, Algorithm, Artifact, DiscontinuityRecipe, + DiscontinuityRequest, Dtype, Error, ExchangeRecipe, ExchangeRequest, Layout, NativeRequest, Operation, Provider, ReduceAlgorithm, ReduceBatchedRecipe, ReduceBatchedRequest, ReduceForm, ReduceOperator, ReduceRecipe, ReduceRequest, Request, Result, ScanRecipe, ScanRequest, ShuffleRecipe, ShuffleRequest, @@ -194,6 +199,32 @@ struct SpecializationRequest { impl SpecializationRequest { fn native(&self) -> Result { + if self.operation.starts_with("discontinuity_") { + let recipe = DiscontinuityRecipe::parse(&self.operation, &self.algorithm)?; + if self.warp_width.is_some() { + return Err(Error::message("Discontinuity requires a Block scope")); + } + return DiscontinuityRequest::block( + self.dtype.dtype()?, + recipe, + self.block, + self.items_per_thread, + ) + .map(NativeRequest::Discontinuity); + } + if self.operation.starts_with("adjacent_difference_") { + let recipe = AdjacentDifferenceRecipe::parse(&self.operation, &self.algorithm)?; + if self.warp_width.is_some() { + return Err(Error::message("Adjacent Difference requires a Block scope")); + } + return AdjacentDifferenceRequest::block( + self.dtype.dtype()?, + recipe, + self.block, + self.items_per_thread, + ) + .map(NativeRequest::AdjacentDifference); + } if self.operation.starts_with("shuffle_") { let recipe = ShuffleRecipe::parse(&self.operation, &self.algorithm)?; if self.warp_width.is_some() { @@ -360,6 +391,46 @@ struct BatchResponse { fn arguments(request: NativeRequest) -> Vec { let request = match request { + NativeRequest::Discontinuity(request) => { + let dtype = Scalar::from(request.dtype); + let mut arguments = vec![ + Argument::ConstPointer(dtype), + Argument::MutPointer(Scalar::Signed(32)), + ]; + if request.recipe.paired() { + arguments.push(Argument::MutPointer(Scalar::Signed(32))); + } + arguments.extend([ + Argument::ConstPointer(Scalar::Unsigned(8)), + Argument::MutPointer(Scalar::Unsigned(8)), + ]); + if request.recipe.paired() { + arguments.push(Argument::MutPointer(Scalar::Unsigned(8))); + } + if request.recipe.predecessor() { + arguments.push(Argument::Value(dtype)); + } + if request.recipe.successor() { + arguments.push(Argument::Value(dtype)); + } + return arguments; + } + NativeRequest::AdjacentDifference(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.counted() { + arguments.push(Argument::Value(Scalar::Signed(32))); + } + if request.recipe.has_neighbor() { + arguments.push(Argument::Value(dtype)); + } + return arguments; + } NativeRequest::Shuffle(request) => { let dtype = Scalar::from(request.dtype); return vec![ @@ -671,6 +742,64 @@ pub fn shuffle_registration(executable: &Path, recipe: ShuffleRecipe) -> Result< .map_err(|error| Error::message(format!("cannot encode provider registration: {error}"))) } +/// Encode a Block Adjacent Difference recipe for the existing provider executable. +/// +/// Dtype, exact Block geometry and payload extent arrive with specialization. +/// Four typed payload pointers describe preserved input and fully initialized +/// fresh output with packed validity. Partial recipes append an `i32` count; +/// recipes with a tile predecessor or successor append that value in the exact +/// input dtype, after the count when present. Scratch is inserted by the +/// existing library-call ABI. See [`AdjacentDifferenceRequest`] for capacities, +/// boundary behavior and participation requirements. +/// +/// # Errors +/// Rejects a nonabsolute executable or a path that cannot be encoded as JSON. +pub fn adjacent_difference_registration( + executable: &Path, + recipe: AdjacentDifferenceRecipe, +) -> 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: "block", + }) + .map_err(|error| Error::message(format!("cannot encode provider registration: {error}"))) +} + +/// Encode a Block Discontinuity recipe for the existing provider executable. +/// +/// Input dtype, exact Block geometry and payload extent arrive with +/// specialization. Single-result recipes use four typed payload pointers; +/// paired recipes use six, ordered input, heads, tails, input validity, heads +/// validity, tails validity. Outputs are i32 regardless of input dtype. +/// Optional predecessor and successor controls follow, in that order, with +/// exactly the input type. Scratch is inserted at native parameter two by the +/// existing ABI. See [`DiscontinuityRequest`] for initialization and reuse. +/// +/// # Errors +/// Rejects a nonabsolute executable or a path that cannot be encoded as JSON. +pub fn discontinuity_registration( + executable: &Path, + recipe: DiscontinuityRecipe, +) -> 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: "block", + }) + .map_err(|error| Error::message(format!("cannot encode provider registration: {error}"))) +} + #[derive(Deserialize)] #[serde(deny_unknown_fields)] struct InputRequest { @@ -993,6 +1122,133 @@ mod tests { ScanOperator, }; + #[test] + fn discontinuity_registration_preserves_mixed_output_abi() { + use DiscontinuityRecipe::*; + for (recipe, paired, neighbors) in [ + (Heads, false, 0), + (HeadsPredecessor, false, 1), + (Tails, false, 0), + (TailsSuccessor, false, 1), + (HeadsAndTails, true, 0), + (HeadsAndTailsPredecessor, true, 1), + (HeadsAndTailsSuccessor, true, 1), + (HeadsAndTailsBoth, true, 2), + ] { + let registration: serde_json::Value = serde_json::from_str( + &discontinuity_registration(Path::new("/provider"), recipe).unwrap(), + ) + .unwrap(); + let mut request = SpecializationRequest { + operation: registration["operation"].as_str().unwrap().into(), + algorithm: registration["algorithm"].as_str().unwrap().into(), + dtype: Scalar::Float(64), + block: [4, 3, 2], + items_per_thread: 4, + warp_width: None, + }; + let native = request.native().unwrap(); + assert_eq!( + native, + NativeRequest::Discontinuity( + DiscontinuityRequest::block(Dtype::F64, recipe, [4, 3, 2], 4).unwrap() + ) + ); + let mut expected = vec![ + Argument::ConstPointer(Scalar::Float(64)), + Argument::MutPointer(Scalar::Signed(32)), + ]; + if paired { + expected.push(Argument::MutPointer(Scalar::Signed(32))); + } + expected.extend([ + Argument::ConstPointer(Scalar::Unsigned(8)), + Argument::MutPointer(Scalar::Unsigned(8)), + ]); + if paired { + expected.push(Argument::MutPointer(Scalar::Unsigned(8))); + } + expected.extend((0..neighbors).map(|_| Argument::Value(Scalar::Float(64)))); + assert_eq!(arguments(native), expected); + assert!(!native.scratch_free() && !native.mutates_payload()); + assert_eq!(registration["version"], 1); + assert_eq!(registration["provider"], "cuda_coop"); + request.warp_width = Some(32); + assert!(request.native().is_err()); + } + assert!(discontinuity_registration(Path::new("relative"), Heads).is_err()); + assert!( + DiscontinuityRecipe::parse("discontinuity_heads_successor_validity_bits_v1", "block") + .is_err() + ); + assert!( + DiscontinuityRecipe::parse("discontinuity_tails_predecessor_validity_bits_v1", "block") + .is_err() + ); + assert!(DiscontinuityRecipe::parse(Heads.token(), "warp").is_err()); + assert!(DiscontinuityRecipe::parse(Heads.token(), "sort").is_err()); + } + + #[test] + fn adjacent_difference_registration_preserves_typed_controls() { + use AdjacentDifferenceRecipe::*; + for (recipe, controls) in [ + (LeftFull, vec![]), + (LeftFullPredecessor, vec![Scalar::Float(64)]), + (LeftPartial, vec![Scalar::Signed(32)]), + ( + LeftPartialPredecessor, + vec![Scalar::Signed(32), Scalar::Float(64)], + ), + (RightFull, vec![]), + (RightFullSuccessor, vec![Scalar::Float(64)]), + (RightPartial, vec![Scalar::Signed(32)]), + ] { + let registration: serde_json::Value = serde_json::from_str( + &adjacent_difference_registration(Path::new("/provider"), recipe).unwrap(), + ) + .unwrap(); + let mut request = SpecializationRequest { + operation: registration["operation"].as_str().unwrap().into(), + algorithm: registration["algorithm"].as_str().unwrap().into(), + dtype: Scalar::Float(64), + block: [4, 3, 2], + items_per_thread: 9, + warp_width: None, + }; + let native = request.native().unwrap(); + assert_eq!( + native, + NativeRequest::AdjacentDifference( + AdjacentDifferenceRequest::block(Dtype::F64, recipe, [4, 3, 2], 9).unwrap() + ) + ); + let mut expected = vec![ + Argument::ConstPointer(Scalar::Float(64)), + Argument::MutPointer(Scalar::Float(64)), + Argument::ConstPointer(Scalar::Unsigned(8)), + Argument::MutPointer(Scalar::Unsigned(8)), + ]; + expected.extend(controls.into_iter().map(Argument::Value)); + assert_eq!(arguments(native), expected); + assert!(!native.scratch_free() && !native.mutates_payload()); + assert_eq!(registration["version"], 1); + assert_eq!(registration["provider"], "cuda_coop"); + request.warp_width = Some(32); + assert!(request.native().is_err()); + } + assert!(adjacent_difference_registration(Path::new("relative"), LeftFull).is_err()); + assert!( + AdjacentDifferenceRecipe::parse( + "adjacent_difference_right_partial_successor_validity_bits_v1", + "block" + ) + .is_err() + ); + assert!(AdjacentDifferenceRecipe::parse(LeftFull.token(), "warp").is_err()); + assert!(AdjacentDifferenceRecipe::parse(RightFull.token(), "transpose").is_err()); + } + #[test] fn shuffle_registration_preserves_scope_extent_and_pointer_abi() { for recipe in [ShuffleRecipe::Down, ShuffleRecipe::Up] { diff --git a/rust/cuda-coop-codegen/tests/native.rs b/rust/cuda-coop-codegen/tests/native.rs index ce3cbddfe0ea..7c2024768da8 100644 --- a/rust/cuda-coop-codegen/tests/native.rs +++ b/rust/cuda-coop-codegen/tests/native.rs @@ -4,13 +4,262 @@ //! Real native compiler coverage; ordinary host tests need no CUDA installation. use cuda_coop_codegen::{ - Algorithm, Dtype, ExchangeMode, ExchangeRecipe, ExchangeRequest, Layout, NativeRequest, - Operation, Provider, ReduceAlgorithm, ReduceBatchedLayout, ReduceBatchedRecipe, + AdjacentDifferenceRecipe, AdjacentDifferenceRequest, Algorithm, DiscontinuityRecipe, + DiscontinuityRequest, Dtype, ExchangeMode, ExchangeRecipe, ExchangeRequest, Layout, + NativeRequest, Operation, Provider, ReduceAlgorithm, ReduceBatchedLayout, ReduceBatchedRecipe, ReduceBatchedRequest, ReduceForm, ReduceOperator, ReduceRequest, Request, ScanAlgorithm, ScanOperation, ScanOperator, ScanRequest, ShuffleRecipe, ShuffleRequest, }; use std::{env, path::Path, time::Instant}; +#[test] +#[ignore = "requires CUDA_HOME and a Linux CUDA toolkit with NVRTC LTO support"] +fn native_mixed_discontinuity_recipes() -> Result<(), Box> { + use DiscontinuityRecipe::*; + 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(); + // Eight overloads and all input types; every result uses i32. N9 crosses + // the bitmap boundary, including independent heads/tails validity planes. + for (dtype, block, items, recipe) in [ + (Dtype::F64, [4, 3, 2], 4, HeadsAndTailsBoth), + (Dtype::I8, [1, 1, 1], 1, Heads), + (Dtype::U8, [1, 1, 1], 1, Tails), + (Dtype::I16, [33, 1, 1], 4, HeadsPredecessor), + (Dtype::U16, [4, 3, 2], 4, TailsSuccessor), + (Dtype::I32, [33, 1, 1], 9, HeadsAndTails), + (Dtype::U32, [64, 1, 1], 4, HeadsAndTailsPredecessor), + (Dtype::I64, [64, 1, 1], 1, HeadsAndTailsSuccessor), + (Dtype::U64, [4, 3, 2], 9, Heads), + (Dtype::F32, [128, 1, 1], 4, HeadsAndTailsBoth), + ] { + requests.push(NativeRequest::Discontinuity(DiscontinuityRequest::block( + dtype, recipe, block, items, + )?)); + } + requests.extend([ + NativeRequest::LoadStore(Request::new_typed( + Dtype::I32, + Operation::LoadFullValidity, + Algorithm::Direct, + 64, + 4, + )?), + NativeRequest::Scan(ScanRequest::warp( + Dtype::I32, + ScanOperation::ExclusiveSum, + [64, 1, 1], + 8, + )?), + NativeRequest::Reduce(ReduceRequest::warp( + Dtype::I32, + ReduceOperator::Sum, + ReduceForm::Payload, + [64, 1, 1], + 8, + 4, + )?), + NativeRequest::ReduceBatched(ReduceBatchedRequest::warp( + Dtype::I32, + ReduceBatchedRecipe::new(ReduceOperator::Sum, ReduceBatchedLayout::Striped), + [64, 1, 1], + 8, + 4, + )?), + NativeRequest::Exchange(ExchangeRequest::warp( + Dtype::I32, + ExchangeRecipe::Warp(ExchangeMode::StripedToBlocked), + [64, 1, 1], + 32, + 4, + )?), + NativeRequest::Shuffle(ShuffleRequest::block( + Dtype::I32, + ShuffleRecipe::Down, + [64, 1, 1], + 4, + )?), + NativeRequest::AdjacentDifference(AdjacentDifferenceRequest::block( + Dtype::I32, + AdjacentDifferenceRecipe::LeftPartialPredecessor, + [64, 1, 1], + 4, + )?), + ]); + let artifact = provider.compile_native(&requests, &target)?; + assert_eq!(artifact.specializations.len(), 17); + 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 + ); + } + // Actual CUB TempStorage contains both endpoint planes, even for Heads. + assert_eq!( + artifact.specializations[0].scratch, + Layout { + size: 384, + alignment: 8 + } + ); + assert_eq!( + artifact.specializations[1].scratch, + Layout { + size: 2, + alignment: 1 + } + ); + assert!( + provider + .compile_native(&[requests[0], requests[0]], &target) + .unwrap_err() + .to_string() + .contains("duplicate") + ); + eprintln!( + "{} mixed Discontinuity/Adjacent Difference/Load/Scan/Reduce/batched Reduce/Exchange/Shuffle 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_adjacent_difference_recipes() -> Result<(), Box> { + use AdjacentDifferenceRecipe::*; + 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(); + // Distribute all types and seven overloads over small, odd and XYZ Blocks; + // N9 crosses the packed-validity byte boundary without a Cartesian matrix. + for (dtype, block, items, recipe) in [ + (Dtype::I8, [1, 1, 1], 1, LeftFull), + (Dtype::U8, [1, 1, 1], 1, RightFull), + (Dtype::I16, [33, 1, 1], 1, LeftFullPredecessor), + (Dtype::U16, [4, 3, 2], 4, RightFullSuccessor), + (Dtype::I32, [33, 1, 1], 9, LeftPartial), + (Dtype::U32, [64, 1, 1], 4, RightPartial), + (Dtype::I64, [4, 3, 2], 4, LeftPartialPredecessor), + (Dtype::U64, [64, 1, 1], 1, LeftFull), + (Dtype::F32, [128, 1, 1], 4, RightFull), + (Dtype::F64, [4, 3, 2], 9, LeftPartialPredecessor), + ] { + requests.push(NativeRequest::AdjacentDifference( + AdjacentDifferenceRequest::block(dtype, recipe, block, items)?, + )); + } + requests.extend([ + NativeRequest::LoadStore(Request::new_typed( + Dtype::I32, + Operation::LoadFullValidity, + Algorithm::Direct, + 64, + 4, + )?), + NativeRequest::Scan(ScanRequest::warp( + Dtype::I32, + ScanOperation::ExclusiveSum, + [64, 1, 1], + 8, + )?), + NativeRequest::Reduce(ReduceRequest::warp( + Dtype::I32, + ReduceOperator::Sum, + ReduceForm::Payload, + [64, 1, 1], + 8, + 4, + )?), + NativeRequest::ReduceBatched(ReduceBatchedRequest::warp( + Dtype::I32, + ReduceBatchedRecipe::new(ReduceOperator::Sum, ReduceBatchedLayout::Striped), + [64, 1, 1], + 8, + 4, + )?), + NativeRequest::Exchange(ExchangeRequest::warp( + Dtype::I32, + ExchangeRecipe::Warp(ExchangeMode::StripedToBlocked), + [64, 1, 1], + 32, + 4, + )?), + NativeRequest::Shuffle(ShuffleRequest::block( + Dtype::I32, + ShuffleRecipe::Down, + [64, 1, 1], + 4, + )?), + ]); + let artifact = provider.compile_native(&requests, &target)?; + assert_eq!(artifact.specializations.len(), 16); + 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 + ); + } + // Both endpoint planes belong to CUB TempStorage, even for one direction. + assert_eq!( + artifact.specializations[0].scratch, + Layout { + size: 2, + alignment: 1 + } + ); + assert_eq!( + artifact.specializations[9].scratch, + Layout { + size: 384, + alignment: 8 + } + ); + assert!( + provider + .compile_native(&[requests[0], requests[0]], &target) + .unwrap_err() + .to_string() + .contains("duplicate") + ); + eprintln!( + "{} mixed Adjacent Difference/Load/Scan/Reduce/batched Reduce/Exchange/Shuffle 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_shuffle_recipes() -> Result<(), Box> { diff --git a/rust/cuda-coop/build.rs b/rust/cuda-coop/build.rs index c1144e841372..d470b30e1e34 100644 --- a/rust/cuda-coop/build.rs +++ b/rust/cuda-coop/build.rs @@ -3,9 +3,10 @@ //! Register the owning host provider for native cooperative operations. use cuda_coop_codegen::{ - Algorithm, ExchangeMode, ExchangeRecipe, Operation, ReduceAlgorithm, ReduceBatchedLayout, - ReduceBatchedRecipe, ReduceForm, ReduceOperator, ReduceRecipe, ScanAlgorithm, ScanOperation, - ScanOperator, ScanRecipe, ShuffleRecipe, oxide, + AdjacentDifferenceRecipe, Algorithm, DiscontinuityRecipe, ExchangeMode, ExchangeRecipe, + Operation, ReduceAlgorithm, ReduceBatchedLayout, ReduceBatchedRecipe, ReduceForm, + ReduceOperator, ReduceRecipe, ScanAlgorithm, ScanOperation, ScanOperator, ScanRecipe, + ShuffleRecipe, oxide, }; fn main() { @@ -24,6 +25,51 @@ fn main() { let callable = oxide::reduce_callable_registration(&executable) .expect("cuda.coop callable Reduce registration"); println!("cargo::rustc-env=COOP_REDUCE_RAKING_CALLABLE_PAYLOAD={callable}"); + for (recipe, name) in [ + (DiscontinuityRecipe::Heads, "HEADS"), + (DiscontinuityRecipe::HeadsPredecessor, "HEADS_PREDECESSOR"), + (DiscontinuityRecipe::Tails, "TAILS"), + (DiscontinuityRecipe::TailsSuccessor, "TAILS_SUCCESSOR"), + (DiscontinuityRecipe::HeadsAndTails, "HEADS_AND_TAILS"), + ( + DiscontinuityRecipe::HeadsAndTailsPredecessor, + "HEADS_AND_TAILS_PREDECESSOR", + ), + ( + DiscontinuityRecipe::HeadsAndTailsSuccessor, + "HEADS_AND_TAILS_SUCCESSOR", + ), + ( + DiscontinuityRecipe::HeadsAndTailsBoth, + "HEADS_AND_TAILS_BOTH", + ), + ] { + let value = oxide::discontinuity_registration(&executable, recipe) + .expect("cuda.coop discontinuity registration"); + println!("cargo::rustc-env=COOP_DISCONTINUITY_{name}={value}"); + } + for (recipe, name) in [ + (AdjacentDifferenceRecipe::LeftFull, "LEFT_FULL"), + ( + AdjacentDifferenceRecipe::LeftFullPredecessor, + "LEFT_FULL_PREDECESSOR", + ), + (AdjacentDifferenceRecipe::LeftPartial, "LEFT_PARTIAL"), + ( + AdjacentDifferenceRecipe::LeftPartialPredecessor, + "LEFT_PARTIAL_PREDECESSOR", + ), + (AdjacentDifferenceRecipe::RightFull, "RIGHT_FULL"), + ( + AdjacentDifferenceRecipe::RightFullSuccessor, + "RIGHT_FULL_SUCCESSOR", + ), + (AdjacentDifferenceRecipe::RightPartial, "RIGHT_PARTIAL"), + ] { + let value = oxide::adjacent_difference_registration(&executable, recipe) + .expect("cuda.coop adjacent difference registration"); + println!("cargo::rustc-env=COOP_ADJACENT_DIFFERENCE_{name}={value}"); + } for (recipe, name) in [(ShuffleRecipe::Down, "DOWN"), (ShuffleRecipe::Up, "UP")] { let value = oxide::shuffle_registration(&executable, recipe) .expect("cuda.coop shuffle registration"); diff --git a/rust/cuda-coop/examples/adjacent_difference/.gitignore b/rust/cuda-coop/examples/adjacent_difference/.gitignore new file mode 100644 index 000000000000..74bffa8af4a4 --- /dev/null +++ b/rust/cuda-coop/examples/adjacent_difference/.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/adjacent_difference/Cargo.lock b/rust/cuda-coop/examples/adjacent_difference/Cargo.lock new file mode 100644 index 000000000000..f59254e531cf --- /dev/null +++ b/rust/cuda-coop/examples/adjacent_difference/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-adjacent-difference" +version = "0.0.0" +dependencies = [ + "coop-adjacent-difference-helper", + "cuda-async", + "cuda-coop", + "cuda-core", + "cuda-device", + "cuda-host", +] + +[[package]] +name = "coop-adjacent-difference-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/adjacent_difference/Cargo.toml b/rust/cuda-coop/examples/adjacent_difference/Cargo.toml new file mode 100644 index 000000000000..dd8edf6de22e --- /dev/null +++ b/rust/cuda-coop/examples/adjacent_difference/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-adjacent-difference" +version = "0.0.0" +edition = "2024" +publish = false + +[workspace] +members = ["helper"] + +[dependencies] +cuda-coop = { path = "../.." } +coop-adjacent-difference-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/adjacent_difference/helper/Cargo.toml b/rust/cuda-coop/examples/adjacent_difference/helper/Cargo.toml new file mode 100644 index 000000000000..bb80215c53ae --- /dev/null +++ b/rust/cuda-coop/examples/adjacent_difference/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-adjacent-difference-helper" +version = "0.0.0" +edition = "2024" +publish = false + +[dependencies] +cuda-coop = { path = "../../.." } diff --git a/rust/cuda-coop/examples/adjacent_difference/helper/src/lib.rs b/rust/cuda-coop/examples/adjacent_difference/helper/src/lib.rs new file mode 100644 index 000000000000..0e52f8a9e6ab --- /dev/null +++ b/rust/cuda-coop/examples/adjacent_difference/helper/src/lib.rs @@ -0,0 +1,56 @@ +// Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +//! Compute adjacent differences with inferred payload types and a static extent. +#![no_std] +use cuda_coop as coop; + +/// Load a fully readable tile, subtract adjacent active items, and store the active prefix. +/// +/// `zero` initializes out-of-bounds input slots before the operation. The count +/// limits subtraction, not input readability; the initialized suffix is copied. +/// Every thread participates, including threads beyond the final active item. +pub fn difference_tile( + block: coop::ThreadGroup, + input: &[T], + output: &mut coop::DisjointSlice<'_, T, coop::CollectiveTiles>, + offset: usize, + zero: T, + direction: D, +) where + T: coop::Numeric, + D: coop::AdjacentDifferenceDirection, +{ + let capacity = block.count() as usize * ITEMS_PER_THREAD; + let remaining = if offset <= input.len() { + input.len() - offset + } else { + 0 + }; + let valid = remaining.min(capacity); + 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 differences = coop::adjacent_difference( + block, + &values, + coop::AdjacentDifferenceOptions::default() + .direction(direction) + .valid_items(valid), + ); + coop::store( + block, + output, + &differences, + coop::StoreOptions::default() + .offset(offset) + .valid_items(valid), + ); +} diff --git a/rust/cuda-coop/examples/adjacent_difference/src/main.rs b/rust/cuda-coop/examples/adjacent_difference/src/main.rs new file mode 100644 index 000000000000..2c1ec7da1c5a --- /dev/null +++ b/rust/cuda-coop/examples/adjacent_difference/src/main.rs @@ -0,0 +1,101 @@ +// 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 difference_tile( + input: &[i32], + mut output: coop::DisjointSlice>, + ) { + let block = coop::this_block(); + let offset = + block.rank_at::() as usize * block.count() as usize * ITEMS_PER_THREAD; + if RIGHT { + coop_adjacent_difference_helper::difference_tile::<_, _, ITEMS_PER_THREAD>( + block, + input, + &mut output, + offset, + 0_i32, + coop::Right, + ); + } else { + coop_adjacent_difference_helper::difference_tile::<_, _, ITEMS_PER_THREAD>( + block, + input, + &mut output, + offset, + 0_i32, + coop::Left, + ); + } + } +} + +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, $right:literal) => {{ + let tile = 128 * $items; + let input: Vec = (0..3 * tile - 7) + .map(|i| (i * 17 % 97) as i32 - 48) + .collect(); + let expected: Vec = input + .chunks(tile) + .flat_map(|values| { + values.iter().enumerate().map(|(i, &value)| { + if $right { + if i + 1 < values.len() { + value - values[i + 1] + } else { + value + } + } else if i > 0 { + value - values[i - 1] + } else { + value + } + }) + }) + .collect(); + let source = DeviceBuffer::from_host(&stream, &input)?; + let target = DeviceBuffer::from_host(&stream, &vec![-1_i32; input.len()])?; + let plan = module + .prepare_difference_tile::<$items, $right>(LaunchConfig1D::auto_shared(3, 128))?; + let (_, target) = module + .difference_tile_async_owned::<$items, $right, _, _>(&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={}, right={}: {} differences checked", + $items, + $right, + output.len() + ); + }}; + } + run!(1, false); + run!(4, false); + run!(1, true); + run!(4, true); + Ok(()) +} diff --git a/rust/cuda-coop/examples/discontinuity/.gitignore b/rust/cuda-coop/examples/discontinuity/.gitignore new file mode 100644 index 000000000000..74bffa8af4a4 --- /dev/null +++ b/rust/cuda-coop/examples/discontinuity/.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/discontinuity/Cargo.lock b/rust/cuda-coop/examples/discontinuity/Cargo.lock new file mode 100644 index 000000000000..3915d3d92e99 --- /dev/null +++ b/rust/cuda-coop/examples/discontinuity/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-discontinuity" +version = "0.0.0" +dependencies = [ + "coop-discontinuity-helper", + "cuda-async", + "cuda-coop", + "cuda-core", + "cuda-device", + "cuda-host", +] + +[[package]] +name = "coop-discontinuity-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/discontinuity/Cargo.toml b/rust/cuda-coop/examples/discontinuity/Cargo.toml new file mode 100644 index 000000000000..f40edd821f66 --- /dev/null +++ b/rust/cuda-coop/examples/discontinuity/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-discontinuity" +version = "0.0.0" +edition = "2024" +publish = false + +[workspace] +members = ["helper"] + +[dependencies] +cuda-coop = { path = "../.." } +coop-discontinuity-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/discontinuity/helper/Cargo.toml b/rust/cuda-coop/examples/discontinuity/helper/Cargo.toml new file mode 100644 index 000000000000..b1c1c3a5cb6d --- /dev/null +++ b/rust/cuda-coop/examples/discontinuity/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-discontinuity-helper" +version = "0.0.0" +edition = "2024" +publish = false + +[dependencies] +cuda-coop = { path = "../../.." } diff --git a/rust/cuda-coop/examples/discontinuity/helper/src/lib.rs b/rust/cuda-coop/examples/discontinuity/helper/src/lib.rs new file mode 100644 index 000000000000..65923da7c984 --- /dev/null +++ b/rust/cuda-coop/examples/discontinuity/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 + +//! Infer input types while returning distinct head and tail flags for each tile. +#![no_std] +use cuda_coop as coop; + +/// Load a complete tile and return head flags, tail flags and preserved input. +/// +/// Every input slot participates. Outside neighbors have the exact input type +/// and must be uniform within the Block. The two fresh flag payloads contain i32. +pub fn flag_tile( + block: coop::ThreadGroup, + input: &[T], + offset: usize, + predecessor: T, + successor: T, +) -> ( + coop::ThreadData, + coop::ThreadData, + coop::ThreadData, +) { + let mut values = coop::ThreadData::<_, ITEMS_PER_THREAD>::new(); + coop::load( + block, + input, + &mut values, + coop::LoadOptions::default().offset(offset), + ); + let (heads, tails) = coop::discontinuity( + block, + &values, + coop::DiscontinuityOptions::default() + .mode(coop::HeadsAndTails) + .tile_predecessor_item(predecessor) + .tile_successor_item(successor), + ); + (heads, tails, values) +} diff --git a/rust/cuda-coop/examples/discontinuity/src/main.rs b/rust/cuda-coop/examples/discontinuity/src/main.rs new file mode 100644 index 000000000000..674d50f782e5 --- /dev/null +++ b/rust/cuda-coop/examples/discontinuity/src/main.rs @@ -0,0 +1,128 @@ +// 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 = 3, block = (4, 3, 2), dynamic_shared = inferred)] + pub fn flag_tile( + input: &[f64], + mut heads: coop::DisjointSlice>, + mut tails: coop::DisjointSlice>, + mut preserved: coop::DisjointSlice>, + predecessor: f64, + successor: f64, + ) { + let block = coop::this_block(); + let offset = + block.rank_at::() as usize * block.count() as usize * ITEMS_PER_THREAD; + let (head_flags, tail_flags, values) = coop_discontinuity_helper::flag_tile::< + _, + ITEMS_PER_THREAD, + >( + block, input, offset, predecessor, successor + ); + coop::store( + block, + &mut heads, + &head_flags, + coop::StoreOptions::default().offset(offset), + ); + coop::store( + block, + &mut tails, + &tail_flags, + coop::StoreOptions::default().offset(offset), + ); + coop::store( + block, + &mut preserved, + &values, + coop::StoreOptions::default().offset(offset), + ); + } +} + +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, 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)? }; + macro_rules! run { + ($items:literal) => {{ + let tile = 24 * $items; + let input: Vec = (0..3 * tile) + .map(|i| ((i % tile / 3 + 2 * (i / tile)) % 7) as f64 + 0.25) + .collect(); + let predecessor = 0.25_f64; + let successor = 3.25_f64; + let mut expected_heads = Vec::with_capacity(input.len()); + let mut expected_tails = Vec::with_capacity(input.len()); + for values in input.chunks_exact(tile) { + for (i, &value) in values.iter().enumerate() { + let previous = if i == 0 { predecessor } else { values[i - 1] }; + let next = if i + 1 == tile { + successor + } else { + values[i + 1] + }; + expected_heads.push(i32::from(previous != value)); + expected_tails.push(i32::from(value != next)); + } + } + let source = DeviceBuffer::from_host(&stream, &input)?; + let head_target = DeviceBuffer::from_host(&stream, &vec![-1_i32; input.len()])?; + let tail_target = DeviceBuffer::from_host(&stream, &vec![-1_i32; input.len()])?; + let preserved_target = DeviceBuffer::from_host(&stream, &vec![-1_f64; input.len()])?; + let plan = module + .prepare_flag_tile::<$items>(LaunchConfig3D::auto_shared((3, 1, 1), (4, 3, 2)))?; + let (source, head_target, tail_target, preserved_target) = module + .flag_tile_async_owned::<$items, _, _, _, _>( + &plan, + source, + head_target, + tail_target, + preserved_target, + predecessor, + successor, + ) + .sync()?; + let mut heads = vec![0_i32; input.len()]; + let mut tails = vec![0_i32; input.len()]; + let mut preserved = vec![0_f64; input.len()]; + let mut source_after = vec![0_f64; input.len()]; + head_target.copy_to_host(&stream, &mut heads)?; + tail_target.copy_to_host(&stream, &mut tails)?; + preserved_target.copy_to_host(&stream, &mut preserved)?; + source.copy_to_host(&stream, &mut source_after)?; + stream.synchronize()?; + assert_eq!(heads, expected_heads); + assert_eq!(tails, expected_tails); + for i in 0..input.len() { + assert_eq!(preserved[i].to_bits(), input[i].to_bits()); + assert_eq!(source_after[i].to_bits(), input[i].to_bits()); + } + println!( + "items/thread={}: {} head flags, {} tail flags and {} preserved inputs checked", + $items, + heads.len(), + tails.len(), + preserved.len() + ); + }}; + } + run!(1); + run!(4); + Ok(()) +} diff --git a/rust/cuda-coop/src/adjacent_difference.rs b/rust/cuda-coop/src/adjacent_difference.rs new file mode 100644 index 000000000000..43bc6965ea70 --- /dev/null +++ b/rust/cuda-coop/src/adjacent_difference.rs @@ -0,0 +1,264 @@ +// 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, Numeric, SupportedAlignment, TempStorage, TempStorageLike, ThreadData, + ThreadGroup, +}; + +mod boundary; + +/// Subtract the previous item in blocked order (the default direction). +/// +/// Without an outside predecessor, copy the first active item unchanged. +#[derive(Clone, Copy, Debug, Default)] +pub struct Left; + +/// Subtract the next item in blocked order. +/// +/// Without an outside successor, copy the last active item unchanged. +#[derive(Clone, Copy, Debug, Default)] +pub struct Right; + +/// A sealed legal combination of direction, count and outside neighbors. +/// +/// [`Left`] permits an omitted count or `usize` and an omitted predecessor or +/// exactly `T`; its successor must be omitted. [`Right`] permits an omitted +/// count with an omitted successor or exactly `T`, or a `usize` count with no +/// successor. Its predecessor must be omitted. Omitted fields have type `()`. +/// A supplied count remains the partial form even when it equals tile capacity. +/// The device compiler checks control uniformity separately from these types. +pub trait AdjacentDifferenceDirection: + sealed::Direction +{ +} + +/// Static direction and typed controls for [`adjacent_difference`]. +/// +/// [`Default`] selects [`Left`], the full tile, no outside neighbor and implicit +/// scratch. `T` is inferred from the input. `D`, `Count`, `Pred`, `Succ` and `S` +/// record consuming builder selections. Builders preserve every other field; +/// only the seven combinations admitted by [`AdjacentDifferenceDirection`] +/// can be passed to the operation. There is no algorithm or operator selector. +/// +/// A right successor cannot be combined with any supplied count: +/// ```compile_fail,E0277 +/// use cuda_coop as coop; +/// fn invalid(input: &coop::ThreadData) { +/// coop::adjacent_difference(coop::this_block(), input, +/// coop::AdjacentDifferenceOptions::default().direction(coop::Right) +/// .valid_items(128).tile_successor_item(0)); +/// } +/// ``` +/// Neighbors must have the input's exact type: +/// ```compile_fail,E0308 +/// use cuda_coop as coop; +/// fn invalid(input: &coop::ThreadData) { +/// coop::adjacent_difference(coop::this_block(), input, +/// coop::AdjacentDifferenceOptions::default().tile_predecessor_item(0_i32)); +/// } +/// ``` +pub struct AdjacentDifferenceOptions { + count: Count, + predecessor: Pred, + successor: Succ, + storage: S, + selection: PhantomData<(T, D)>, +} + +impl Default for AdjacentDifferenceOptions { + fn default() -> Self { + Self { + count: (), + predecessor: (), + successor: (), + storage: (), + selection: PhantomData, + } + } +} + +impl AdjacentDifferenceOptions { + /// Select [`Left`] or [`Right`], retaining the count and both neighbor fields. + /// + /// Changing direction does not reinterpret a predecessor as a successor. + pub fn direction( + self, + _direction: D2, + ) -> AdjacentDifferenceOptions { + AdjacentDifferenceOptions { + count: self.count, + predecessor: self.predecessor, + successor: self.successor, + storage: self.storage, + selection: PhantomData, + } + } + + /// Limit subtraction to a Block-uniform prefix of `count` items. + /// + /// The count must be in `0..=block_threads * ITEMS_PER_THREAD`. The initialized + /// suffix is copied unchanged; zero copies the entire input. All input slots + /// must still be initialized. Any supplied count excludes a right successor. + pub fn valid_items( + self, + count: usize, + ) -> AdjacentDifferenceOptions { + AdjacentDifferenceOptions { + count, + predecessor: self.predecessor, + successor: self.successor, + storage: self.storage, + selection: PhantomData, + } + } + + /// Supply the exact-T, Block-uniform predecessor for [`Left`]. + /// + /// The first active result is input minus this value; at count zero no item + /// is active. This field is unsupported for [`Right`]. + pub fn tile_predecessor_item( + self, + value: T, + ) -> AdjacentDifferenceOptions { + AdjacentDifferenceOptions { + count: self.count, + predecessor: value, + successor: self.successor, + storage: self.storage, + selection: PhantomData, + } + } + + /// Supply the exact-T, Block-uniform successor for a full [`Right`] tile. + /// + /// The last result is input minus this value. A predecessor, any supplied + /// count or the [`Left`] direction makes this combination unsupported. + pub fn tile_successor_item( + self, + value: T, + ) -> AdjacentDifferenceOptions { + AdjacentDifferenceOptions { + count: self.count, + predecessor: self.predecessor, + successor: value, + storage: self.storage, + selection: PhantomData, + } + } + + /// Borrow explicit scratch with its existing sharing/completion policy. + /// + /// Manual reuse requires a full Block barrier. An auto-sync descriptor + /// inserts its existing trailing barrier. Options are consumed without a + /// Copy requirement, so a later call can use an ordinary mutable reborrow. + pub fn temp_storage( + self, + storage: &mut TempStorage, + ) -> AdjacentDifferenceOptions> { + AdjacentDifferenceOptions { + count: self.count, + predecessor: self.predecessor, + successor: self.successor, + storage, + selection: PhantomData, + } + } +} + +/// Subtract adjacent blocked items, preserving input and returning a fresh payload. +/// +/// `group` must be the complete [`Block`], including every lane of an XYZ block +/// in x-fastest rank order. All participants use the same static direction and +/// Block-uniform count/neighbors. `input` contains positive `ITEMS_PER_THREAD` +/// initialized items per thread. `T` supports all ten [`Numeric`] types and is +/// inferred from input; the result preserves its extent and minimum `ALIGN`. +/// Every output slot is initialized. Shared input and mutable reborrows work; +/// this signature requires no consumer `generic_const_exprs` feature. +/// +/// Default [`Left`] computes current minus previous, copying the first item +/// unless a predecessor is supplied. [`Right`] computes current minus next, +/// copying the last active item unless a full-tile successor is supplied. +/// With `valid_items`, items outside the active prefix are copied unchanged. +/// Count zero copies all input even with a predecessor. **Every input slot**, +/// including that copied suffix, must be initialized before any native read. +/// Subtraction follows CUB's typed minus; no saturation is added. +/// +/// Omitting storage selects compiler-planned scratch and synchronization before +/// dirty reuse. Explicit [`TempStorage`] retains its manual/auto-sync policy. +/// This operation adds no promise to complete unrelated global accesses; use +/// `cuda_device::thread::sync_threads()` when those accesses require completion. +/// +/// Unsupported groups, scalar inputs, directions or neighbor combinations fail +/// compilation. Invalid counts and unreadable input trap during device execution. +/// Count is validated before narrowing to CUB's signed index type; static tile +/// capacity must fit `i32`. Execution requires an Oxide-compiled CUDA kernel. +/// +/// ```no_run +/// use cuda_coop as coop; +/// fn differences( +/// block: coop::ThreadGroup, +/// input: &coop::ThreadData, +/// ) -> coop::ThreadData +/// where coop::Alignment: coop::SupportedAlignment { +/// coop::adjacent_difference(block, input, Default::default()) +/// } +/// ``` +#[inline(always)] +pub fn adjacent_difference< + T, + D, + Count, + Pred, + Succ, + S, + const ITEMS_PER_THREAD: usize, + const ALIGN: usize, +>( + _group: ThreadGroup, + input: &ThreadData, + options: AdjacentDifferenceOptions, +) -> ThreadData +where + T: Numeric, + D: AdjacentDifferenceDirection, + S: TempStorageLike, + Alignment: SupportedAlignment, +{ + boundary::adjacent_difference(input, options) +} + +mod sealed { + pub trait Direction { + type Uniform: Copy; + const KIND: u8; + fn controls( + count: Count, + predecessor: Pred, + successor: Succ, + ) -> Self::Uniform; + } +} + +// Only these seven forms can reach a registration with the matching scalar ABI. +macro_rules! direction { + ($direction:ty, $count:ty, $pred:ty, $succ:ty, $kind:literal, + |$items:ident; $c:ident, $p:ident, $s:ident| -> $uniform:ty $body:block) => { + impl sealed::Direction for $direction { + type Uniform = $uniform; + const KIND: u8 = $kind; + #[inline(always)] + fn controls($c: $count, $p: $pred, $s: $succ) -> Self::Uniform $body + } + impl AdjacentDifferenceDirection for $direction {} + }; +} +direction!(Left, (), (), (), 0, |ITEMS; _c, _p, _s| -> () {}); +direction!(Left, (), T, (), 1, |ITEMS; _c, p, _s| -> (T,) { (p,) }); +direction!(Left, usize, (), (), 2, |ITEMS; c, _p, _s| -> (i32,) { (boundary::count::(c),) }); +direction!(Left, usize, T, (), 3, |ITEMS; c, p, _s| -> (i32, T) { (boundary::count::(c), p) }); +direction!(Right, (), (), (), 4, |ITEMS; _c, _p, _s| -> () {}); +direction!(Right, (), (), T, 5, |ITEMS; _c, _p, s| -> (T,) { (s,) }); +direction!(Right, usize, (), (), 6, |ITEMS; c, _p, _s| -> (i32,) { (boundary::count::(c),) }); diff --git a/rust/cuda-coop/src/adjacent_difference/boundary.rs b/rust/cuda-coop/src/adjacent_difference/boundary.rs new file mode 100644 index 000000000000..835346260efb --- /dev/null +++ b/rust/cuda-coop/src/adjacent_difference/boundary.rs @@ -0,0 +1,97 @@ +// Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +use super::{AdjacentDifferenceDirection, AdjacentDifferenceOptions}; +use crate::{Alignment, Numeric, SupportedAlignment, TempStorageLike, ThreadData}; +use cuda_device::library; + +macro_rules! registrations { + ($(($name:ident, $variable:literal)),+ $(,)?) => {$ ( + fn $name() -> &'static str { library::registration!($variable) } + )+}; +} +registrations!( + (left_full, "COOP_ADJACENT_DIFFERENCE_LEFT_FULL"), + ( + left_full_predecessor, + "COOP_ADJACENT_DIFFERENCE_LEFT_FULL_PREDECESSOR" + ), + (left_partial, "COOP_ADJACENT_DIFFERENCE_LEFT_PARTIAL"), + ( + left_partial_predecessor, + "COOP_ADJACENT_DIFFERENCE_LEFT_PARTIAL_PREDECESSOR" + ), + (right_full, "COOP_ADJACENT_DIFFERENCE_RIGHT_FULL"), + ( + right_full_successor, + "COOP_ADJACENT_DIFFERENCE_RIGHT_FULL_SUCCESSOR" + ), + (right_partial, "COOP_ADJACENT_DIFFERENCE_RIGHT_PARTIAL"), +); + +#[inline(always)] +pub(super) fn count(count: usize) -> i32 { + let threads = crate::this_block().count() as usize; + // Validate capacity without multiplication overflow, then narrow the count. + assert!(ITEMS <= i32::MAX as usize / threads); + assert!(count <= threads * ITEMS); + count as i32 +} + +#[inline(always)] +pub(super) fn adjacent_difference< + T, + D, + Count, + Pred, + Succ, + S, + const ITEMS: usize, + const ALIGN: usize, +>( + input: &ThreadData, + options: AdjacentDifferenceOptions, +) -> ThreadData +where + T: Numeric, + D: AdjacentDifferenceDirection, + S: TempStorageLike, + Alignment: SupportedAlignment, +{ + let mut output = ThreadData::new(); + let (input_pointer, input_bits) = input.block_io_read_parts(); + let (output_pointer, output_bits) = output.block_io_parts(); + let payload = (input_pointer, output_pointer, input_bits, output_bits); + let uniform = D::controls::(options.count, options.predecessor, options.successor); + let storage = options.storage.into_raw(); + macro_rules! call { + ($registration:ident) => { + library::checked_block_payload_call_v1::( + $registration, + storage, + uniform, + payload, + ) + }; + } + // SAFETY: input and fresh output are distinct live aligned lane-local + // numeric buffers. The exact recipe validates every input bit, including + // copied suffix items at count zero, before numeric reads. It preserves + // input, initializes all output slots and then marks their bits. Pointers + // cannot escape the synchronous call. Sealed selections give the exact + // scalar ABI; the marker checks uniform controls, full participation and + // genuine storage origin/layout/reuse under the existing policy. + unsafe { + match D::KIND { + 0 => call!(left_full), + 1 => call!(left_full_predecessor), + 2 => call!(left_partial), + 3 => call!(left_partial_predecessor), + 4 => call!(right_full), + 5 => call!(right_full_successor), + 6 => call!(right_partial), + _ => panic!(), + } + } + output +} diff --git a/rust/cuda-coop/src/discontinuity.rs b/rust/cuda-coop/src/discontinuity.rs new file mode 100644 index 000000000000..3ba7a106bed9 --- /dev/null +++ b/rust/cuda-coop/src/discontinuity.rs @@ -0,0 +1,276 @@ +// 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, Numeric, SupportedAlignment, TempStorage, TempStorageLike, ThreadData, + ThreadGroup, +}; + +mod boundary; + +/// Flag items unequal to their predecessor (the default mode). +/// +/// Without an outside predecessor, the first tile flag is one. +#[derive(Clone, Copy, Debug, Default)] +pub struct Heads; + +/// Flag items unequal to their successor. +/// +/// Without an outside successor, the last tile flag is one. +#[derive(Clone, Copy, Debug, Default)] +pub struct Tails; + +/// Return separate head and tail flags, in that order. +/// +/// Each omitted outside neighbor makes its corresponding boundary flag one. +#[derive(Clone, Copy, Debug, Default)] +pub struct HeadsAndTails; + +/// A sealed Discontinuity mode and legal combination of typed outside neighbors. +/// +/// `T` is the input type; `Pred` and `Succ` are `()` when omitted or exactly `T`. +/// [`Heads`] allows only a predecessor, [`Tails`] only a successor, and +/// [`HeadsAndTails`] allows either, both or neither. Every output flag is `i32`; +/// the static item count and minimum alignment are inherited from the input. +/// Generic mode-forwarding helpers can name this trait's [`Output`](Self::Output). +pub trait DiscontinuityMode< + T: Numeric, + Pred, + Succ, + const ITEMS_PER_THREAD: usize, + const ALIGN: usize, +>: sealed::Mode where + Alignment: SupportedAlignment, +{ + /// One flag payload for Heads/Tails, or `(heads, tails)` for HeadsAndTails. + type Output; + + #[doc(hidden)] + fn apply( + input: &ThreadData, + options: DiscontinuityOptions, + ) -> Self::Output + where + Self: Sized; +} + +/// Static mode, outside neighbors and scratch policy for [`discontinuity`]. +/// +/// [`Default`] selects [`Heads`], no outside neighbor and implicit scratch. +/// `T` is inferred from the input. `M`, `Pred`, `Succ` and `S` record consuming +/// builder selections. Builders retain all other fields; the operation accepts +/// only the eight combinations admitted by [`DiscontinuityMode`]. There is no +/// count, layout, algorithm or operator selector. +/// +/// A successor cannot be used with Heads, even when it was set before the mode: +/// ```compile_fail,E0277 +/// use cuda_coop as coop; +/// fn invalid(input: &coop::ThreadData) { +/// coop::discontinuity(coop::this_block(), input, +/// coop::DiscontinuityOptions::default().tile_successor_item(7).mode(coop::Heads)); +/// } +/// ``` +/// Outside neighbors retain the input's exact type: +/// ```compile_fail,E0308 +/// use cuda_coop as coop; +/// fn invalid(input: &coop::ThreadData) { +/// coop::discontinuity(coop::this_block(), input, +/// coop::DiscontinuityOptions::default().tile_predecessor_item(1_f32)); +/// } +/// ``` +pub struct DiscontinuityOptions { + predecessor: Pred, + successor: Succ, + storage: S, + selection: PhantomData<(T, M)>, +} + +impl Default for DiscontinuityOptions { + fn default() -> Self { + Self { + predecessor: (), + successor: (), + storage: (), + selection: PhantomData, + } + } +} + +impl DiscontinuityOptions { + /// Select Heads, Tails or HeadsAndTails, retaining both neighbor fields. + /// + /// Changing mode does not reinterpret a predecessor as a successor. + pub fn mode(self, _mode: M2) -> DiscontinuityOptions { + DiscontinuityOptions { + predecessor: self.predecessor, + successor: self.successor, + storage: self.storage, + selection: PhantomData, + } + } + + /// Supply the exact-T, Block-uniform predecessor for Heads or HeadsAndTails. + /// + /// The first head flag then compares the predecessor and first input item. + pub fn tile_predecessor_item(self, value: T) -> DiscontinuityOptions { + DiscontinuityOptions { + predecessor: value, + successor: self.successor, + storage: self.storage, + selection: PhantomData, + } + } + + /// Supply the exact-T, Block-uniform successor for Tails or HeadsAndTails. + /// + /// The last tail flag then compares the last input item and successor. + pub fn tile_successor_item(self, value: T) -> DiscontinuityOptions { + DiscontinuityOptions { + predecessor: self.predecessor, + successor: value, + storage: self.storage, + selection: PhantomData, + } + } + + /// Borrow explicit scratch with its existing sharing and completion policy. + /// + /// Manual reuse requires a full Block barrier. An auto-sync descriptor + /// supplies its existing trailing barrier. Options are consumed, so later + /// calls can use ordinary mutable reborrows without requiring Copy storage. + pub fn temp_storage( + self, + storage: &mut TempStorage, + ) -> DiscontinuityOptions> { + DiscontinuityOptions { + predecessor: self.predecessor, + successor: self.successor, + storage, + selection: PhantomData, + } + } +} + +/// Flag unequal adjacent blocked items, preserving input and returning fresh flags. +/// +/// `group` must be the complete [`Block`], including all lanes of an XYZ block +/// in x-fastest rank order. Every participant uses the same static mode and +/// Block-uniform outside neighbors. `input` contains positive `ITEMS_PER_THREAD` +/// initialized items per thread. All ten [`Numeric`] types are supported and T +/// is inferred from input. Shared input and mutable reborrows are accepted. +/// +/// Heads compares each item with its predecessor; Tails compares it with its +/// successor. Without an outside neighbor, the first head or last tail is one. +/// Otherwise the flag is one for inequality and zero for equality. Floating +/// comparisons use native inequality, including NaNs and signed zero. +/// +/// Heads and Tails return `ThreadData`. +/// HeadsAndTails returns `(heads, tails)` with two independent payloads of that +/// type. Every result slot is initialized. N and minimum ALIGN are preserved; +/// output storage also satisfies i32's natural alignment. Input values and +/// validity are unchanged. Same-extent helpers require no consumer GCE feature. +/// There is no partial operation: every initialized Load fill participates in +/// comparisons, and every input slot must be readable before this call. +/// +/// Implicit scratch uses managed synchronization before dirty reuse. Explicit +/// [`TempStorage`] retains its manual/auto-sync policy. This operation adds no +/// promise to complete unrelated global accesses; use +/// `cuda_device::thread::sync_threads()` when they require Block completion. +/// +/// Unsupported groups, scalar inputs, modes or neighbor combinations fail +/// compilation. Unreadable input traps during device execution. The positive +/// tile capacity must fit CUB's signed 32-bit indexing range. Execution requires +/// an Oxide-compiled CUDA kernel. +/// +/// ```no_run +/// use cuda_coop as coop; +/// fn heads( +/// block: coop::ThreadGroup, +/// input: &coop::ThreadData, +/// ) -> coop::ThreadData +/// where coop::Alignment: coop::SupportedAlignment { +/// coop::discontinuity(block, input, Default::default()) +/// } +/// ``` +#[inline(always)] +pub fn discontinuity( + _group: ThreadGroup, + input: &ThreadData, + options: DiscontinuityOptions, +) -> M::Output +where + T: Numeric, + M: DiscontinuityMode, + S: TempStorageLike, + Alignment: SupportedAlignment, +{ + M::apply(input, options) +} + +mod sealed { + pub trait Mode { + type Uniform: Copy; + const KIND: u8; + fn controls(predecessor: Pred, successor: Succ) -> Self::Uniform; + } +} + +macro_rules! mode { + ($mode:ty, $pred:ty, $succ:ty, $kind:literal, + |$p:ident, $s:ident| -> $uniform:ty $body:block) => { + impl sealed::Mode for $mode { + type Uniform = $uniform; + const KIND: u8 = $kind; + #[inline(always)] + fn controls($p: $pred, $s: $succ) -> Self::Uniform $body + } + }; +} +mode!(Heads, (), (), 0, |_p, _s| -> () {}); +mode!(Heads, T, (), 1, |p, _s| -> (T,) { (p,) }); +mode!(Tails, (), (), 2, |_p, _s| -> () {}); +mode!(Tails, (), T, 3, |_p, s| -> (T,) { (s,) }); +mode!(HeadsAndTails, (), (), 4, |_p, _s| -> () {}); +mode!(HeadsAndTails, T, (), 5, |p, _s| -> (T,) { (p,) }); +mode!(HeadsAndTails, (), T, 6, |_p, s| -> (T,) { (s,) }); +mode!(HeadsAndTails, T, T, 7, |p, s| -> (T, T) { (p, s) }); + +macro_rules! single_output { + ($mode:ty) => { + impl + DiscontinuityMode for $mode + where + Alignment: SupportedAlignment, + Self: sealed::Mode, + { + type Output = ThreadData; + #[inline(always)] + fn apply( + input: &ThreadData, + options: DiscontinuityOptions, + ) -> Self::Output { + boundary::single(input, options) + } + } + }; +} +single_output!(Heads); +single_output!(Tails); + +impl + DiscontinuityMode for HeadsAndTails +where + Alignment: SupportedAlignment, + Self: sealed::Mode, +{ + type Output = (ThreadData, ThreadData); + #[inline(always)] + fn apply( + input: &ThreadData, + options: DiscontinuityOptions, + ) -> Self::Output { + boundary::pair(input, options) + } +} diff --git a/rust/cuda-coop/src/discontinuity/boundary.rs b/rust/cuda-coop/src/discontinuity/boundary.rs new file mode 100644 index 000000000000..e14f6d5b9e83 --- /dev/null +++ b/rust/cuda-coop/src/discontinuity/boundary.rs @@ -0,0 +1,132 @@ +// Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +use super::{DiscontinuityOptions, sealed}; +use crate::{Alignment, Numeric, SupportedAlignment, TempStorageLike, ThreadData}; +use cuda_device::library; + +macro_rules! registrations { + ($(($name:ident, $variable:literal)),+ $(,)?) => {$ ( + fn $name() -> &'static str { library::registration!($variable) } + )+}; +} +registrations!( + (heads, "COOP_DISCONTINUITY_HEADS"), + (heads_predecessor, "COOP_DISCONTINUITY_HEADS_PREDECESSOR"), + (tails, "COOP_DISCONTINUITY_TAILS"), + (tails_successor, "COOP_DISCONTINUITY_TAILS_SUCCESSOR"), + (heads_and_tails, "COOP_DISCONTINUITY_HEADS_AND_TAILS"), + ( + heads_and_tails_predecessor, + "COOP_DISCONTINUITY_HEADS_AND_TAILS_PREDECESSOR" + ), + ( + heads_and_tails_successor, + "COOP_DISCONTINUITY_HEADS_AND_TAILS_SUCCESSOR" + ), + ( + heads_and_tails_both, + "COOP_DISCONTINUITY_HEADS_AND_TAILS_BOTH" + ), +); + +#[inline(always)] +pub(super) fn single( + input: &ThreadData, + options: DiscontinuityOptions, +) -> ThreadData +where + T: Numeric, + M: sealed::Mode, + S: TempStorageLike, + Alignment: SupportedAlignment, +{ + let mut output = ThreadData::new(); + let (input_pointer, input_bits) = input.block_io_read_parts(); + let (output_pointer, output_bits) = output.block_io_parts(); + let payload = (input_pointer, output_pointer, input_bits, output_bits); + let uniform = M::controls(options.predecessor, options.successor); + let storage = options.storage.into_raw(); + macro_rules! call { + ($registration:ident) => { + library::checked_block_payload_call_v1::( + $registration, + storage, + uniform, + payload, + ) + }; + } + // SAFETY: input and fresh outputs are distinct live aligned lane-local + // buffers. The exact recipe checks every input bit before numeric reads, + // preserves input and initializes all i32 flags before marking their bits. + // Each output has its own validity buffer. Pointers cannot escape the + // synchronous call. Sealed selections give the exact typed control ABI; + // the marker checks uniformity, complete Block participation and genuine + // storage origin, layout and reuse under the existing policy. + unsafe { + match M::KIND { + 0 => call!(heads), + 1 => call!(heads_predecessor), + 2 => call!(tails), + 3 => call!(tails_successor), + _ => panic!(), + } + } + output +} + +#[inline(always)] +pub(super) fn pair( + input: &ThreadData, + options: DiscontinuityOptions, +) -> (ThreadData, ThreadData) +where + T: Numeric, + M: sealed::Mode, + S: TempStorageLike, + Alignment: SupportedAlignment, +{ + let mut output = ThreadData::new(); + let (input_pointer, input_bits) = input.block_io_read_parts(); + let (output_pointer, output_bits) = output.block_io_parts(); + let mut tails = ThreadData::new(); + let (tail_pointer, tail_bits) = tails.block_io_parts(); + let payload = ( + input_pointer, + output_pointer, + tail_pointer, + input_bits, + output_bits, + tail_bits, + ); + let uniform = M::controls(options.predecessor, options.successor); + let storage = options.storage.into_raw(); + macro_rules! call { + ($registration:ident) => { + library::checked_block_payload_call_v1::( + $registration, + storage, + uniform, + payload, + ) + }; + } + // SAFETY: input and fresh outputs are distinct live aligned lane-local + // buffers. The exact recipe checks every input bit before numeric reads, + // preserves input and initializes all i32 flags before marking their bits. + // Each output has its own validity buffer. Pointers cannot escape the + // synchronous call. Sealed selections give the exact typed control ABI; + // the marker checks uniformity, complete Block participation and genuine + // storage origin, layout and reuse under the existing policy. + unsafe { + match M::KIND { + 4 => call!(heads_and_tails), + 5 => call!(heads_and_tails_predecessor), + 6 => call!(heads_and_tails_successor), + 7 => call!(heads_and_tails_both), + _ => panic!(), + } + } + (output, tails) +} diff --git a/rust/cuda-coop/src/lib.rs b/rust/cuda-coop/src/lib.rs index da3458e2e44d..e55f911b6b76 100644 --- a/rust/cuda-coop/src/lib.rs +++ b/rust/cuda-coop/src/lib.rs @@ -59,6 +59,12 @@ //! returns a fresh payload with the same type, extent and alignment. //! [`shuffle`] shifts a complete block's values by one item, preserving input //! and returning the same payload shape with one invalid boundary slot. +//! [`adjacent_difference`] subtracts adjacent blocked items on a complete Block. +//! It preserves full-readable input and returns a fully initialized fresh payload, +//! copying any inactive count suffix unchanged. [`AdjacentDifferenceOptions`] +//! selects direction, typed outside neighbors, count and existing scratch policy. +//! [`discontinuity`] returns fresh i32 head or tail flags, or both payloads, +//! preserving the full-readable input's extent and minimum alignment. //! Host payload tests need no GPU; //! executing hardware queries requires an Oxide-compiled CUDA kernel. @@ -66,7 +72,9 @@ #![feature(generic_const_exprs)] #![allow(incomplete_features)] +mod adjacent_difference; mod alignment; +mod discontinuity; mod exchange; mod group; mod load_store; @@ -77,8 +85,14 @@ mod shuffle; mod temp_storage; mod thread_data; +pub use adjacent_difference::{ + AdjacentDifferenceDirection, AdjacentDifferenceOptions, Left, Right, adjacent_difference, +}; pub use alignment::{Alignment, SupportedAlignment}; pub use cuda_device::{CollectiveTiles, DisjointSlice}; +pub use discontinuity::{ + DiscontinuityMode, DiscontinuityOptions, Heads, HeadsAndTails, Tails, discontinuity, +}; pub use exchange::{ BlockedToStriped, ExchangeGroup, ExchangeMode, ExchangeOptions, StripedToBlocked, exchange, }; diff --git a/rust/cuda-coop/src/shuffle.rs b/rust/cuda-coop/src/shuffle.rs index b55ae77831b7..a63807d93e04 100644 --- a/rust/cuda-coop/src/shuffle.rs +++ b/rust/cuda-coop/src/shuffle.rs @@ -104,7 +104,7 @@ impl Distance { /// contribute to a valid output; otherwise execution traps before native /// numeric reads. Oxide plans implicit shared scratch from the actual CUB /// layout and synchronizes before reuse. Shuffle does not promise global-memory -/// completion: use the block's `sync()` for dependent global accesses. +/// completion: use `cuda_device::thread::sync_threads()` for dependent global accesses. /// Unsupported groups, modes or distances fail compilation. Execution requires /// an Oxide-compiled CUDA kernel; host code can check the signatures. /// diff --git a/rust/cuda-coop/tests/adjacent_difference/.gitignore b/rust/cuda-coop/tests/adjacent_difference/.gitignore new file mode 100644 index 000000000000..74bffa8af4a4 --- /dev/null +++ b/rust/cuda-coop/tests/adjacent_difference/.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/adjacent_difference/Cargo.lock b/rust/cuda-coop/tests/adjacent_difference/Cargo.lock new file mode 100644 index 000000000000..ddb98d3f403a --- /dev/null +++ b/rust/cuda-coop/tests/adjacent_difference/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-adjacent-difference-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/adjacent_difference/Cargo.toml b/rust/cuda-coop/tests/adjacent_difference/Cargo.toml new file mode 100644 index 000000000000..7b3215eb1185 --- /dev/null +++ b/rust/cuda-coop/tests/adjacent_difference/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-adjacent-difference-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/adjacent_difference/src/main.rs b/rust/cuda-coop/tests/adjacent_difference/src/main.rs new file mode 100644 index 000000000000..8c6201c12cf4 --- /dev/null +++ b/rust/cuda-coop/tests/adjacent_difference/src/main.rs @@ -0,0 +1,600 @@ +// 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; + +const CHECKS: u32 = 7; + +#[cuda_module] +mod kernels { + use super::*; + + fn same( + left: &coop::ThreadData, + right: &coop::ThreadData, + ) -> bool + where + coop::Alignment: coop::SupportedAlignment, + { + let mut matches = left.len() == N; + let mut i = 0; + while i < N { + matches &= left.get(i) == right.get(i); + i += 1; + } + matches + } + + fn observe( + input: &coop::ThreadData, + before: &coop::ThreadData, + result: &coop::ThreadData, + expected: &[T], + tid: usize, + ) -> u32 + where + coop::Alignment: coop::SupportedAlignment, + { + let mut matches = result.len() == N; + let mut i = 0; + while i < N { + // The inactive suffix is copied and initialized, never repaired. + matches &= result.get(i) == Some(expected[tid * N + i]); + i += 1; + } + matches as u32 | ((same(input, before) as u32) << 1) + } + + fn invoke( + block: coop::ThreadGroup, + input: &coop::ThreadData, + mode: u32, + count: usize, + outside: T, + ) -> coop::ThreadData + where + coop::Alignment: coop::SupportedAlignment, + { + match FORMS { + 0 => match mode { + 0 => coop::adjacent_difference(block, input, Default::default()), + 1 => coop::adjacent_difference( + block, + input, + coop::AdjacentDifferenceOptions::default().tile_predecessor_item(outside), + ), + 2 => coop::adjacent_difference( + block, + input, + coop::AdjacentDifferenceOptions::default().valid_items(count), + ), + 3 => coop::adjacent_difference( + block, + input, + coop::AdjacentDifferenceOptions::default() + .valid_items(count) + .tile_predecessor_item(outside), + ), + 4 => coop::adjacent_difference( + block, + input, + coop::AdjacentDifferenceOptions::default().direction(coop::Right), + ), + 5 => coop::adjacent_difference( + block, + input, + coop::AdjacentDifferenceOptions::default() + .direction(coop::Right) + .tile_successor_item(outside), + ), + 6 => coop::adjacent_difference( + block, + input, + coop::AdjacentDifferenceOptions::default() + .direction(coop::Right) + .valid_items(count), + ), + _ => panic!(), + }, + 1 => match mode { + 0 => coop::adjacent_difference( + block, + input, + coop::AdjacentDifferenceOptions::default().direction(coop::Right), + ), + 1 => coop::adjacent_difference( + block, + input, + coop::AdjacentDifferenceOptions::default().tile_predecessor_item(outside), + ), + _ => panic!(), + }, + 2 => match mode { + 0 => coop::adjacent_difference( + block, + input, + coop::AdjacentDifferenceOptions::default() + .valid_items(count) + .tile_predecessor_item(outside), + ), + 1 => coop::adjacent_difference( + block, + input, + coop::AdjacentDifferenceOptions::default() + .direction(coop::Right) + .tile_successor_item(outside), + ), + _ => panic!(), + }, + _ => panic!(), + } + } + + fn write_results( + block: coop::ThreadGroup, + output: &mut coop::DisjointSlice>, + reports: &mut coop::DisjointSlice>, + result: &coop::ThreadData, + word: u32, + block_index: usize, + ) where + coop::Alignment: coop::SupportedAlignment, + { + coop::store( + block, + output, + result, + coop::StoreOptions::default().offset(block_index * B * N), + ); + let mut report = coop::ThreadData::<_, 1>::new(); + report.write(0, word); + coop::store( + block, + reports, + &report, + coop::StoreOptions::default().offset(block_index * B), + ); + } + + #[allow(clippy::too_many_arguments)] + fn exercise( + input: &[T], + expected: &[T], + output: &mut coop::DisjointSlice>, + reports: &mut coop::DisjointSlice>, + rank: usize, + mode: u32, + count: usize, + outside: T, + invalid: u32, + ) where + T: coop::Numeric + PartialEq + Default, + coop::Alignment: coop::SupportedAlignment, + { + let block = coop::this_block(); + let block_index = block.rank_at::() as usize; + let tid = block_index * B + rank; + let mut values = coop::ThreadData::<_, N, A>::new(); + if invalid == 0 { + coop::load( + block, + input, + &mut values, + coop::LoadOptions::default().offset(block_index * B * N), + ); + } else { + let mut i = 0; + while i < N { + if !(invalid >= 3 && rank == B - 1 && i == N - 1) { + values.write(i, input[tid * N + i]); + } + i += 1; + } + let _result = coop::adjacent_difference( + block, + &values, + coop::AdjacentDifferenceOptions::default().valid_items(count), + ); + // Unexpected return must not inspect the result or trigger a + // downstream result guard that could counterfeit the expected trap. + let mut report = coop::ThreadData::<_, 1>::new(); + report.write(0, 99_u32); + coop::store( + block, + reports, + &report, + coop::StoreOptions::default().offset(block_index * B), + ); + return; + } + let before = values; + let mut first = invoke::<_, N, A, FORMS>(block, &values, mode, count, outside); + let mut word = observe(&values, &before, &first, expected, tid); + first.write(0, T::default()); + let mutated = first; + let second = invoke::<_, N, A, FORMS>(block, &values, mode, count, outside); + word &= observe(&values, &before, &second, expected, tid); + word |= ((same(&first, &mutated) && first.get(0) == Some(T::default())) as u32) << 2; + write_results::<_, N, A, B>(block, output, reports, &second, word, block_index); + } + + #[allow(clippy::too_many_arguments)] + #[kernel] + #[launch_contract(domain = 1, block = (33, 1, 1), dynamic_shared = inferred)] + pub fn forms_i32( + input: &[i32], + expected: &[i32], + mut output: coop::DisjointSlice>, + mut reports: coop::DisjointSlice>, + mode: u32, + count: usize, + outside: i32, + invalid: u32, + ) { + exercise::<_, 4, 16, 33, 0>( + input, + expected, + &mut output, + &mut reports, + thread::threadIdx_x() as usize, + mode, + count, + outside, + invalid, + ); + } + + #[kernel] + #[launch_contract(domain = 3, block = (5, 3, 2), dynamic_shared = inferred)] + pub fn xyz_n1( + input: &[i8], + expected: &[i8], + mut output: coop::DisjointSlice>, + mut reports: coop::DisjointSlice>, + mode: u32, + count: usize, + outside: i8, + ) { + let rank = thread::threadIdx_x() + 5 * (thread::threadIdx_y() + 3 * thread::threadIdx_z()); + exercise::<_, 1, 1, 30, 1>( + input, + expected, + &mut output, + &mut reports, + rank as usize, + mode, + count, + outside, + 0, + ); + } + + #[kernel] + #[launch_contract(domain = 1, block = (32, 1, 1), dynamic_shared = inferred)] + pub fn tail_n9( + input: &[f64], + expected: &[f64], + mut output: coop::DisjointSlice>, + mut reports: coop::DisjointSlice>, + mode: u32, + count: usize, + outside: f64, + ) { + exercise::<_, 9, 1, 32, 2>( + input, + expected, + &mut output, + &mut reports, + thread::threadIdx_x() as usize, + mode, + count, + outside, + 0, + ); + } + + #[kernel] + #[launch_contract(domain = 1, block = (1, 1, 1), dynamic_shared = inferred)] + pub fn single_n1( + input: &[u16], + expected: &[u16], + mut output: coop::DisjointSlice>, + mut reports: coop::DisjointSlice>, + mode: u32, + count: usize, + outside: u16, + ) { + exercise::<_, 1, 1, 1, 2>( + input, + expected, + &mut output, + &mut reports, + thread::threadIdx_x() as usize, + mode, + count, + outside, + 0, + ); + } + + #[allow(clippy::too_many_arguments)] + fn with_storage( + input: &[i32], + expected: &[i32], + output: &mut coop::DisjointSlice>, + reports: &mut coop::DisjointSlice>, + count: usize, + outside: i32, + scratch: &mut coop::TempStorage<0, 0, AUTO, false>, + ) { + let block = coop::this_block(); + let block_index = block.rank_at::() as usize; + let tid = block_index * 33 + thread::threadIdx_x() as usize; + let mut values = coop::ThreadData::<_, 4, 16>::new(); + coop::load( + block, + input, + &mut values, + coop::LoadOptions::default().offset(block_index * 132), + ); + let before = values; + let mut first = coop::adjacent_difference( + block, + &values, + coop::AdjacentDifferenceOptions::default() + .valid_items(count) + .tile_predecessor_item(outside) + .temp_storage(scratch), + ); + let mut word = observe(&values, &before, &first, expected, tid); + first.write(0, 0_i32); + let mutated = first; + if !AUTO { + thread::sync_threads(); + } + let second = coop::adjacent_difference( + block, + &values, + coop::AdjacentDifferenceOptions::default() + .valid_items(count) + .tile_predecessor_item(outside) + .temp_storage(scratch), + ); + word &= observe(&values, &before, &second, expected, tid); + word |= ((same(&first, &mutated) && first.get(0) == Some(0_i32)) as u32) << 2; + if !AUTO { + thread::sync_threads(); + } + write_results::<_, 4, 16, 33>(block, output, reports, &second, word, block_index); + } + + #[kernel] + #[launch_contract(domain = 1, block = (33, 1, 1), dynamic_shared = inferred)] + pub fn manual_storage( + input: &[i32], + expected: &[i32], + mut output: coop::DisjointSlice>, + mut reports: coop::DisjointSlice>, + _mode: u32, + count: usize, + outside: i32, + ) { + let mut scratch = coop::TempStorage::new(); + with_storage( + input, + expected, + &mut output, + &mut reports, + count, + outside, + &mut scratch, + ); + } + + #[kernel] + #[launch_contract(domain = 1, block = (33, 1, 1), dynamic_shared = inferred)] + pub fn auto_storage( + input: &[i32], + expected: &[i32], + mut output: coop::DisjointSlice>, + mut reports: coop::DisjointSlice>, + _mode: u32, + count: usize, + outside: i32, + ) { + let mut scratch = coop::TempStorage::new().auto_sync(); + with_storage( + input, + expected, + &mut output, + &mut reports, + count, + outside, + &mut scratch, + ); + } +} + +fn oracle>( + input: &[T], + tile: usize, + recipe: u32, + count: usize, + outside: T, +) -> Vec { + let active = if matches!(recipe, 2 | 3 | 6) { + count + } else { + tile + }; + let mut expected = input.to_vec(); + for (source, target) in input + .chunks_exact(tile) + .zip(expected.chunks_exact_mut(tile)) + { + for i in 0..active { + if recipe < 4 { + if i != 0 { + target[i] = source[i] - source[i - 1]; + } else if matches!(recipe, 1 | 3) { + target[i] = source[i] - outside; + } + } else if i + 1 < active { + target[i] = source[i] - source[i + 1]; + } else if recipe == 5 { + target[i] = source[i] - outside; + } + } + } + expected +} + +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}; + + let invalid = match std::env::args().nth(1).as_deref() { + None => 0, + Some("count-over-capacity") => 1, + Some("count-max") => 2, + Some("missing-suffix-zero") => 3, + Some("missing-suffix-partial") => 4, + Some(other) => panic!("unknown test mode {other}"), + }; + init_device_contexts(0, 1)?; + let context = with_cuda_context(0, |context| context.clone())?; + let stream = context.default_stream(); + // SAFETY: typed declarations and embedded kernels are compiled together. + let module = unsafe { kernels::load_async(0)? }; + if invalid != 0 { + let count = match invalid { + 1 => 133, + 2 => usize::MAX, + 3 => 0, + _ => 5, + }; + let source = DeviceBuffer::from_host(&stream, &[20_i32; 132])?; + let expected = DeviceBuffer::from_host(&stream, &[0_i32; 132])?; + let target = DeviceBuffer::from_host(&stream, &[-1_i32; 132])?; + let reports = DeviceBuffer::from_host(&stream, &[0_u32; 33])?; + let plan = module.prepare_forms_i32(LaunchConfig1D::auto_shared(1, 33))?; + module + .forms_i32_async_owned( + &plan, source, expected, target, reports, 2, count, 0, invalid, + ) + .sync()?; + stream.synchronize()?; + panic!("invalid Adjacent Difference mode {invalid} unexpectedly completed"); + } + + let mut cases = 0; + let mut reports = 0; + let mut values = 0; + macro_rules! run { + ($ty:ty, $n:literal, $b:literal, $scale:literal, $prepare:ident, $launch:ident, $config:expr, $cases:expr $(, $extra:expr)?) => {{ + let input: Vec<$ty> = (0..3 * $b * $n).map(|i| { + let block = i / ($b * $n); + let slot = i % ($b * $n); + (20 + (7 * slot + 3 * block) % 31) as $ty / $scale as $ty + }).collect(); + for (mode, recipe, count, outside) in $cases { + let expected = oracle(&input, $b * $n, recipe, count, outside); + let source = DeviceBuffer::from_host(&stream, &input)?; + let expected_device = DeviceBuffer::from_host(&stream, &expected)?; + let target = DeviceBuffer::from_host(&stream, &vec![<$ty>::MAX; input.len()])?; + let words = DeviceBuffer::from_host(&stream, &[0_u32; 3 * $b])?; + let plan = module.$prepare($config)?; + let (_, _, target, words) = module.$launch(&plan, source, expected_device, target, words, + mode, count, outside $(, $extra)?).sync()?; + assert_eq!(target.to_host_vec(&stream)?, expected, "{} mode{mode} count{count}", stringify!($prepare)); + assert_eq!(words.to_host_vec(&stream)?, vec![CHECKS; 3 * $b], "{} mode{mode} count{count}", stringify!($prepare)); + cases += 1; + reports += 3 * $b; + values += input.len(); + } + }}; + } + run!( + i32, + 4, + 33, + 1, + prepare_forms_i32, + forms_i32_async_owned, + LaunchConfig1D::auto_shared(3, 33), + [ + (0, 0, 0, 0), + (1, 1, 0, -11), + (2, 2, 0, 0), + (2, 2, 1, 0), + (2, 2, 5, 0), + (2, 2, 132, 0), + (3, 3, 0, -11), + (3, 3, 5, -11), + (4, 4, 0, 0), + (5, 5, 0, 13), + (6, 6, 0, 0), + (6, 6, 1, 0), + (6, 6, 131, 0), + (6, 6, 132, 0) + ], + 0 + ); + run!( + i8, + 1, + 30, + 1, + prepare_xyz_n1, + xyz_n1_async_owned, + LaunchConfig3D::auto_shared((3, 1, 1), (5, 3, 2)), + [(0, 4, 0, 0_i8), (1, 1, 0, -7_i8)] + ); + run!( + f64, + 9, + 32, + 2, + prepare_tail_n9, + tail_n9_async_owned, + LaunchConfig1D::auto_shared(3, 32), + [(0, 3, 10, 0.5_f64), (1, 5, 0, -3.5_f64)] + ); + run!( + u16, + 1, + 1, + 1, + prepare_single_n1, + single_n1_async_owned, + LaunchConfig1D::auto_shared(3, 1), + [(0, 3, 0, 3_u16), (1, 5, 0, 3_u16)] + ); + run!( + i32, + 4, + 33, + 1, + prepare_manual_storage, + manual_storage_async_owned, + LaunchConfig1D::auto_shared(3, 33), + [(3, 3, 5, -11_i32)] + ); + run!( + i32, + 4, + 33, + 1, + prepare_auto_storage, + auto_storage_async_owned, + LaunchConfig1D::auto_shared(3, 33), + [(3, 3, 5, -11_i32)] + ); + println!( + "PASS: {cases} public Adjacent Difference cases, {values} values, {reports} lane reports" + ); + Ok(()) +} diff --git a/rust/cuda-coop/tests/check_load_store.py b/rust/cuda-coop/tests/check_load_store.py index e75be5e0a10a..69b7fa259711 100644 --- a/rust/cuda-coop/tests/check_load_store.py +++ b/rust/cuda-coop/tests/check_load_store.py @@ -1318,6 +1318,244 @@ ), ) +# Adjacent Difference keeps counted suffixes initialized and requires uniform +# neighbor/count controls. Explicit scratch retains its existing completion rules. +adjacent = "let result=coop::adjacent_difference(block,&values,Default::default());" +adjacent_manual = "values=coop::adjacent_difference(block,&values,coop::AdjacentDifferenceOptions::default().temp_storage(&mut scratch));" +cases.update( + adjacent_generic_defaults=( + load + "fn forward(" + "block:coop::ThreadGroup,input:&coop::ThreadData)->coop::ThreadData " + "where coop::Alignment:coop::SupportedAlignment {" + "coop::adjacent_difference(block,input,Default::default())}" + "let default_result=forward(block,&values);" + "let mut one=coop::ThreadData::<_,1>::new();one.write(0,7_i32);" + "let one_result=forward(block,&mut one);" + "let mut aligned=coop::ThreadData::<_,4,16>::new();" + "aligned.write(0,0_i32);aligned.write(1,1);aligned.write(2,2);aligned.write(3,3);" + "let mut aligned_result:coop::ThreadData<_,4,16>=forward(block,&mut aligned);" + "aligned_result.write(0,default_result[0]^one_result[0]);" + + store.replace("&values", "&aligned_result"), + None, + ), + adjacent_recipes_partial_block=( + load + "let a=coop::adjacent_difference(block,&values,Default::default());" + "let b=coop::adjacent_difference(block,&values,coop::AdjacentDifferenceOptions::default().tile_predecessor_item(7));" + "let c=coop::adjacent_difference(block,&values,coop::AdjacentDifferenceOptions::default().valid_items(valid));" + "let d=coop::adjacent_difference(block,&values,coop::AdjacentDifferenceOptions::default().tile_predecessor_item(7).valid_items(0));" + "let e=coop::adjacent_difference(block,&values,coop::AdjacentDifferenceOptions::default().direction(coop::Right));" + "let f=coop::adjacent_difference(block,&values,coop::AdjacentDifferenceOptions::default().tile_successor_item(7).direction(coop::Right));" + "let g=coop::adjacent_difference(block,&values,coop::AdjacentDifferenceOptions::default().valid_items(valid).direction(coop::Right));" + "values.write(0,a[0]^b[0]^c[0]^d[0]^e[0]^f[0]^g[0]);" + store, + None, + ), + adjacent_right_count_successor=( + load + "let result=coop::adjacent_difference(block,&values," + "coop::AdjacentDifferenceOptions::default().direction(coop::Right).valid_items(512).tile_successor_item(0));", + "AdjacentDifferenceDirection", + ), + adjacent_predecessor_then_right=( + load + "let result=coop::adjacent_difference(block,&values," + "coop::AdjacentDifferenceOptions::default().tile_predecessor_item(0).direction(coop::Right));", + "AdjacentDifferenceDirection", + ), + adjacent_left_successor=( + load + "let result=coop::adjacent_difference(block,&values," + "coop::AdjacentDifferenceOptions::default().tile_successor_item(0));", + "AdjacentDifferenceDirection", + ), + adjacent_neighbor_type=( + load + "let result=coop::adjacent_difference(block,&values," + "coop::AdjacentDifferenceOptions::default().tile_predecessor_item(0_i64));", + "mismatched types", + ), + adjacent_scalar=( + load + "let result=coop::adjacent_difference(block,&3_i32,Default::default());", + "mismatched types", + ), + adjacent_warp=( + load + + "let result=coop::adjacent_difference(coop::this_warp(),&values,Default::default());", + "mismatched types", + ), + adjacent_partition=( + load + + "let result=coop::adjacent_difference(block.group_by::<2>(),&values,Default::default());", + "mismatched types", + ), + adjacent_varying_count=( + load + "let result=coop::adjacent_difference(block,&values," + "coop::AdjacentDifferenceOptions::default().valid_items(thread::threadIdx_x() as usize));", + "not proved block-uniform", + ), + adjacent_varying_predecessor=( + load + "let result=coop::adjacent_difference(block,&values," + "coop::AdjacentDifferenceOptions::default().tile_predecessor_item(thread::threadIdx_x() as i32));", + "not proved block-uniform", + ), + adjacent_lane_divergence=( + load + "if thread::threadIdx_x()&1==0 {" + adjacent + "}", + "lane-varying", + ), + adjacent_result_control=( + load + adjacent + "if result[0]==0 {" + adjacent + "}", + "lane-varying", + ), + adjacent_storage_missing_barrier=( + load + new_storage + adjacent_manual + adjacent_manual, + "requires a full-block barrier", + ), + adjacent_storage_completed=( + load + new_storage + adjacent_manual + sync + adjacent_manual + store, + None, + ), + adjacent_storage_auto=( + load + + new_storage.replace("new()", "new().auto_sync()") + + adjacent_manual + + adjacent_manual + + store, + None, + ), + adjacent_read_epoch=( + reload.replace(".valid_items(valid)", ".valid_items(valid).oob_default(0)") + + adjacent + + store, + "requires a barrier covering both accesses", + ), + adjacent_read_completed=( + reload.replace(".valid_items(valid)", ".valid_items(valid).oob_default(0)") + + adjacent + + sync + + store, + None, + ), +) + +# Discontinuity preserves input and supplies one or two independent i32 results. +# Scratch completion and unrelated global-memory completion remain separate. +discontinuity = "let result=coop::discontinuity(block,&values,Default::default());" +discontinuity_pair = "let (_heads,_tails)=coop::discontinuity(block,&values,coop::DiscontinuityOptions::default().mode(coop::HeadsAndTails));" +discontinuity_manual = "let (heads,_tails)=coop::discontinuity(block,&values,coop::DiscontinuityOptions::default().mode(coop::HeadsAndTails).temp_storage(&mut scratch));values=heads;" +cases.update( + discontinuity_generic_defaults=( + load + "fn forward(" + "block:coop::ThreadGroup,input:&coop::ThreadData)->coop::ThreadData " + "where coop::Alignment:coop::SupportedAlignment {" + "coop::discontinuity(block,input,Default::default())}" + "fn project(" + "block:coop::ThreadGroup,input:&coop::ThreadData,mode:M)->M::Output " + "where coop::Alignment:coop::SupportedAlignment,M:coop::DiscontinuityMode {" + "coop::discontinuity(block,input,coop::DiscontinuityOptions::default().mode(mode))}" + "fn pair(" + "block:coop::ThreadGroup,input:&coop::ThreadData)->(coop::ThreadData,coop::ThreadData) " + "where coop::Alignment:coop::SupportedAlignment {project(block,input,coop::HeadsAndTails)}" + "let default_result=forward(block,&values);" + "let mut one=coop::ThreadData::<_,1>::new();one.write(0,7_u8);" + "let one_result=forward(block,&mut one);" + "let mut aligned=coop::ThreadData::<_,4,16>::new();" + "aligned.write(0,0_f64);aligned.write(1,1.0);aligned.write(2,1.0);aligned.write(3,2.0);" + "let (mut heads,tails):(coop::ThreadData<_,4,16>,coop::ThreadData<_,4,16>)=pair(block,&mut aligned);" + "heads.write(0,default_result[0]^one_result[0]^tails[0]);" + + store.replace("&values", "&heads"), + None, + ), + discontinuity_recipes_partial_block=( + load + "let a=coop::discontinuity(block,&values,Default::default());" + "let b=coop::discontinuity(block,&values,coop::DiscontinuityOptions::default().tile_predecessor_item(7));" + "let c=coop::discontinuity(block,&values,coop::DiscontinuityOptions::default().mode(coop::Tails));" + "let d=coop::discontinuity(block,&values,coop::DiscontinuityOptions::default().tile_successor_item(7).mode(coop::Tails));" + "let (eh,et)=coop::discontinuity(block,&values,coop::DiscontinuityOptions::default().mode(coop::HeadsAndTails));" + "let (fh,ft)=coop::discontinuity(block,&values,coop::DiscontinuityOptions::default().mode(coop::HeadsAndTails).tile_predecessor_item(7));" + "let (gh,gt)=coop::discontinuity(block,&values,coop::DiscontinuityOptions::default().tile_successor_item(7).mode(coop::HeadsAndTails));" + "let (hh,ht)=coop::discontinuity(block,&values,coop::DiscontinuityOptions::default().tile_successor_item(7).tile_predecessor_item(7).mode(coop::HeadsAndTails));" + "values.write(0,a[0]^b[0]^c[0]^d[0]^eh[0]^et[0]^fh[0]^ft[0]^gh[0]^gt[0]^hh[0]^ht[0]);" + + store, + None, + ), + discontinuity_heads_successor=( + load + "let result=coop::discontinuity(block,&values," + "coop::DiscontinuityOptions::default().tile_successor_item(0));", + "DiscontinuityMode", + ), + discontinuity_predecessor_then_tails=( + load + "let result=coop::discontinuity(block,&values," + "coop::DiscontinuityOptions::default().tile_predecessor_item(0).mode(coop::Tails));", + "DiscontinuityMode", + ), + discontinuity_neighbor_type=( + load + "let result=coop::discontinuity(block,&values," + "coop::DiscontinuityOptions::default().tile_predecessor_item(0_i64));", + "mismatched types", + ), + discontinuity_no_count=( + load + "let result=coop::discontinuity(block,&values," + "coop::DiscontinuityOptions::default().valid_items(valid));", + "no method named `valid_items`", + ), + discontinuity_scalar=( + load + "let result=coop::discontinuity(block,&3_i32,Default::default());", + "mismatched types", + ), + discontinuity_warp=( + load + + "let result=coop::discontinuity(coop::this_warp(),&values,Default::default());", + "mismatched types", + ), + discontinuity_partition=( + load + + "let result=coop::discontinuity(block.group_by::<2>(),&values,Default::default());", + "mismatched types", + ), + discontinuity_varying_predecessor=( + load + "let result=coop::discontinuity(block,&values," + "coop::DiscontinuityOptions::default().tile_predecessor_item(thread::threadIdx_x() as i32));", + "not proved block-uniform", + ), + discontinuity_varying_successor_pair=( + load + "let (_heads,_tails)=coop::discontinuity(block,&values," + "coop::DiscontinuityOptions::default().mode(coop::HeadsAndTails).tile_successor_item(thread::threadIdx_x() as i32));", + "not proved block-uniform", + ), + discontinuity_lane_divergence=( + load + "if thread::threadIdx_x()&1==0 {" + discontinuity + "}", + "lane-varying", + ), + discontinuity_result_control=( + load + discontinuity + "if result[0]==0 {" + discontinuity + "}", + "lane-varying", + ), + discontinuity_storage_missing_barrier=( + load + new_storage + discontinuity_manual + discontinuity_manual, + "requires a full-block barrier", + ), + discontinuity_storage_completed=( + load + new_storage + discontinuity_manual + sync + discontinuity_manual + store, + None, + ), + discontinuity_storage_auto=( + load + + new_storage.replace("new()", "new().auto_sync()") + + discontinuity_manual + + discontinuity_manual + + store, + None, + ), + discontinuity_read_epoch=( + reload.replace(".valid_items(valid)", ".valid_items(valid).oob_default(0)") + + discontinuity_pair + + store, + "requires a barrier covering both accesses", + ), + discontinuity_read_completed=( + reload.replace(".valid_items(valid)", ".valid_items(valid).oob_default(0)") + + discontinuity_pair + + sync + + store, + None, + ), +) + if args.case: unknown = set(args.case).difference(cases) if unknown: diff --git a/rust/cuda-coop/tests/discontinuity/.gitignore b/rust/cuda-coop/tests/discontinuity/.gitignore new file mode 100644 index 000000000000..74bffa8af4a4 --- /dev/null +++ b/rust/cuda-coop/tests/discontinuity/.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/discontinuity/Cargo.lock b/rust/cuda-coop/tests/discontinuity/Cargo.lock new file mode 100644 index 000000000000..a56b9f7c5f1d --- /dev/null +++ b/rust/cuda-coop/tests/discontinuity/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-discontinuity-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/discontinuity/Cargo.toml b/rust/cuda-coop/tests/discontinuity/Cargo.toml new file mode 100644 index 000000000000..ed813ae83105 --- /dev/null +++ b/rust/cuda-coop/tests/discontinuity/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-discontinuity-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/discontinuity/src/main.rs b/rust/cuda-coop/tests/discontinuity/src/main.rs new file mode 100644 index 000000000000..0982d8464ebc --- /dev/null +++ b/rust/cuda-coop/tests/discontinuity/src/main.rs @@ -0,0 +1,1023 @@ +// 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; + +const CHECKS: u32 = 15; + +#[cuda_module] +mod kernels { + use super::*; + + fn observe( + flags: &coop::ThreadData, + expected: &[i32], + offset: usize, + ) -> bool + where + coop::Alignment: coop::SupportedAlignment, + { + let mut valid = flags.len() == N; + let mut i = 0; + while i < N { + valid &= flags.get(i) == Some(expected[offset + i]); + i += 1; + } + valid + } + + fn retained( + flags: &coop::ThreadData, + expected: &[i32], + offset: usize, + slot: usize, + marker: i32, + ) -> bool + where + coop::Alignment: coop::SupportedAlignment, + { + let mut valid = flags.len() == N; + let mut i = 0; + while i < N { + let value = if i == slot { + marker + } else { + expected[offset + i] + }; + valid &= flags.get(i) == Some(value); + i += 1; + } + valid + } + + #[allow(clippy::too_many_arguments)] + fn exercise( + input: &[T], + expected: &[i32], + heads_output: &mut coop::DisjointSlice>, + tails_output: &mut coop::DisjointSlice>, + preserved: &mut coop::DisjointSlice>, + reports: &mut coop::DisjointSlice>, + rank: usize, + mode: u32, + predecessor: T, + successor: T, + replacement: T, + invalid: u32, + ) where + coop::Alignment: coop::SupportedAlignment, + { + let block = coop::this_block(); + let block_index = block.rank_at::() as usize; + let offset = block_index * B * N; + let local = offset + rank * N; + let len = input.len(); + let mut values = coop::ThreadData::<_, N, A>::new(); + if FORM == 8 && invalid != 0 { + let mut i = 0; + while i < N { + if !(rank == B - 1 && i == N - 1) { + values.write(i, input[local + i]); + } + i += 1; + } + if invalid == 1 { + let _ = coop::discontinuity(block, &values, Default::default()); + } else { + let _ = coop::discontinuity( + block, + &values, + coop::DiscontinuityOptions::default().mode(coop::HeadsAndTails), + ); + } + // A lost preflight must reach the host's unexpected-return panic. + // No result read or value-dependent device assertion can mimic it. + let mut report = coop::ThreadData::<_, 1>::new(); + report.write(0, 99_u32); + coop::store( + block, + reports, + &report, + coop::StoreOptions::default().offset(block_index * B), + ); + return; + } + coop::load( + block, + input, + &mut values, + coop::LoadOptions::default().offset(offset), + ); + let mut changed = values; + changed.write(0, replacement); + let mut word: u32; + macro_rules! single { + ($options:expr, $plane:expr) => {{ + let plane = $plane * len; + let mut first = coop::discontinuity(block, &values, $options); + word = observe(&first, expected, plane + local) as u32; + first.write(0, 7); + let second = coop::discontinuity(block, &changed, $options); + word |= (observe(&second, expected, 2 * len + plane + local) as u32) << 1; + word |= (retained(&first, expected, plane + local, 0, 7) as u32) << 2; + word |= 8; + if $plane == 0 { + coop::store( + block, + heads_output, + &second, + coop::StoreOptions::default().offset(offset), + ); + } else { + coop::store( + block, + tails_output, + &second, + coop::StoreOptions::default().offset(offset), + ); + } + }}; + } + macro_rules! pair { + ($options:expr, $manual:expr) => {{ + let (mut heads, mut tails) = coop::discontinuity(block, &values, $options); + word = (observe(&heads, expected, local) && observe(&tails, expected, len + local)) + as u32; + heads.write(0, 7); + let separate = observe(&tails, expected, len + local); + tails.write(N - 1, 9); + if $manual { + thread::sync_threads(); + } + let (next_heads, next_tails) = coop::discontinuity(block, &changed, $options); + word |= ((observe(&next_heads, expected, 2 * len + local) + && observe(&next_tails, expected, 3 * len + local)) as u32) + << 1; + word |= ((retained(&heads, expected, local, 0, 7) + && retained(&tails, expected, len + local, N - 1, 9)) as u32) + << 2; + word |= (separate as u32) << 3; + if $manual { + thread::sync_threads(); + } + coop::store( + block, + heads_output, + &next_heads, + coop::StoreOptions::default().offset(offset), + ); + coop::store( + block, + tails_output, + &next_tails, + coop::StoreOptions::default().offset(offset), + ); + }}; + } + if FORM == 9 { + let mut scratch = coop::TempStorage::new(); + pair!( + coop::DiscontinuityOptions::default() + .mode(coop::HeadsAndTails) + .tile_predecessor_item(predecessor) + .tile_successor_item(successor) + .temp_storage(&mut scratch), + true + ); + } else if FORM == 10 { + let mut scratch = coop::TempStorage::new().auto_sync(); + pair!( + coop::DiscontinuityOptions::default() + .mode(coop::HeadsAndTails) + .tile_predecessor_item(predecessor) + .tile_successor_item(successor) + .temp_storage(&mut scratch), + false + ); + } else { + let recipe = if FORM == 8 { mode } else { FORM as u32 }; + match recipe { + 0 => single!(Default::default(), 0), + 1 => single!( + coop::DiscontinuityOptions::default().tile_predecessor_item(predecessor), + 0 + ), + 2 => single!(coop::DiscontinuityOptions::default().mode(coop::Tails), 1), + 3 => single!( + coop::DiscontinuityOptions::default() + .mode(coop::Tails) + .tile_successor_item(successor), + 1 + ), + 4 => pair!( + coop::DiscontinuityOptions::default().mode(coop::HeadsAndTails), + false + ), + 5 => pair!( + coop::DiscontinuityOptions::default() + .mode(coop::HeadsAndTails) + .tile_predecessor_item(predecessor), + false + ), + 6 => pair!( + coop::DiscontinuityOptions::default() + .mode(coop::HeadsAndTails) + .tile_successor_item(successor), + false + ), + 7 => pair!( + coop::DiscontinuityOptions::default() + .mode(coop::HeadsAndTails) + .tile_predecessor_item(predecessor) + .tile_successor_item(successor), + false + ), + _ => panic!(), + } + } + // Store original numeric bytes, including NaN payloads and zero signs. + coop::store( + block, + preserved, + &values, + coop::StoreOptions::default().offset(offset), + ); + let mut report = coop::ThreadData::<_, 1>::new(); + report.write(0, word); + coop::store( + block, + reports, + &report, + coop::StoreOptions::default().offset(block_index * B), + ); + } + + #[allow(clippy::too_many_arguments)] + #[kernel] + #[launch_contract(domain = 1, block = (33, 1, 1), dynamic_shared = inferred)] + pub fn forms_i32( + input: &[i32], + expected: &[i32], + mut heads_output: coop::DisjointSlice>, + mut tails_output: coop::DisjointSlice>, + mut preserved: coop::DisjointSlice>, + mut reports: coop::DisjointSlice>, + mode: u32, + predecessor: i32, + successor: i32, + replacement: i32, + invalid: u32, + ) { + let rank = thread::threadIdx_x(); + exercise::<_, 4, 16, 33, 8>( + input, + expected, + &mut heads_output, + &mut tails_output, + &mut preserved, + &mut reports, + rank as usize, + mode, + predecessor, + successor, + replacement, + invalid, + ); + } + + #[allow(clippy::too_many_arguments)] + #[kernel] + #[launch_contract(domain = 3, block = (4, 3, 2), dynamic_shared = inferred)] + pub fn xyz_f64( + input: &[f64], + expected: &[i32], + mut heads_output: coop::DisjointSlice>, + mut tails_output: coop::DisjointSlice>, + mut preserved: coop::DisjointSlice>, + mut reports: coop::DisjointSlice>, + mode: u32, + predecessor: f64, + successor: f64, + replacement: f64, + invalid: u32, + ) { + let rank = thread::threadIdx_x() + 4 * (thread::threadIdx_y() + 3 * thread::threadIdx_z()); + exercise::<_, 4, 1, 24, 7>( + input, + expected, + &mut heads_output, + &mut tails_output, + &mut preserved, + &mut reports, + rank as usize, + mode, + predecessor, + successor, + replacement, + invalid, + ); + } + + #[allow(clippy::too_many_arguments)] + #[kernel] + #[launch_contract(domain = 1, block = (32, 1, 1), dynamic_shared = inferred)] + pub fn tail_f32( + input: &[f32], + expected: &[i32], + mut heads_output: coop::DisjointSlice>, + mut tails_output: coop::DisjointSlice>, + mut preserved: coop::DisjointSlice>, + mut reports: coop::DisjointSlice>, + mode: u32, + predecessor: f32, + successor: f32, + replacement: f32, + invalid: u32, + ) { + let rank = thread::threadIdx_x(); + exercise::<_, 9, 1, 32, 4>( + input, + expected, + &mut heads_output, + &mut tails_output, + &mut preserved, + &mut reports, + rank as usize, + mode, + predecessor, + successor, + replacement, + invalid, + ); + } + + #[allow(clippy::too_many_arguments)] + #[kernel] + #[launch_contract(domain = 1, block = (33, 1, 1), dynamic_shared = inferred)] + pub fn heads_i8( + input: &[i8], + expected: &[i32], + mut heads_output: coop::DisjointSlice>, + mut tails_output: coop::DisjointSlice>, + mut preserved: coop::DisjointSlice>, + mut reports: coop::DisjointSlice>, + mode: u32, + predecessor: i8, + successor: i8, + replacement: i8, + invalid: u32, + ) { + let rank = thread::threadIdx_x(); + exercise::<_, 1, 1, 33, 1>( + input, + expected, + &mut heads_output, + &mut tails_output, + &mut preserved, + &mut reports, + rank as usize, + mode, + predecessor, + successor, + replacement, + invalid, + ); + } + + #[allow(clippy::too_many_arguments)] + #[kernel] + #[launch_contract(domain = 1, block = (1, 1, 1), dynamic_shared = inferred)] + pub fn single_u8( + input: &[u8], + expected: &[i32], + mut heads_output: coop::DisjointSlice>, + mut tails_output: coop::DisjointSlice>, + mut preserved: coop::DisjointSlice>, + mut reports: coop::DisjointSlice>, + mode: u32, + predecessor: u8, + successor: u8, + replacement: u8, + invalid: u32, + ) { + let rank = thread::threadIdx_x(); + exercise::<_, 1, 1, 1, 4>( + input, + expected, + &mut heads_output, + &mut tails_output, + &mut preserved, + &mut reports, + rank as usize, + mode, + predecessor, + successor, + replacement, + invalid, + ); + } + + #[allow(clippy::too_many_arguments)] + #[kernel] + #[launch_contract(domain = 3, block = (4, 3, 2), dynamic_shared = inferred)] + pub fn xyz_i16( + input: &[i16], + expected: &[i32], + mut heads_output: coop::DisjointSlice>, + mut tails_output: coop::DisjointSlice>, + mut preserved: coop::DisjointSlice>, + mut reports: coop::DisjointSlice>, + mode: u32, + predecessor: i16, + successor: i16, + replacement: i16, + invalid: u32, + ) { + let rank = thread::threadIdx_x() + 4 * (thread::threadIdx_y() + 3 * thread::threadIdx_z()); + exercise::<_, 4, 1, 24, 3>( + input, + expected, + &mut heads_output, + &mut tails_output, + &mut preserved, + &mut reports, + rank as usize, + mode, + predecessor, + successor, + replacement, + invalid, + ); + } + + #[allow(clippy::too_many_arguments)] + #[kernel] + #[launch_contract(domain = 1, block = (32, 1, 1), dynamic_shared = inferred)] + pub fn tail_u16( + input: &[u16], + expected: &[i32], + mut heads_output: coop::DisjointSlice>, + mut tails_output: coop::DisjointSlice>, + mut preserved: coop::DisjointSlice>, + mut reports: coop::DisjointSlice>, + mode: u32, + predecessor: u16, + successor: u16, + replacement: u16, + invalid: u32, + ) { + let rank = thread::threadIdx_x(); + exercise::<_, 9, 1, 32, 5>( + input, + expected, + &mut heads_output, + &mut tails_output, + &mut preserved, + &mut reports, + rank as usize, + mode, + predecessor, + successor, + replacement, + invalid, + ); + } + + #[allow(clippy::too_many_arguments)] + #[kernel] + #[launch_contract(domain = 1, block = (1, 1, 1), dynamic_shared = inferred)] + pub fn single_u32( + input: &[u32], + expected: &[i32], + mut heads_output: coop::DisjointSlice>, + mut tails_output: coop::DisjointSlice>, + mut preserved: coop::DisjointSlice>, + mut reports: coop::DisjointSlice>, + mode: u32, + predecessor: u32, + successor: u32, + replacement: u32, + invalid: u32, + ) { + let rank = thread::threadIdx_x(); + exercise::<_, 4, 1, 1, 7>( + input, + expected, + &mut heads_output, + &mut tails_output, + &mut preserved, + &mut reports, + rank as usize, + mode, + predecessor, + successor, + replacement, + invalid, + ); + } + + #[allow(clippy::too_many_arguments)] + #[kernel] + #[launch_contract(domain = 1, block = (32, 1, 1), dynamic_shared = inferred)] + pub fn heads_i64( + input: &[i64], + expected: &[i32], + mut heads_output: coop::DisjointSlice>, + mut tails_output: coop::DisjointSlice>, + mut preserved: coop::DisjointSlice>, + mut reports: coop::DisjointSlice>, + mode: u32, + predecessor: i64, + successor: i64, + replacement: i64, + invalid: u32, + ) { + let rank = thread::threadIdx_x(); + exercise::<_, 1, 1, 32, 0>( + input, + expected, + &mut heads_output, + &mut tails_output, + &mut preserved, + &mut reports, + rank as usize, + mode, + predecessor, + successor, + replacement, + invalid, + ); + } + + #[allow(clippy::too_many_arguments)] + #[kernel] + #[launch_contract(domain = 1, block = (33, 1, 1), dynamic_shared = inferred)] + pub fn tails_u64( + input: &[u64], + expected: &[i32], + mut heads_output: coop::DisjointSlice>, + mut tails_output: coop::DisjointSlice>, + mut preserved: coop::DisjointSlice>, + mut reports: coop::DisjointSlice>, + mode: u32, + predecessor: u64, + successor: u64, + replacement: u64, + invalid: u32, + ) { + let rank = thread::threadIdx_x(); + exercise::<_, 4, 1, 33, 2>( + input, + expected, + &mut heads_output, + &mut tails_output, + &mut preserved, + &mut reports, + rank as usize, + mode, + predecessor, + successor, + replacement, + invalid, + ); + } + + #[allow(clippy::too_many_arguments)] + #[kernel] + #[launch_contract(domain = 1, block = (33, 1, 1), dynamic_shared = inferred)] + pub fn manual_storage( + input: &[i32], + expected: &[i32], + mut heads_output: coop::DisjointSlice>, + mut tails_output: coop::DisjointSlice>, + mut preserved: coop::DisjointSlice>, + mut reports: coop::DisjointSlice>, + mode: u32, + predecessor: i32, + successor: i32, + replacement: i32, + invalid: u32, + ) { + let rank = thread::threadIdx_x(); + exercise::<_, 4, 16, 33, 9>( + input, + expected, + &mut heads_output, + &mut tails_output, + &mut preserved, + &mut reports, + rank as usize, + mode, + predecessor, + successor, + replacement, + invalid, + ); + } + + #[allow(clippy::too_many_arguments)] + #[kernel] + #[launch_contract(domain = 1, block = (33, 1, 1), dynamic_shared = inferred)] + pub fn auto_storage( + input: &[i32], + expected: &[i32], + mut heads_output: coop::DisjointSlice>, + mut tails_output: coop::DisjointSlice>, + mut preserved: coop::DisjointSlice>, + mut reports: coop::DisjointSlice>, + mode: u32, + predecessor: i32, + successor: i32, + replacement: i32, + invalid: u32, + ) { + let rank = thread::threadIdx_x(); + exercise::<_, 4, 16, 33, 10>( + input, + expected, + &mut heads_output, + &mut tails_output, + &mut preserved, + &mut reports, + rank as usize, + mode, + predecessor, + successor, + replacement, + invalid, + ); + } +} + +fn oracle( + input: &[T], + tile: usize, + recipe: u32, + pred: T, + succ: T, +) -> Vec { + let mut flags = vec![-17; 2 * input.len()]; + for (block, values) in input.chunks_exact(tile).enumerate() { + for i in 0..tile { + if !(2..4).contains(&recipe) { + flags[block * tile + i] = if i != 0 { + values[i - 1] != values[i] + } else if matches!(recipe, 1 | 5 | 7) { + pred != values[i] + } else { + true + } as i32; + } + if recipe >= 2 { + flags[input.len() + block * tile + i] = if i + 1 != tile { + values[i] != values[i + 1] + } else if matches!(recipe, 3 | 6 | 7) { + values[i] != succ + } else { + true + } as i32; + } + } + } + flags +} + +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}; + let invalid = match std::env::args().nth(1).as_deref() { + None => 0, + Some("missing-head-input") => 1, + Some("missing-pair-input") => 2, + Some(other) => panic!("unknown test mode {other}"), + }; + init_device_contexts(0, 1)?; + let context = with_cuda_context(0, |context| context.clone())?; + let stream = context.default_stream(); + // SAFETY: typed declarations and embedded kernels are compiled together. + let module = unsafe { kernels::load_async(0)? }; + if invalid != 0 { + let source = DeviceBuffer::from_host(&stream, &[3_i32; 132])?; + let expected = DeviceBuffer::from_host(&stream, &[0_i32; 528])?; + let heads = DeviceBuffer::from_host(&stream, &[-17_i32; 132])?; + let tails = DeviceBuffer::from_host(&stream, &[-17_i32; 132])?; + let preserved = DeviceBuffer::from_host(&stream, &[0_i32; 132])?; + let report = DeviceBuffer::from_host(&stream, &[0_u32; 33])?; + let plan = module.prepare_forms_i32(LaunchConfig1D::auto_shared(1, 33))?; + module + .forms_i32_async_owned( + &plan, source, expected, heads, tails, preserved, report, 0, 3, 3, 4, invalid, + ) + .sync()?; + stream.synchronize()?; + panic!("invalid Discontinuity mode {invalid} unexpectedly completed"); + } + let mut cases = 0; + let mut flags = 0; + let mut values = 0; + let mut reports = 0; + macro_rules! run { + ($ty:ty, $n:literal, $b:literal, $prepare:ident, $launch:ident, $config:expr, + $recipe:expr, $input:expr, $pred:expr, $succ:expr, $replacement:expr, $bits:expr) => {{ + let input: Vec<$ty> = $input; + let recipe = $recipe; + let (pred, succ, replacement): ($ty, $ty, $ty) = ($pred, $succ, $replacement); + let mut changed = input.clone(); + for slot in changed.iter_mut().step_by($n) { + *slot = replacement; + } + let mut expected = oracle(&input, $b * $n, recipe, pred, succ); + let second = oracle(&changed, $b * $n, recipe, pred, succ); + expected.extend_from_slice(&second); + let source = DeviceBuffer::from_host(&stream, &input)?; + let expected_device = DeviceBuffer::from_host(&stream, &expected)?; + let heads = DeviceBuffer::from_host(&stream, &vec![-17_i32; input.len()])?; + let tails = DeviceBuffer::from_host(&stream, &vec![-17_i32; input.len()])?; + let preserved = DeviceBuffer::from_host(&stream, &vec![0 as $ty; input.len()])?; + let words = DeviceBuffer::from_host(&stream, &[0_u32; 3 * $b])?; + let plan = module.$prepare($config)?; + let (source, _, heads, tails, preserved, words) = module + .$launch( + &plan, + source, + expected_device, + heads, + tails, + preserved, + words, + recipe, + pred, + succ, + replacement, + 0, + ) + .sync()?; + assert_eq!( + [heads.to_host_vec(&stream)?, tails.to_host_vec(&stream)?].concat(), + second, + "{} recipe{recipe}", + stringify!($prepare) + ); + let bit_pattern = $bits; + let original_bits: Vec<_> = input.iter().copied().map(bit_pattern).collect(); + assert_eq!( + source + .to_host_vec(&stream)? + .into_iter() + .map(bit_pattern) + .collect::>(), + original_bits + ); + assert_eq!( + preserved + .to_host_vec(&stream)? + .into_iter() + .map(bit_pattern) + .collect::>(), + original_bits + ); + assert_eq!( + words.to_host_vec(&stream)?, + vec![CHECKS; 3 * $b], + "{} recipe{recipe}", + stringify!($prepare) + ); + cases += 1; + flags += input.len() * if recipe < 4 { 1 } else { 2 }; + values += input.len(); + reports += 3 * $b; + }}; + } + for recipe in 0..8 { + run!( + i32, + 4, + 33, + prepare_forms_i32, + forms_i32_async_owned, + LaunchConfig1D::auto_shared(3, 33), + recipe, + (0..3 * 33 * 4) + .map(|i| (i / (33 * 4) + (i % (33 * 4)) / 3) % 5) + .collect(), + 0, + 3, + 11, + |x: i32| x + ); + } + run!( + f64, + 4, + 24, + prepare_xyz_f64, + xyz_f64_async_owned, + LaunchConfig3D::auto_shared((3, 1, 1), (4, 3, 2)), + 7, + (0..288) + .map(|i| 1.0 + ((i / 3) % 5) as f64 * 1.0e-10) + .collect(), + 1.0, + 1.0, + 2.0, + f64::to_bits + ); + run!( + f64, + 4, + 24, + prepare_xyz_f64, + xyz_f64_async_owned, + LaunchConfig3D::auto_shared((3, 1, 1), (4, 3, 2)), + 7, + (0..288) + .map(|i| [ + f64::from_bits(0x7ff8_0000_0000_0042), + f64::NAN, + 2.0, + 0.0, + -0.0, + 2.0, + 2.0, + 3.0 + ][i % 8]) + .collect(), + f64::NAN, + -0.0, + 4.0, + f64::to_bits + ); + run!( + f32, + 9, + 32, + prepare_tail_f32, + tail_f32_async_owned, + LaunchConfig1D::auto_shared(3, 32), + 4, + (0..864) + .map(|i| [ + f32::from_bits(0x7fc0_0042), + f32::NAN, + 2.0, + 2.0, + 3.0, + 3.0, + 3.0, + 2.0, + 0.0, + -0.0 + ][i % 10]) + .collect(), + 0.0, + 0.0, + 4.0, + f32::to_bits + ); + run!( + i8, + 1, + 33, + prepare_heads_i8, + heads_i8_async_owned, + LaunchConfig1D::auto_shared(3, 33), + 1, + (0..99).map(|i| -7 + ((i / 3) % 5) as i8).collect(), + -7, + -3, + -11, + |x: i8| x + ); + run!( + u8, + 1, + 1, + prepare_single_u8, + single_u8_async_owned, + LaunchConfig1D::auto_shared(3, 1), + 4, + vec![2, 3, 4], + 2, + 4, + 7, + |x: u8| x + ); + run!( + i16, + 4, + 24, + prepare_xyz_i16, + xyz_i16_async_owned, + LaunchConfig3D::auto_shared((3, 1, 1), (4, 3, 2)), + 3, + (0..3 * 24 * 4) + .map(|i| ((i / (24 * 4) + (i % (24 * 4)) / 3) % 5) as i16) + .collect(), + 0, + 1, + 17, + |x: i16| x + ); + run!( + u16, + 9, + 32, + prepare_tail_u16, + tail_u16_async_owned, + LaunchConfig1D::auto_shared(3, 32), + 5, + (0..3 * 32 * 9) + .map(|i| ((i / (32 * 9) + (i % (32 * 9)) / 3) % 5) as u16) + .collect(), + 0, + 0, + 17, + |x: u16| x + ); + run!( + u32, + 4, + 1, + prepare_single_u32, + single_u32_async_owned, + LaunchConfig1D::auto_shared(3, 1), + 7, + vec![2, 2, 3, 3, 2, 4, 4, 3, 5, 5, 5, 5], + 2, + 3, + 7, + |x: u32| x + ); + run!( + i64, + 1, + 32, + prepare_heads_i64, + heads_i64_async_owned, + LaunchConfig1D::auto_shared(3, 32), + 0, + (0..96) + .map(|i| (1_i64 << 40) + ((i / 3) % 3) as i64 * (1_i64 << 32)) + .collect(), + 0, + 0, + -1, + |x: i64| x + ); + run!( + u64, + 4, + 33, + prepare_tails_u64, + tails_u64_async_owned, + LaunchConfig1D::auto_shared(3, 33), + 2, + (0..396) + .map(|i| (1_u64 << 40) + ((i / 3) % 3) as u64 * (1_u64 << 32)) + .collect(), + 0, + 0, + 1, + |x: u64| x + ); + run!( + i32, + 4, + 33, + prepare_manual_storage, + manual_storage_async_owned, + LaunchConfig1D::auto_shared(3, 33), + 7, + (0..3 * 33 * 4) + .map(|i| (i / (33 * 4) + (i % (33 * 4)) / 3) % 5) + .collect(), + 0, + 3, + 11, + |x: i32| x + ); + run!( + i32, + 4, + 33, + prepare_auto_storage, + auto_storage_async_owned, + LaunchConfig1D::auto_shared(3, 33), + 7, + (0..3 * 33 * 4) + .map(|i| (i / (33 * 4) + (i % (33 * 4)) / 3) % 5) + .collect(), + 0, + 3, + 11, + |x: i32| x + ); + assert_eq!(cases, 20); + println!( + "discontinuity: {cases} cases, {flags} flags, {values} preserved values, {reports} lane reports passed" + ); + Ok(()) +} From 290ec8b2552ab0df08285e69bf4d1a993d49fde6 Mon Sep 17 00:00:00 2001 From: Trent Nelson Date: Sat, 24 Oct 2026 17:43:12 -0700 Subject: [PATCH 08/10] Add Rust cooperative Histogram Expose Block histogram algorithms with typed bin controls and fresh output ownership through the shared native provider. Signed-off-by: Trent Nelson --- docs/rust/histogram.rst | 88 ++ docs/rust/index.rst | 12 +- rust/cuda-coop-codegen/src/histogram.rs | 319 ++++++ rust/cuda-coop-codegen/src/lib.rs | 29 +- rust/cuda-coop-codegen/src/oxide.rs | 135 ++- rust/cuda-coop-codegen/tests/native.rs | 123 +- rust/cuda-coop/build.rs | 14 +- rust/cuda-coop/examples/histogram/.gitignore | 7 + rust/cuda-coop/examples/histogram/Cargo.lock | 1019 +++++++++++++++++ rust/cuda-coop/examples/histogram/Cargo.toml | 19 + .../examples/histogram/helper/Cargo.toml | 11 + .../examples/histogram/helper/src/lib.rs | 36 + rust/cuda-coop/examples/histogram/src/main.rs | 95 ++ rust/cuda-coop/src/histogram.rs | 238 ++++ rust/cuda-coop/src/histogram/boundary.rs | 145 +++ rust/cuda-coop/src/lib.rs | 8 + rust/cuda-coop/tests/check_load_store.py | 90 ++ rust/cuda-coop/tests/histogram/.gitignore | 7 + rust/cuda-coop/tests/histogram/Cargo.lock | 1011 ++++++++++++++++ rust/cuda-coop/tests/histogram/Cargo.toml | 17 + rust/cuda-coop/tests/histogram/src/main.rs | 736 ++++++++++++ 21 files changed, 4135 insertions(+), 24 deletions(-) create mode 100644 docs/rust/histogram.rst create mode 100644 rust/cuda-coop-codegen/src/histogram.rs create mode 100644 rust/cuda-coop/examples/histogram/.gitignore create mode 100644 rust/cuda-coop/examples/histogram/Cargo.lock create mode 100644 rust/cuda-coop/examples/histogram/Cargo.toml create mode 100644 rust/cuda-coop/examples/histogram/helper/Cargo.toml create mode 100644 rust/cuda-coop/examples/histogram/helper/src/lib.rs create mode 100644 rust/cuda-coop/examples/histogram/src/main.rs create mode 100644 rust/cuda-coop/src/histogram.rs create mode 100644 rust/cuda-coop/src/histogram/boundary.rs create mode 100644 rust/cuda-coop/tests/histogram/.gitignore create mode 100644 rust/cuda-coop/tests/histogram/Cargo.lock create mode 100644 rust/cuda-coop/tests/histogram/Cargo.toml create mode 100644 rust/cuda-coop/tests/histogram/src/main.rs diff --git a/docs/rust/histogram.rst b/docs/rust/histogram.rst new file mode 100644 index 000000000000..a639e4b7e4cb --- /dev/null +++ b/docs/rust/histogram.rst @@ -0,0 +1,88 @@ +.. Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. +.. SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +Counting samples into bins +========================== + +``histogram`` counts a complete Block's samples and returns fresh counters. +The bin count is required; defaults select ``i32`` counters, one counter per +thread and the ``Atomic`` algorithm: + +.. code-block:: rust + + let counts = coop::histogram(block, &samples, + coop::HistogramOptions::<32>::default()); + +Every sample is an integer in ``[0, BINS)``. All input slots participate, +including any fill supplied to Load. Zero-filled tails count toward bin zero. +Unreadable or out-of-range samples trap before native counting. Each call +starts its counters from zero and preserves the input. + +A complete program +------------------ + +This helper loads a tile with an inferred sample type and a generic +``ITEMS_PER_THREAD``. It counts 65 bins using ``Sort`` and returns two +``u64`` counters per thread, together with the preserved input: + +.. literalinclude:: ../../rust/cuda-coop/examples/histogram/helper/src/lib.rs + :language: rust + :start-after: // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +The host uses 64-thread Blocks and instantiates one and four input items per +thread. Each Block owns 128 output slots: 65 bins and 63 initialized padding +zeros. It checks every counter and the preserved input: + +.. literalinclude:: ../../rust/cuda-coop/examples/histogram/src/main.rs + :language: rust + :start-after: // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +With the pinned Oxide revision and Rust ``1.99.1`` +installed, run from a complete CCCL checkout: + +.. code-block:: shell + + cd rust/cuda-coop/examples/histogram + 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. ``auto_shared`` uses the compiler's +plan for the complete temporary-storage requirement. + +Independent input and output shapes +----------------------------------- + +Samples support ``u8``, ``i32``, ``u32``, ``i64`` and ``u64``. Counters support +``i32``, ``u32``, ``i64`` and ``u64``. ``counter_type::()`` selects their +type and ``bins_per_thread::()`` selects their per-thread extent, independently +of the input's type, item count and alignment. The result is +``ThreadData`` with its default alignment, including C's natural alignment. + +For a Block of size B, rank t owns bins ``t + i * B`` for slots i in ``0..K``. +Every slot is readable; slots beyond ``BINS`` contain zero. Use ``Striped`` +Store and an output ownership grain of K to write increasing bin order. +The complete output capacity ``B * K`` must cover ``BINS``. The bin count, +input tile size and output capacity must fit signed 32-bit indexing. + +``Atomic`` counts with atomic additions. ``Sort`` counts by sorting a private +copy; it also preserves the borrowed input. Both support complete +one-dimensional Blocks with positive input and output extents. Every thread +participates with the same static selections. Generic helpers need no +``generic_const_exprs`` feature. + +Temporary storage +----------------- + +Scratch includes CUB's temporary storage and the intermediate counters for +both algorithms. The operation synchronizes before reading these counters. +Omitted storage uses managed synchronization before later scratch reuse. + +Use ``temp_storage(&mut scratch)`` to borrow an explicit descriptor. Manual +reuse requires ``cuda_device::thread::sync_threads()`` after the call; +``auto_sync()`` supplies a trailing barrier. Exclusive storage retains its +existing allocation policy. Unrelated global accesses that require Block +completion still need their own synchronization. + +See :doc:`payload`, :doc:`groups` and :doc:`load_store` for checked access, +participation and storage. Generate the API reference with +``cargo doc --workspace --no-deps`` from ``rust/`` for the full signatures. diff --git a/docs/rust/index.rst b/docs/rust/index.rst index 9f10b3ff44ef..74b9e45c50d4 100644 --- a/docs/rust/index.rst +++ b/docs/rust/index.rst @@ -45,16 +45,19 @@ group's tile, and must be uniform across that group. The public operations include :doc:`load_store`, :doc:`scan`, :doc:`reduce`, :doc:`reduce_batched`, :doc:`exchange`, :doc:`shuffle`, -:doc:`adjacent_difference` and :doc:`discontinuity`. +:doc:`adjacent_difference`, :doc:`discontinuity` and :doc:`histogram`. Scan returns fresh per-group prefixes; Reduce returns a total only at each group's root. Batched reduction returns a fresh payload whose extent and valid slots follow its Warp width and output layout. Exchange converts between blocked and striped per-thread layouts. Shuffle shifts blocked values by one item, leaving the exposed edge invalid. Discontinuity returns head flags, tail flags or both as fresh ``i32`` payloads. -Scan, Reduce, batched reduction, Exchange, Shuffle, Adjacent Difference and -Discontinuity preserve their input. Additional primitive families are under -development. The crate is unpublished and its interface can change. +Histogram returns fresh striped counters with an independent type and extent. +It requires generic native-registration support beyond the currently pinned +Oxide revision; see :doc:`histogram` for the compiler requirement. +Scan, Reduce, batched reduction, Exchange, Shuffle, Adjacent Difference, +Discontinuity and Histogram preserve their input. Additional primitive families +are under development. The crate is unpublished and its interface can change. Build and run ------------- @@ -128,6 +131,7 @@ Read next shuffle adjacent_difference discontinuity + histogram provider Generate the API reference with ``cargo doc --workspace --no-deps`` from diff --git a/rust/cuda-coop-codegen/src/histogram.rs b/rust/cuda-coop-codegen/src/histogram.rs new file mode 100644 index 000000000000..3a1d0d1636e1 --- /dev/null +++ b/rust/cuda-coop-codegen/src/histogram.rs @@ -0,0 +1,319 @@ +// Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +//! Validated Block histograms with independent counter types and output extents. + +use crate::{Dtype, Error, Result}; + +/// The CUB implementation used to construct a fresh Block histogram. +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub enum HistogramAlgorithm { + /// Accumulate samples using atomics, the common API default. + Atomic, + /// Sort a private copy of the samples before counting runs. + Sort, +} + +impl HistogramAlgorithm { + pub(crate) fn token(self) -> &'static str { + match self { + Self::Atomic => "atomic", + Self::Sort => "sort", + } + } + + fn native(self) -> &'static str { + match self { + Self::Atomic => "cub::BLOCK_HISTO_ATOMIC", + Self::Sort => "cub::BLOCK_HISTO_SORT", + } + } +} + +pub(crate) const TEMPLATE_OPERATION: &str = + "histogram_xxx_bins0000000000_k0000000000_validity_bits_v1"; + +/// Static histogram choices independent of the input type and extent. +/// +/// `counter` is one of i32/u32/i64/u64; `bins` is the block-wide bin count; +/// `output_items` is each lane's output capacity. A request additionally checks +/// that the complete Block has capacity for every bin. Every fresh count fits +/// all four counter types because the input tile is bounded by i32::MAX. +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub struct HistogramRecipe { + algorithm: HistogramAlgorithm, + pub(crate) counter: Dtype, + bins: u32, + output_items: u32, +} + +impl HistogramRecipe { + /// Select the algorithm, counter type, total bins and per-lane output extent. + /// + /// # Errors + /// Rejects counters other than i32/u32/i64/u64 and zero or greater-than-i32::MAX + /// bin counts or output extents. Block capacity is checked by + /// [`HistogramRequest::block`]. + pub fn new( + algorithm: HistogramAlgorithm, + counter: Dtype, + bins: u32, + output_items: u32, + ) -> Result { + if !matches!(counter, Dtype::I32 | Dtype::U32 | Dtype::I64 | Dtype::U64) { + return Err(Error::message("Histogram counters must be i32/u32/i64/u64")); + } + if bins == 0 + || bins > i32::MAX as u32 + || output_items == 0 + || output_items > i32::MAX as u32 + { + return Err(Error::message( + "Histogram bins and output extent must be positive and no larger than i32::MAX", + )); + } + Ok(Self { + algorithm, + counter, + bins, + output_items, + }) + } + + pub(crate) fn token(self) -> String { + let counter = match self.counter { + Dtype::I32 => "i32", + Dtype::U32 => "u32", + Dtype::I64 => "i64", + Dtype::U64 => "u64", + _ => unreachable!("validated histogram counter"), + }; + format!( + "histogram_{counter}_bins{:010}_k{:010}_validity_bits_v1", + self.bins, self.output_items + ) + } + + pub(crate) fn parse(operation: &str, algorithm: &str) -> Result { + let algorithm = match algorithm { + "atomic" => HistogramAlgorithm::Atomic, + "sort" => HistogramAlgorithm::Sort, + _ => return Err(Error::message("Histogram algorithm must be atomic or sort")), + }; + let invalid = || Error::message("unsupported Histogram operation"); + let body = operation + .strip_prefix("histogram_") + .and_then(|s| s.strip_suffix("_validity_bits_v1")) + .ok_or_else(invalid)?; + let (counter, extents) = body.split_once("_bins").ok_or_else(invalid)?; + let counter = match counter { + "i32" => Dtype::I32, + "u32" => Dtype::U32, + "i64" => Dtype::I64, + "u64" => Dtype::U64, + _ => return Err(invalid()), + }; + let (bins, output_items) = extents.split_once("_k").ok_or_else(invalid)?; + let decimal = |text: &str| { + if text.len() != 10 || !text.bytes().all(|b| b.is_ascii_digit()) { + return Err(invalid()); + } + text.parse::().map_err(|_| invalid()) + }; + Self::new(algorithm, counter, decimal(bins)?, decimal(output_items)?) + } +} + +/// A complete one-dimensional Block histogram with preserved input. +/// +/// Input samples are u8/i32/u32/i64/u64, with `items` slots per lane. The native +/// ABI is `(const T* input, C* output, u8* scratch, const u8* input_bits, +/// u8* output_bits)`, with no scalar controls. Input and output bitmaps contain +/// respectively ceil(items/8) and ceil(output_items/8) low-bit-first bytes. +/// Every input bit is checked before any numeric read. Samples outside +/// `[0, bins)` trap before CUB; signed negatives and wide integers are checked +/// without narrowing. All buffers must be disjoint, live and naturally aligned. +/// +/// CUB operates on a private sample copy. Each rank t receives bins t+i*BlockX; +/// padding slots are zero. Every output slot is initialized and marked valid, +/// with unused high bitmap bits cleared. Input values and validity are preserved. +/// Scratch includes both CUB TempStorage and unsigned `counters[bins]`, using the +/// public counter width. The returned layout includes their padding/alignment. +/// +/// Every Block member participates in converged control flow. The provider +/// synchronizes after CUB counting and before reading the counters. Callers +/// must still complete the Block collective before scratch reuse: output +/// distribution has no trailing completion barrier. +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub struct HistogramRequest { + pub(crate) recipe: HistogramRecipe, + pub(crate) dtype: Dtype, + block: [u32; 3], + items: u32, +} + +impl HistogramRequest { + /// Validate input samples, complete 1D Block geometry and output capacity. + /// + /// `dtype` selects u8/i32/u32/i64/u64. `items` is the positive input extent; + /// the independent counter type and output extent come from `recipe`. + /// + /// # Errors + /// Rejects unsupported samples, non-1D or nonpositive Blocks, more than 1024 + /// threads, a zero input extent, input/output tiles exceeding i32::MAX, or + /// insufficient total output capacity for the selected bins. + pub fn block( + dtype: Dtype, + recipe: HistogramRecipe, + block: [u32; 3], + items: u32, + ) -> Result { + if !matches!( + dtype, + Dtype::U8 | Dtype::I32 | Dtype::U32 | Dtype::I64 | Dtype::U64 + ) { + return Err(Error::message( + "Histogram samples must be u8/i32/u32/i64/u64", + )); + } + let [threads, y, z] = block; + if !(1..=1024).contains(&threads) || y != 1 || z != 1 { + return Err(Error::message( + "Histogram requires a 1D Block with 1..=1024 threads", + )); + } + let bounded_tile = |extent: u32| { + threads + .checked_mul(extent) + .filter(|n| (1..=i32::MAX as u32).contains(n)) + }; + if bounded_tile(items).is_none() { + return Err(Error::message( + "Histogram input tile must be positive and no larger than i32::MAX", + )); + } + let capacity = bounded_tile(recipe.output_items).ok_or_else(|| { + Error::message("Histogram output tile must be positive and no larger than i32::MAX") + })?; + if capacity < recipe.bins { + return Err(Error::message( + "Histogram output capacity is smaller than its bin count", + )); + } + Ok(Self { + recipe, + dtype, + block, + items, + }) + } + + pub(crate) fn block_shape(self) -> [u32; 3] { + self.block + } + + pub(crate) fn source(self, symbol: &str) -> String { + let (sample, _) = self.dtype.native(); + let (counter, _) = self.recipe.counter.native(); + let unsigned_counter = if matches!(self.recipe.counter, Dtype::I64 | Dtype::U64) { + "unsigned long long" + } else { + "unsigned int" + }; + let negative = if matches!(self.dtype, Dtype::I32 | Dtype::I64) { + "if (value < 0) __trap();" + } else { + "" + }; + let threads = self.block[0]; + let items = self.items; + let HistogramRecipe { + algorithm, + bins, + output_items, + .. + } = self.recipe; + let algorithm = algorithm.native(); + let output_bytes = output_items / 8 + u32::from(output_items % 8 != 0); + format!( + r#" +#include +namespace {symbol}_detail {{ +using value_type = {sample}; +using counter_type = {counter}; +using internal_counter_type = {unsigned_counter}; +using primitive = cub::BlockHistogram; +struct scratch_storage {{ + primitive::TempStorage cub_storage; + internal_counter_type counters[{bins}]; +}}; +}} +extern "C" __device__ void {symbol}(const {sample}* input, {counter}* output, + unsigned char* scratch, const unsigned char* input_bits, unsigned char* output_bits) {{ + using namespace {symbol}_detail; +#pragma unroll + for (int i = 0; i < {output_bytes}; ++i) output_bits[i] = 0; +#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) {{ + const value_type value = input[i]; + {negative} + if (static_cast(value) >= {bins}ull) __trap(); + local_input[i] = value; + }} + auto& storage = *reinterpret_cast(scratch); + primitive{{storage.cub_storage}}.Histogram(local_input, storage.counters); + __syncthreads(); +#pragma unroll + for (int i = 0; i < {output_items}; ++i) {{ + const int bin = static_cast(threadIdx.x) + i * {threads}; + output[i] = bin < {bins} ? static_cast(storage.counters[bin]) : counter_type{{0}}; + output_bits[i / 8] |= static_cast(1u << (i % 8)); + }} +}} +"# + ) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn histogram_geometry_types_and_independent_capacity() { + let recipe = HistogramRecipe::new(HistogramAlgorithm::Sort, Dtype::U64, 65, 2).unwrap(); + assert!(HistogramRequest::block(Dtype::U8, recipe, [64, 1, 1], 4).is_ok()); + assert!( + HistogramRequest::block(Dtype::U8, recipe, [32, 1, 1], 4) + .unwrap_err() + .to_string() + .contains("capacity") + ); + for shape in [[0, 1, 1], [1025, 1, 1], [32, 2, 1], [32, 1, 2]] { + assert!(HistogramRequest::block(Dtype::U8, recipe, shape, 4).is_err()); + } + for dtype in [Dtype::I8, Dtype::I16, Dtype::U16, Dtype::F32, Dtype::F64] { + assert!(HistogramRequest::block(dtype, recipe, [64, 1, 1], 4).is_err()); + assert!(HistogramRecipe::new(HistogramAlgorithm::Atomic, dtype, 1, 1).is_err()); + } + for (bins, output) in [(0, 1), (1, 0), (u32::MAX, 1), (1, u32::MAX)] { + assert!( + HistogramRecipe::new(HistogramAlgorithm::Atomic, Dtype::I32, bins, output).is_err() + ); + } + let single = HistogramRecipe::new(HistogramAlgorithm::Atomic, Dtype::I32, 1, 1).unwrap(); + assert!(HistogramRequest::block(Dtype::I64, single, [1, 1, 1], i32::MAX as u32).is_ok()); + for items in [0, i32::MAX as u32, u32::MAX] { + assert!(HistogramRequest::block(Dtype::I64, single, [2, 1, 1], items).is_err()); + } + let large_output = + HistogramRecipe::new(HistogramAlgorithm::Atomic, Dtype::I64, 1, i32::MAX as u32) + .unwrap(); + assert!(HistogramRequest::block(Dtype::U64, large_output, [1, 1, 1], 1).is_ok()); + assert!(HistogramRequest::block(Dtype::U64, large_output, [2, 1, 1], 1).is_err()); + } +} diff --git a/rust/cuda-coop-codegen/src/lib.rs b/rust/cuda-coop-codegen/src/lib.rs index b88e65b594c0..dce4530209fa 100644 --- a/rust/cuda-coop-codegen/src/lib.rs +++ b/rust/cuda-coop-codegen/src/lib.rs @@ -4,7 +4,8 @@ //! Native, host-only CUB specialization for Rust cooperative algorithms. //! //! This private development crate compiles Load/Store, Scan, Reduce, batched -//! Reduce, Exchange, Shuffle, Adjacent Difference and Discontinuity with NVRTC. +//! Reduce, Exchange, Shuffle, Adjacent Difference, Discontinuity and Histogram +//! 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. @@ -53,6 +54,8 @@ //! into a fully initialized fresh output. //! [`DiscontinuityRequest`] compares neighboring Block items, preserving input //! and returning fresh i32 head flags, tail flags or both, with complete validity. +//! [`HistogramRequest`] counts bounded integer samples in a complete 1D Block, +//! with independent counter type/output extent and initialized zero padding. //! [`Provider::compile_native`] batches these families in one native compilation. //! //! [`Provider::from_checkout`] snapshots the selected CCCL and CUDA headers; @@ -86,6 +89,7 @@ mod adjacent_difference; mod discontinuity; mod exchange; mod headers; +mod histogram; mod reduce; mod reduce_batched; mod scan; @@ -93,6 +97,7 @@ mod shuffle; pub use adjacent_difference::{AdjacentDifferenceRecipe, AdjacentDifferenceRequest}; pub use discontinuity::{DiscontinuityRecipe, DiscontinuityRequest}; pub use exchange::{ExchangeMode, ExchangeRecipe, ExchangeRequest}; +pub use histogram::{HistogramAlgorithm, HistogramRecipe, HistogramRequest}; pub use reduce::{ReduceAlgorithm, ReduceForm, ReduceOperator, ReduceRecipe, ReduceRequest}; pub use reduce_batched::{ReduceBatchedLayout, ReduceBatchedRecipe, ReduceBatchedRequest}; pub use scan::{ScanAlgorithm, ScanOperation, ScanOperator, ScanRecipe, ScanRequest}; @@ -491,6 +496,8 @@ pub enum NativeRequest { AdjacentDifference(AdjacentDifferenceRequest), /// Block head/tail flags with preserved input and one or two fresh i32 outputs. Discontinuity(DiscontinuityRequest), + /// Bounded integer samples with independent counter type and output extent. + Histogram(HistogramRequest), } impl NativeRequest { @@ -504,6 +511,7 @@ impl NativeRequest { Self::Shuffle(request) => request.block_shape(), Self::AdjacentDifference(request) => request.block_shape(), Self::Discontinuity(request) => request.block_shape(), + Self::Histogram(request) => request.block_shape(), } } @@ -517,6 +525,7 @@ impl NativeRequest { Self::Shuffle(request) => request.source(symbol), Self::AdjacentDifference(request) => request.source(symbol), Self::Discontinuity(request) => request.source(symbol), + Self::Histogram(request) => request.source(symbol), } } @@ -530,6 +539,7 @@ impl NativeRequest { Self::Shuffle(request) => request.recipe.token().into(), Self::AdjacentDifference(request) => request.recipe.token().into(), Self::Discontinuity(request) => request.recipe.token().into(), + Self::Histogram(request) => request.recipe.token(), } } @@ -542,7 +552,8 @@ impl NativeRequest { | Self::Exchange(_) | Self::Shuffle(_) | Self::AdjacentDifference(_) - | Self::Discontinuity(_) => false, + | Self::Discontinuity(_) + | Self::Histogram(_) => false, } } @@ -558,7 +569,8 @@ impl NativeRequest { | Self::Exchange(_) | Self::Shuffle(_) | Self::AdjacentDifference(_) - | Self::Discontinuity(_) => false, + | Self::Discontinuity(_) + | Self::Histogram(_) => false, } } } @@ -577,11 +589,11 @@ pub struct Layout { /// # Device ABI obligations /// /// The memory-tile and payload-layout requirements below describe Load/Store. -/// Scan, Reduce, batched Reduce, Exchange, Shuffle, Adjacent Difference and -/// Discontinuity use lane-private input and fresh output. Their ABI and validity contracts are -/// documented by [`ScanRequest`], [`ReduceRequest`], [`ReduceBatchedRequest`], -/// [`ExchangeRequest`], [`ShuffleRequest`], [`AdjacentDifferenceRequest`] and -/// [`DiscontinuityRequest`]. +/// Scan, Reduce, batched Reduce, Exchange, Shuffle, Adjacent Difference, +/// Discontinuity and Histogram use lane-private input and fresh output. Their +/// ABI and validity contracts are documented by [`ScanRequest`], [`ReduceRequest`], +/// [`ReduceBatchedRequest`], [`ExchangeRequest`], [`ShuffleRequest`], +/// [`AdjacentDifferenceRequest`], [`DiscontinuityRequest`] and [`HistogramRequest`]. /// /// 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. @@ -805,6 +817,7 @@ impl Provider { include_bytes!("shuffle.rs"), include_bytes!("adjacent_difference.rs"), include_bytes!("discontinuity.rs"), + include_bytes!("histogram.rs"), ] { hash_field(&mut inputs, bytes); } diff --git a/rust/cuda-coop-codegen/src/oxide.rs b/rust/cuda-coop-codegen/src/oxide.rs index 9975c44ea084..c958469fabaa 100644 --- a/rust/cuda-coop-codegen/src/oxide.rs +++ b/rust/cuda-coop-codegen/src/oxide.rs @@ -20,7 +20,8 @@ //! [`AdjacentDifferenceRecipe`] selects Block neighbor subtraction with typed //! optional count and predecessor/successor controls. //! [`DiscontinuityRecipe`] selects Block head/tail flags with one or two i32 -//! outputs and optional neighbors of the input type. +//! outputs and optional neighbors of the input type. [`HistogramRecipe`] selects +//! bounded integer samples with an independent counter type and output extent. //! 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 @@ -29,10 +30,10 @@ use crate::{ AdjacentDifferenceRecipe, AdjacentDifferenceRequest, Algorithm, Artifact, DiscontinuityRecipe, - DiscontinuityRequest, Dtype, Error, ExchangeRecipe, ExchangeRequest, Layout, NativeRequest, - Operation, Provider, ReduceAlgorithm, ReduceBatchedRecipe, ReduceBatchedRequest, ReduceForm, - ReduceOperator, ReduceRecipe, ReduceRequest, Request, Result, ScanRecipe, ScanRequest, - ShuffleRecipe, ShuffleRequest, + DiscontinuityRequest, Dtype, Error, ExchangeRecipe, ExchangeRequest, HistogramAlgorithm, + HistogramRecipe, HistogramRequest, Layout, NativeRequest, Operation, Provider, ReduceAlgorithm, + ReduceBatchedRecipe, ReduceBatchedRequest, ReduceForm, ReduceOperator, ReduceRecipe, + ReduceRequest, Request, Result, ScanRecipe, ScanRequest, ShuffleRecipe, ShuffleRequest, }; use serde::{Deserialize, Serialize}; use std::io::{Read, Write}; @@ -199,6 +200,19 @@ struct SpecializationRequest { impl SpecializationRequest { fn native(&self) -> Result { + if self.operation.starts_with("histogram_") { + let recipe = HistogramRecipe::parse(&self.operation, &self.algorithm)?; + if self.warp_width.is_some() { + return Err(Error::message("Histogram requires a Block scope")); + } + return HistogramRequest::block( + self.dtype.dtype()?, + recipe, + self.block, + self.items_per_thread, + ) + .map(NativeRequest::Histogram); + } if self.operation.starts_with("discontinuity_") { let recipe = DiscontinuityRecipe::parse(&self.operation, &self.algorithm)?; if self.warp_width.is_some() { @@ -391,6 +405,14 @@ struct BatchResponse { fn arguments(request: NativeRequest) -> Vec { let request = match request { + NativeRequest::Histogram(request) => { + return vec![ + Argument::ConstPointer(Scalar::from(request.dtype)), + Argument::MutPointer(Scalar::from(request.recipe.counter)), + Argument::ConstPointer(Scalar::Unsigned(8)), + Argument::MutPointer(Scalar::Unsigned(8)), + ]; + } NativeRequest::Discontinuity(request) => { let dtype = Scalar::from(request.dtype); let mut arguments = vec![ @@ -800,6 +822,38 @@ pub fn discontinuity_registration( .map_err(|error| Error::message(format!("cannot encode provider registration: {error}"))) } +/// Encode a tracked Histogram registration template for one algorithm. +/// +/// The operation contains `xxx` for a sealed i32/u32/i64/u64 counter tag and +/// ten zero digits each for the total bin count and per-lane output extent. +/// A device getter fills those fixed-width fields in its associated constant, +/// retaining the [`serve_inputs`] dependency query from `registration!`. +/// The unfilled template is not a valid specialization. +/// +/// Input type/extent and exact Block shape arrive through the genuine marker. +/// Output type/extent are independent. Native arguments are input, output, +/// scratch, input bitmap, output bitmap; there are no scalar controls. +/// See [`HistogramRequest`] for preflight, output validity and scratch reuse. +/// +/// # Errors +/// Rejects a nonabsolute provider executable or an unencodable path. +pub fn histogram_registration_template( + executable: &Path, + algorithm: HistogramAlgorithm, +) -> 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: crate::histogram::TEMPLATE_OPERATION, + algorithm: algorithm.token(), + }) + .map_err(|error| Error::message(format!("cannot encode provider registration: {error}"))) +} + #[derive(Deserialize)] #[serde(deny_unknown_fields)] struct InputRequest { @@ -1122,6 +1176,77 @@ mod tests { ScanOperator, }; + #[test] + fn histogram_template_decoding_and_independent_counter_abi() { + for (algorithm, counter, bins, output_items) in [ + (HistogramAlgorithm::Sort, Dtype::U64, 65, 2), + (HistogramAlgorithm::Atomic, Dtype::I32, 1, 9), + (HistogramAlgorithm::Sort, Dtype::I64, 33, 1), + (HistogramAlgorithm::Atomic, Dtype::U32, 65, 3), + ] { + let template = + histogram_registration_template(Path::new("/provider"), algorithm).unwrap(); + let mut registration: serde_json::Value = serde_json::from_str(&template).unwrap(); + assert_eq!( + registration["operation"], + crate::histogram::TEMPLATE_OPERATION + ); + let recipe = HistogramRecipe::new(algorithm, counter, bins, output_items).unwrap(); + registration["operation"] = recipe.token().into(); + let request = SpecializationRequest { + operation: registration["operation"].as_str().unwrap().into(), + algorithm: registration["algorithm"].as_str().unwrap().into(), + dtype: Scalar::Unsigned(8), + block: [64, 1, 1], + items_per_thread: 4, + warp_width: None, + }; + let native = request.native().unwrap(); + assert_eq!( + native, + NativeRequest::Histogram( + HistogramRequest::block(Dtype::U8, recipe, [64, 1, 1], 4).unwrap() + ) + ); + assert_eq!( + arguments(native), + vec![ + Argument::ConstPointer(Scalar::Unsigned(8)), + Argument::MutPointer(Scalar::from(counter)), + Argument::ConstPointer(Scalar::Unsigned(8)), + Argument::MutPointer(Scalar::Unsigned(8)), + ] + ); + assert!( + SpecializationRequest { + warp_width: Some(32), + ..request + } + .native() + .is_err() + ); + } + for operation in [ + crate::histogram::TEMPLATE_OPERATION, + "histogram_u64_bins65_k0000000002_validity_bits_v1", + "histogram_u64_bins0000000065_k0000000002_validity_bits_v2", + "histogram_f64_bins0000000065_k0000000002_validity_bits_v1", + "histogram_u64_bins0000000000_k0000000002_validity_bits_v1", + "histogram_u64_bins2147483648_k0000000002_validity_bits_v1", + "histogram_u64_bins9999999999_k0000000002_validity_bits_v1", + "histogram_u64_bins0000000065_k+000000002_validity_bits_v1", + "histogram_u64_bins0000000065_k0000000000_validity_bits_v1", + ] { + assert!(HistogramRecipe::parse(operation, "sort").is_err()); + } + let valid = "histogram_u64_bins0000000065_k0000000002_validity_bits_v1"; + assert!(HistogramRecipe::parse(valid, "block").is_err()); + assert!( + histogram_registration_template(Path::new("relative"), HistogramAlgorithm::Atomic) + .is_err() + ); + } + #[test] fn discontinuity_registration_preserves_mixed_output_abi() { use DiscontinuityRecipe::*; diff --git a/rust/cuda-coop-codegen/tests/native.rs b/rust/cuda-coop-codegen/tests/native.rs index 7c2024768da8..3cb7cb988141 100644 --- a/rust/cuda-coop-codegen/tests/native.rs +++ b/rust/cuda-coop-codegen/tests/native.rs @@ -5,13 +5,128 @@ use cuda_coop_codegen::{ AdjacentDifferenceRecipe, AdjacentDifferenceRequest, Algorithm, DiscontinuityRecipe, - DiscontinuityRequest, Dtype, ExchangeMode, ExchangeRecipe, ExchangeRequest, Layout, - NativeRequest, Operation, Provider, ReduceAlgorithm, ReduceBatchedLayout, ReduceBatchedRecipe, - ReduceBatchedRequest, ReduceForm, ReduceOperator, ReduceRequest, Request, ScanAlgorithm, - ScanOperation, ScanOperator, ScanRequest, ShuffleRecipe, ShuffleRequest, + DiscontinuityRequest, Dtype, ExchangeMode, ExchangeRecipe, ExchangeRequest, HistogramAlgorithm, + HistogramRecipe, HistogramRequest, Layout, NativeRequest, Operation, Provider, ReduceAlgorithm, + ReduceBatchedLayout, ReduceBatchedRecipe, ReduceBatchedRequest, ReduceForm, ReduceOperator, + ReduceRequest, Request, ScanAlgorithm, ScanOperation, ScanOperator, ScanRequest, ShuffleRecipe, + ShuffleRequest, }; use std::{env, path::Path, time::Instant}; +#[test] +#[ignore = "requires CUDA_HOME and a Linux CUDA toolkit with NVRTC LTO support"] +fn native_mixed_histogram_recipes() -> Result<(), Box> { + use HistogramAlgorithm::{Atomic, Sort}; + 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)?; + // All sample/counter types, independent N/K, both algorithms and bitmap + // tails. The first row is the required public mixed-output checkpoint. + let rows = [ + (Dtype::U8, Dtype::U64, 64, 4, 65, 2, Sort), + (Dtype::I32, Dtype::I32, 1, 1, 1, 1, Atomic), + (Dtype::U32, Dtype::U32, 33, 9, 33, 1, Sort), + (Dtype::U64, Dtype::I64, 1, 9, 5, 9, Atomic), + (Dtype::I64, Dtype::U64, 33, 4, 65, 3, Atomic), + (Dtype::U8, Dtype::I32, 33, 1, 257, 9, Sort), + (Dtype::I32, Dtype::U32, 64, 4, 32, 1, Atomic), + (Dtype::U8, Dtype::U64, 64, 4, 65, 2, Atomic), + (Dtype::U8, Dtype::U64, 64, 4, 65, 3, Sort), + (Dtype::U64, Dtype::U64, 64, 4, 65, 2, Sort), + ]; + let mut requests = Vec::new(); + for (sample, counter, threads, items, bins, output, algorithm) in rows { + requests.push(NativeRequest::Histogram(HistogramRequest::block( + sample, + HistogramRecipe::new(algorithm, counter, bins, output)?, + [threads, 1, 1], + items, + )?)); + } + requests.extend([ + NativeRequest::LoadStore(Request::new_typed( + Dtype::I32, + Operation::LoadFullValidity, + Algorithm::Direct, + 64, + 4, + )?), + NativeRequest::Discontinuity(DiscontinuityRequest::block( + Dtype::F64, + DiscontinuityRecipe::HeadsAndTailsBoth, + [4, 3, 2], + 4, + )?), + NativeRequest::AdjacentDifference(AdjacentDifferenceRequest::block( + Dtype::I32, + AdjacentDifferenceRecipe::LeftPartialPredecessor, + [64, 1, 1], + 4, + )?), + ]); + let artifact = provider.compile_native(&requests, &target)?; + assert_eq!(artifact.specializations.len(), 13); + assert!(!artifact.ltoir.is_empty()); + let mut identities = std::collections::HashSet::new(); + let mut symbols = std::collections::HashSet::new(); + for (index, (request, entry)) in requests.iter().zip(&artifact.specializations).enumerate() { + assert!(identities.insert(entry.identity)); + assert!(symbols.insert(&entry.symbol)); + assert!(!entry.mutates_payload); + if index < rows.len() { + let (_, counter, _, _, bins, _, _) = rows[index]; + let width = if matches!(counter, Dtype::I64 | Dtype::U64) { + 8 + } else { + 4 + }; + let alignment = entry.cub_layout.alignment.max(width); + let counter_offset = entry.cub_layout.size.next_multiple_of(width); + assert_eq!( + entry.scratch, + Layout { + size: (counter_offset + bins as usize * width).next_multiple_of(alignment), + alignment, + } + ); + assert!(entry.scratch.size > entry.cub_layout.size); + assert!(!entry.scratch_free); + } else 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_free); + } + eprintln!( + "{request:?}: CUB {:?}, scratch {:?}", + entry.cub_layout, entry.scratch + ); + } + // K changes the output contract and identity without changing shared storage. + assert_ne!( + artifact.specializations[0].identity, + artifact.specializations[8].identity + ); + assert_eq!( + artifact.specializations[0].scratch, + artifact.specializations[8].scratch + ); + // The algorithm also participates in identity; Atomic still owns counters. + assert_ne!( + artifact.specializations[0].identity, + artifact.specializations[7].identity + ); + assert!(artifact.specializations[7].scratch.size >= 65 * 8); + eprintln!( + "{} mixed Histogram/Load/Discontinuity/Adjacent Difference 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_discontinuity_recipes() -> Result<(), Box> { diff --git a/rust/cuda-coop/build.rs b/rust/cuda-coop/build.rs index d470b30e1e34..650ef7dfe1f7 100644 --- a/rust/cuda-coop/build.rs +++ b/rust/cuda-coop/build.rs @@ -4,9 +4,9 @@ //! Register the owning host provider for native cooperative operations. use cuda_coop_codegen::{ AdjacentDifferenceRecipe, Algorithm, DiscontinuityRecipe, ExchangeMode, ExchangeRecipe, - Operation, ReduceAlgorithm, ReduceBatchedLayout, ReduceBatchedRecipe, ReduceForm, - ReduceOperator, ReduceRecipe, ScanAlgorithm, ScanOperation, ScanOperator, ScanRecipe, - ShuffleRecipe, oxide, + HistogramAlgorithm, Operation, ReduceAlgorithm, ReduceBatchedLayout, ReduceBatchedRecipe, + ReduceForm, ReduceOperator, ReduceRecipe, ScanAlgorithm, ScanOperation, ScanOperator, + ScanRecipe, ShuffleRecipe, oxide, }; fn main() { @@ -25,6 +25,14 @@ fn main() { let callable = oxide::reduce_callable_registration(&executable) .expect("cuda.coop callable Reduce registration"); println!("cargo::rustc-env=COOP_REDUCE_RAKING_CALLABLE_PAYLOAD={callable}"); + for (algorithm, name) in [ + (HistogramAlgorithm::Atomic, "ATOMIC"), + (HistogramAlgorithm::Sort, "SORT"), + ] { + let value = oxide::histogram_registration_template(&executable, algorithm) + .expect("cuda.coop histogram registration template"); + println!("cargo::rustc-env=COOP_HISTOGRAM_{name}={value}"); + } for (recipe, name) in [ (DiscontinuityRecipe::Heads, "HEADS"), (DiscontinuityRecipe::HeadsPredecessor, "HEADS_PREDECESSOR"), diff --git a/rust/cuda-coop/examples/histogram/.gitignore b/rust/cuda-coop/examples/histogram/.gitignore new file mode 100644 index 000000000000..74bffa8af4a4 --- /dev/null +++ b/rust/cuda-coop/examples/histogram/.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/histogram/Cargo.lock b/rust/cuda-coop/examples/histogram/Cargo.lock new file mode 100644 index 000000000000..a020ced78d4c --- /dev/null +++ b/rust/cuda-coop/examples/histogram/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-histogram" +version = "0.0.0" +dependencies = [ + "coop-histogram-helper", + "cuda-async", + "cuda-coop", + "cuda-core", + "cuda-device", + "cuda-host", +] + +[[package]] +name = "coop-histogram-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/histogram/Cargo.toml b/rust/cuda-coop/examples/histogram/Cargo.toml new file mode 100644 index 000000000000..d1652ae7d819 --- /dev/null +++ b/rust/cuda-coop/examples/histogram/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-histogram" +version = "0.0.0" +edition = "2024" +publish = false + +[workspace] +members = ["helper"] + +[dependencies] +cuda-coop = { path = "../.." } +coop-histogram-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/histogram/helper/Cargo.toml b/rust/cuda-coop/examples/histogram/helper/Cargo.toml new file mode 100644 index 000000000000..53d52b7ee39a --- /dev/null +++ b/rust/cuda-coop/examples/histogram/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-histogram-helper" +version = "0.0.0" +edition = "2024" +publish = false + +[dependencies] +cuda-coop = { path = "../../.." } diff --git a/rust/cuda-coop/examples/histogram/helper/src/lib.rs b/rust/cuda-coop/examples/histogram/helper/src/lib.rs new file mode 100644 index 000000000000..dfd3c8a1bec5 --- /dev/null +++ b/rust/cuda-coop/examples/histogram/helper/src/lib.rs @@ -0,0 +1,36 @@ +// Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +//! Infer sample types while returning independently shaped Histogram counters. +#![no_std] +use cuda_coop as coop; + +/// Load a full tile and count 65 bins into two striped u64 counters per thread. +/// +/// Every sample must be in 0..65. Sort works on a private native copy, so the +/// returned input payload remains readable and unchanged. +pub fn count_tile( + block: coop::ThreadGroup, + input: &[T], + offset: usize, +) -> ( + coop::ThreadData, + coop::ThreadData, +) { + let mut values = coop::ThreadData::<_, ITEMS_PER_THREAD>::new(); + coop::load( + block, + input, + &mut values, + coop::LoadOptions::default().offset(offset), + ); + let counts = coop::histogram( + block, + &values, + coop::HistogramOptions::<65>::default() + .counter_type::() + .bins_per_thread::<2>() + .algorithm(coop::Sort), + ); + (counts, values) +} diff --git a/rust/cuda-coop/examples/histogram/src/main.rs b/rust/cuda-coop/examples/histogram/src/main.rs new file mode 100644 index 000000000000..49ace1edc0b8 --- /dev/null +++ b/rust/cuda-coop/examples/histogram/src/main.rs @@ -0,0 +1,95 @@ +// 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 = (64, 1, 1), dynamic_shared = inferred)] + pub fn count_tile( + input: &[u8], + mut output: coop::DisjointSlice>, + mut preserved: coop::DisjointSlice>, + ) { + let block = coop::this_block(); + let block_index = block.rank_at::() as usize; + let input_offset = block_index * block.count() as usize * ITEMS_PER_THREAD; + let output_offset = block_index * block.count() as usize * 2; + let (counts, values) = + coop_histogram_helper::count_tile::<_, ITEMS_PER_THREAD>(block, input, input_offset); + coop::store( + block, + &mut output, + &counts, + coop::StoreOptions::default() + .algorithm(coop::Striped) + .offset(output_offset), + ); + coop::store( + block, + &mut preserved, + &values, + coop::StoreOptions::default().offset(input_offset), + ); + // Terminal safe reads also cover the initialized zero-padding slot. + assert!(counts.get(0).is_some() && counts.get(1).is_some()); + } +} + +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 = 64 * $items; + let input: Vec = (0..2 * tile) + .map(|i| if i < tile { 64 } else { ((i - tile) % 65) as u8 }) + .collect(); + // Each Block owns all 128 slots: 65 real bins and 63 padding zeros. + let mut expected = vec![0_u64; 2 * 128]; + for (block, samples) in input.chunks_exact(tile).enumerate() { + for &sample in samples { + expected[block * 128 + usize::from(sample)] += 1; + } + } + let source = DeviceBuffer::from_host(&stream, &input)?; + let target = DeviceBuffer::from_host(&stream, &vec![u64::MAX; expected.len()])?; + let preserved_target = DeviceBuffer::from_host(&stream, &vec![255_u8; input.len()])?; + let plan = module.prepare_count_tile::<$items>(LaunchConfig1D::auto_shared(2, 64))?; + let (source, target, preserved_target) = module + .count_tile_async_owned::<$items, _, _, _>( + &plan, source, target, preserved_target, + ) + .sync()?; + let mut output = vec![0_u64; expected.len()]; + let mut preserved = vec![0_u8; input.len()]; + let mut source_after = vec![0_u8; input.len()]; + target.copy_to_host(&stream, &mut output)?; + preserved_target.copy_to_host(&stream, &mut preserved)?; + source.copy_to_host(&stream, &mut source_after)?; + stream.synchronize()?; + assert_eq!(output, expected); + assert_eq!(preserved, input); + assert_eq!(source_after, input); + println!( + "items/thread={}: {} counters (126 padding), {} preserved inputs and {} source values checked", + $items, output.len(), preserved.len(), source_after.len() + ); + }}; + } + run!(1); + run!(4); + Ok(()) +} diff --git a/rust/cuda-coop/src/histogram.rs b/rust/cuda-coop/src/histogram.rs new file mode 100644 index 000000000000..3ce9cd61543e --- /dev/null +++ b/rust/cuda-coop/src/histogram.rs @@ -0,0 +1,238 @@ +// 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, Numeric, SupportedAlignment, TempStorage, TempStorageLike, ThreadData, + ThreadGroup, +}; + +mod boundary; + +/// Count samples with atomic additions (the default Histogram algorithm). +#[derive(Clone, Copy, Debug, Default)] +pub struct Atomic; + +/// Count samples by sorting a private copy, preserving the input payload. +#[derive(Clone, Copy, Debug, Default)] +pub struct Sort; + +/// Sealed Histogram input types: u8, i32, u32, i64 and u64. +/// +/// Every input value must lie in `[0, BINS)`. The native boundary checks this +/// before CUB reads the samples; it does not clamp or filter invalid values. +pub trait HistogramSample: Numeric + sealed::Sample {} + +/// Sealed Histogram counter types: i32, u32, i64 and u64. +/// +/// The counter type is independent of the sample type. A fresh histogram has +/// at most `i32::MAX` samples, so every supported counter represents its count. +pub trait HistogramCounter: Numeric + sealed::Counter {} + +/// Sealed Histogram algorithms: [`Atomic`] and [`Sort`]. +/// +/// Both return the same striped counters and preserve input. They use the +/// complete queried scratch layout, including the intermediate counters. +pub trait HistogramAlgorithm: sealed::Algorithm {} + +/// Static bin count, independent counter shape, algorithm and scratch policy. +/// +/// `BINS` is required. Canonical [`Default`] selects i32 counters, one counter +/// per thread (`K = 1`), [`Atomic`] and implicit scratch. `BINS` and `K` must be +/// positive; the complete Block's `B * K` capacity must cover `BINS`. Bin count, +/// input tile size and output tile capacity must fit signed 32-bit indexing. +/// +/// Builders consume options and preserve other choices. `C` is the counter +/// type, `A` the algorithm and `S` the existing storage selection. Input type, +/// extent and alignment are inferred separately from the borrowed payload. +/// +/// ``` +/// use cuda_coop as coop; +/// let _: coop::HistogramOptions<32> = Default::default(); +/// let _: coop::HistogramOptions<65, u64, 2, coop::Sort> = +/// coop::HistogramOptions::<65>::default() +/// .counter_type::().bins_per_thread::<2>().algorithm(coop::Sort); +/// ``` +/// +/// Counters cannot use narrow integer or floating-point types: +/// ```compile_fail,E0277 +/// use cuda_coop as coop; +/// let _ = coop::HistogramOptions::<32>::default().counter_type::(); +/// ``` +pub struct HistogramOptions { + storage: S, + selection: PhantomData<(C, A)>, +} + +impl Default for HistogramOptions { + fn default() -> Self { + Self { + storage: (), + selection: PhantomData, + } + } +} + +impl HistogramOptions { + /// Select the output counter type independently of the input samples. + pub fn counter_type(self) -> HistogramOptions { + HistogramOptions { + storage: self.storage, + selection: PhantomData, + } + } + + /// Select the number of striped output counters owned by each thread. + /// + /// All `K2` slots are initialized, including zero padding beyond `BINS`. + /// The provider rejects a nonpositive or insufficient output capacity. + pub fn bins_per_thread(self) -> HistogramOptions { + HistogramOptions { + storage: self.storage, + selection: PhantomData, + } + } + + /// Select [`Atomic`] or [`Sort`], retaining counter shape and storage. + pub fn algorithm( + self, + _algorithm: A2, + ) -> HistogramOptions { + HistogramOptions { + storage: self.storage, + selection: PhantomData, + } + } + + /// Borrow explicit scratch with its existing sharing and completion policy. + /// + /// Manual reuse requires `cuda_device::thread::sync_threads()` after the + /// call. Auto-sync storage supplies its existing trailing barrier. Options + /// are consumed, allowing ordinary mutable reborrows on successive calls. + pub fn temp_storage< + const CAPACITY: usize, + const ALIGN: usize, + const AUTO: bool, + const EX: bool, + >( + self, + storage: &mut TempStorage, + ) -> HistogramOptions> { + HistogramOptions { + storage, + selection: PhantomData, + } + } +} + +/// Count a complete one-dimensional Block's samples into fresh striped counters. +/// +/// Every member of `group` participates with `ITEMS_PER_THREAD` initialized +/// samples. Supported samples are [`HistogramSample`]'s five integer types; +/// each value must be in `[0, BINS)`. Unreadable or out-of-range samples trap +/// before native Histogram accesses. There is no partial count or filtering: +/// Load's zero-filled tail samples contribute to bin zero. +/// +/// `BINS` is the number of real bins, `C` the independent counter type, and `K` +/// the per-thread output extent. Rank `t`, slot `i` owns bin `t + i * B`, where +/// `B` is the Block size. Every slot is valid; bins beyond `BINS` contain zero. +/// The output has default alignment, including C's natural alignment. Input +/// values and validity are unchanged, including with [`Sort`]. Shared input +/// and mutable reborrows are accepted. No generic-const-expression feature is +/// needed in an ordinary consumer, including generic T/N/ALIGN/C/K helpers. +/// +/// Store results with [`crate::Striped`] and an output ownership grain of K +/// to obtain increasing bin order. Input T/N and output C/K are independent. +/// Each call starts from zero. Scratch includes both CUB's temporary storage +/// and intermediate counters, even for Atomic. The internal readiness barrier +/// precedes result reads; later scratch reuse follows the existing implicit or +/// explicit [`TempStorage`] policy. Unrelated global accesses gain no completion +/// guarantee from this operation. +/// +/// Unsupported types/groups and invalid static sizes fail compilation. Only +/// complete 1D Blocks are supported, with `BINS`, `B * ITEMS_PER_THREAD` and +/// `B * K` at most `i32::MAX`, and `B * K >= BINS`. Required BINS is specified +/// in options; it cannot be inferred from input. Execution requires an +/// Oxide-compiled CUDA kernel with support for generic registration constants. +/// +/// ```no_run +/// use cuda_coop as coop; +/// fn counts( +/// block: coop::ThreadGroup, +/// input: &coop::ThreadData, +/// ) -> coop::ThreadData +/// where coop::Alignment: coop::SupportedAlignment { +/// coop::histogram(block, input, coop::HistogramOptions::::default() +/// .counter_type::().bins_per_thread::()) +/// } +/// ``` +/// +/// Floating-point samples are unsupported: +/// ```compile_fail,E0277 +/// use cuda_coop as coop; +/// fn invalid(input: &coop::ThreadData) { +/// coop::histogram(coop::this_block(), input, coop::HistogramOptions::<32>::default()); +/// } +/// ``` +#[inline(always)] +pub fn histogram< + T, + C, + A, + S, + const BINS: usize, + const ITEMS_PER_THREAD: usize, + const K: usize, + const ALIGN: usize, +>( + _group: ThreadGroup, + input: &ThreadData, + options: HistogramOptions, +) -> ThreadData +where + T: HistogramSample, + C: HistogramCounter, + A: HistogramAlgorithm, + S: TempStorageLike, + Alignment: SupportedAlignment, +{ + boundary::histogram(input, options) +} + +mod sealed { + pub trait Sample {} + pub trait Counter { + const TAG: [u8; 3]; + } + pub trait Algorithm { + const SORT: bool; + } +} + +macro_rules! samples { + ($($ty:ty),+ $(,)?) => {$ ( + impl sealed::Sample for $ty {} + impl HistogramSample for $ty {} + )+}; +} +samples!(u8, i32, u32, i64, u64); + +macro_rules! counters { + ($(($ty:ty, $tag:literal)),+ $(,)?) => {$ ( + impl sealed::Counter for $ty { const TAG: [u8; 3] = *$tag; } + impl HistogramCounter for $ty {} + )+}; +} +counters!((i32, b"i32"), (u32, b"u32"), (i64, b"i64"), (u64, b"u64")); + +impl sealed::Algorithm for Atomic { + const SORT: bool = false; +} +impl HistogramAlgorithm for Atomic {} +impl sealed::Algorithm for Sort { + const SORT: bool = true; +} +impl HistogramAlgorithm for Sort {} diff --git a/rust/cuda-coop/src/histogram/boundary.rs b/rust/cuda-coop/src/histogram/boundary.rs new file mode 100644 index 000000000000..337ab717a225 --- /dev/null +++ b/rust/cuda-coop/src/histogram/boundary.rs @@ -0,0 +1,145 @@ +// Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +use super::{HistogramAlgorithm, HistogramCounter, HistogramOptions, HistogramSample}; +use crate::{Alignment, SupportedAlignment, TempStorageLike, ThreadData}; +use cuda_device::library; + +const TOKEN: &[u8] = b"\"histogram_xxx_bins0000000000_k0000000000_validity_bits_v1\""; + +const fn decimal(bytes: &mut [u8], offset: usize, mut value: usize) { + let mut remaining = 10; + while remaining > 0 { + remaining -= 1; + bytes[offset + remaining] = b'0' + (value % 10) as u8; + value /= 10; + } + assert!(value == 0, "Histogram decimal field overflow"); +} + +// Preserve registration!'s tracked JSON, changing only the closed operation +// token. Each algorithm's concrete template length is independent of BINS/K. +const fn encode(base: &str, tag: [u8; 3], bins: usize, k: usize) -> [u8; LEN] { + assert!( + bins > 0 && bins <= i32::MAX as usize, + "Histogram requires positive bins <= i32::MAX" + ); + assert!( + k > 0 && k <= i32::MAX as usize, + "Histogram requires positive output extent <= i32::MAX" + ); + let source = base.as_bytes(); + assert!(source.len() == LEN, "Histogram template length changed"); + let mut bytes = [0_u8; LEN]; + let mut i = 0; + while i < LEN { + bytes[i] = source[i]; + i += 1; + } + let mut found = 0; + let mut position = 0; + let mut start = 0; + while start + TOKEN.len() <= LEN { + let mut j = 0; + while j < TOKEN.len() && source[start + j] == TOKEN[j] { + j += 1; + } + if j == TOKEN.len() { + found += 1; + position = start; + } + start += 1; + } + assert!( + found == 1, + "expected exactly one Histogram operation template" + ); + let counter = position + b"\"histogram_".len(); + let bins_offset = counter + 3 + b"_bins".len(); + let k_offset = bins_offset + 10 + b"_k".len(); + i = 0; + while i < 3 { + bytes[counter + i] = tag[i]; + i += 1; + } + decimal(&mut bytes, bins_offset, bins); + decimal(&mut bytes, k_offset, k); + bytes +} + +macro_rules! registration { + ($module:ident, $env:literal) => { + mod $module { + use super::*; + const BASE: &str = library::registration!($env); + const LEN: usize = BASE.len(); + struct Text(core::marker::PhantomData); + impl Text { + const BYTES: [u8; LEN] = encode(BASE, C::TAG, BINS, K); + const REGISTRATION: &'static str = match core::str::from_utf8(&Self::BYTES) { + Ok(text) => text, + Err(_) => panic!("invalid Histogram registration UTF-8"), + }; + } + pub(super) fn registration() + -> &'static str { + Text::::REGISTRATION + } + } + }; +} +registration!(atomic, "COOP_HISTOGRAM_ATOMIC"); +registration!(sort, "COOP_HISTOGRAM_SORT"); + +#[inline(always)] +pub(super) fn histogram< + T, + C, + A, + S, + const BINS: usize, + const N: usize, + const K: usize, + const ALIGN: usize, +>( + input: &ThreadData, + options: HistogramOptions, +) -> ThreadData +where + T: HistogramSample, + C: HistogramCounter, + A: HistogramAlgorithm, + S: TempStorageLike, + Alignment: SupportedAlignment, +{ + let mut output = ThreadData::::new(); + let (input_pointer, input_bits) = input.block_io_read_parts(); + let (output_pointer, output_bits) = output.block_io_parts(); + let payload = (input_pointer, output_pointer, input_bits, output_bits); + let storage = options.storage.into_raw(); + // SAFETY: input and fresh output are distinct live naturally aligned local + // buffers with separate N-bit and K-bit validity maps. The exact recipe + // validates all input bits and sample bounds before CUB, copies input, + // initializes all C/K counters (including padding), then marks output bits. + // No pointer escapes the synchronous call. Sealed static selections encode + // the exact output type/extent. The genuine marker checks full Block + // participation and scratch origin, queried layout and reuse policy. + unsafe { + if A::SORT { + library::checked_block_payload_call_v1::( + sort::registration::, + storage, + (), + payload, + ); + } else { + library::checked_block_payload_call_v1::( + atomic::registration::, + storage, + (), + payload, + ); + } + } + output +} diff --git a/rust/cuda-coop/src/lib.rs b/rust/cuda-coop/src/lib.rs index e55f911b6b76..69ee57255a18 100644 --- a/rust/cuda-coop/src/lib.rs +++ b/rust/cuda-coop/src/lib.rs @@ -65,6 +65,9 @@ //! selects direction, typed outside neighbors, count and existing scratch policy. //! [`discontinuity`] returns fresh i32 head or tail flags, or both payloads, //! preserving the full-readable input's extent and minimum alignment. +//! [`histogram`] counts a complete 1D Block's samples into fresh striped counters. +//! [`HistogramOptions`] selects the required bin count, independent counter type +//! and extent, Atomic or Sort algorithm, and existing scratch policy. //! Host payload tests need no GPU; //! executing hardware queries requires an Oxide-compiled CUDA kernel. @@ -77,6 +80,7 @@ mod alignment; mod discontinuity; mod exchange; mod group; +mod histogram; mod load_store; mod reduce; mod reduce_batched; @@ -96,6 +100,10 @@ pub use discontinuity::{ pub use exchange::{ BlockedToStriped, ExchangeGroup, ExchangeMode, ExchangeOptions, StripedToBlocked, exchange, }; +pub use histogram::{ + Atomic, HistogramAlgorithm, HistogramCounter, HistogramOptions, HistogramSample, Sort, + histogram, +}; pub use load_store::{ Algorithm, Direct, LoadOptions, LoadSource, LoadStoreGroup, StoreOptions, StoreValue, Striped, Transpose, Vectorize, WarpTranspose, WarpTransposeTimesliced, load, store, diff --git a/rust/cuda-coop/tests/check_load_store.py b/rust/cuda-coop/tests/check_load_store.py index 69b7fa259711..d85775cfc065 100644 --- a/rust/cuda-coop/tests/check_load_store.py +++ b/rust/cuda-coop/tests/check_load_store.py @@ -1556,6 +1556,92 @@ ), ) +# Histogram has independent input T/N/ALIGN and counter C/K, with default Atomic. +# Internal counter readiness does not complete unrelated global-memory accesses. +histogram = "let result=coop::histogram(block,&values,coop::HistogramOptions::<129>::default().counter_type::().bins_per_thread::<2>());" +histogram_manual = histogram.replace( + ".bins_per_thread::<2>()", ".bins_per_thread::<2>().temp_storage(&mut scratch)" +) +cases.update( + histogram_generic_defaults=( + load + "fn forward(" + "block:coop::ThreadGroup,input:&coop::ThreadData," + "options:coop::HistogramOptions)->coop::ThreadData " + "where coop::Alignment:coop::SupportedAlignment {" + "coop::histogram(block,input,options)}" + "let mut aligned=coop::ThreadData::<_,4,16>::new();" + "aligned.write(0,values[0]);aligned.write(1,values[1]);" + "aligned.write(2,values[2]);aligned.write(3,values[3]);" + "let result:coop::ThreadData=forward(block,&mut aligned," + "coop::HistogramOptions::<129>::default().counter_type::().bins_per_thread::<2>());" + "values.write(0,result[0] as i32);" + store, + None, + ), + histogram_storage_completed=( + load + new_storage + histogram_manual + sync + histogram_manual + store, + None, + ), + histogram_storage_auto=( + load + + new_storage.replace("new()", "new().auto_sync()") + + histogram_manual + + histogram_manual + + store, + None, + ), + histogram_read_completed=( + reload.replace(".valid_items(valid)", ".valid_items(valid).oob_default(0)") + + histogram + + sync + + store, + None, + ), + histogram_insufficient_capacity=( + load + + "let result=coop::histogram(block,&values,coop::HistogramOptions::<129>::default());", + "Histogram output capacity is smaller than its bin count", + ), + histogram_multidimensional=( + load + histogram, + "Histogram requires a 1D Block with 1..=1024 threads", + ), + histogram_warp=( + load + histogram.replace("histogram(block,", "histogram(coop::this_warp(),"), + "mismatched types", + ), + histogram_zero_bins=( + load + histogram.replace("::<129>", "::<0>"), + "Histogram requires positive bins <= i32::MAX", + ), + histogram_zero_output=( + load + histogram.replace("bins_per_thread::<2>", "bins_per_thread::<0>"), + "ThreadData requires at least one item", + ), + histogram_bins_encoder_bound=( + load + histogram.replace("::<129>", "::<2147483648>"), + "Histogram requires positive bins <= i32::MAX", + ), + histogram_lane_divergence=( + load + "if thread::threadIdx_x()&1==0 {" + histogram + "}", + "lane-varying", + ), + histogram_result_control=( + load + histogram + "if result[0]==0 {" + histogram + "}", + "lane-varying", + ), + histogram_storage_missing_barrier=( + load + new_storage + histogram_manual + histogram_manual, + "requires a full-block barrier", + ), + histogram_read_epoch=( + reload.replace(".valid_items(valid)", ".valid_items(valid).oob_default(0)") + + histogram + + store, + "requires a barrier covering both accesses", + ), +) + if args.case: unknown = set(args.case).difference(cases) if unknown: @@ -1607,6 +1693,10 @@ case_prefix = case_prefix.replace( "dynamic_shared=inferred", "dynamic_shared=512" ) + if name == "histogram_multidimensional": + case_prefix = case_prefix.replace( + "domain=1,block=(128,1,1)", "domain=3,block=(4,4,8)" + ) if name == "warp_partial_geometry": case_prefix = case_prefix.replace("block=(64,1,1)", "block=(48,1,1)") elif name.endswith("_partial_block"): diff --git a/rust/cuda-coop/tests/histogram/.gitignore b/rust/cuda-coop/tests/histogram/.gitignore new file mode 100644 index 000000000000..74bffa8af4a4 --- /dev/null +++ b/rust/cuda-coop/tests/histogram/.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/histogram/Cargo.lock b/rust/cuda-coop/tests/histogram/Cargo.lock new file mode 100644 index 000000000000..f97bda7593e4 --- /dev/null +++ b/rust/cuda-coop/tests/histogram/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-histogram-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/histogram/Cargo.toml b/rust/cuda-coop/tests/histogram/Cargo.toml new file mode 100644 index 000000000000..cf829777fbbe --- /dev/null +++ b/rust/cuda-coop/tests/histogram/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-histogram-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/histogram/src/main.rs b/rust/cuda-coop/tests/histogram/src/main.rs new file mode 100644 index 000000000000..4de995ffddd6 --- /dev/null +++ b/rust/cuda-coop/tests/histogram/src/main.rs @@ -0,0 +1,736 @@ +// 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; + +const CHECKS: u32 = 7; + +#[cuda_module] +mod kernels { + use super::*; + + fn counts( + block: coop::ThreadGroup, + input: &coop::ThreadData, + options: coop::HistogramOptions, + ) -> coop::ThreadData + where + T: coop::HistogramSample, + C: coop::HistogramCounter, + A: coop::HistogramAlgorithm, + S: coop::TempStorageLike, + coop::Alignment: coop::SupportedAlignment, + { + coop::histogram(block, input, options) + } + + fn observe( + counts: &coop::ThreadData, + expected: &[C], + offset: usize, + marker: Option, + ) -> bool { + let mut valid = true; + let mut i = 0; + while i < K { + let value = if i == 0 { + marker.unwrap_or(expected[offset]) + } else { + expected[offset + i * B] + }; + valid &= counts.get(i) == Some(value); + i += 1; + } + valid + } + + // Each root expands exactly one static algorithm/storage choice. Repeated + // expressions create ordinary reborrows of an explicit descriptor. + macro_rules! exercise { + ($n:literal, $align:literal, $b:literal, $k:literal, $options:expr, + $manual:literal, $filled:literal, + $input:ident, $expected:ident, $output:ident, $preserved:ident, + $changed_output:ident, $reports:ident, $replacement:expr, $marker:expr, + $invalid:ident) => {{ + let block = coop::this_block(); + let block_index = block.rank_at::() as usize; + let rank = thread::threadIdx_x() as usize; + let offset = block_index * $b * $n; + let mut values = coop::ThreadData::<_, $n, $align>::new(); + if $invalid != 0 { + let mut i = 0; + while i < $n { + if !($invalid == 6 && rank == $b - 1 && i == $n - 1) { + values.write(i, $input[offset + rank * $n + i]); + } + i += 1; + } + let _ = counts(block, &values, $options); + if $manual { + thread::sync_threads(); + } + // No result read/assertion may counterfeit the preflight trap. + // A lost preflight must return to the host's unconditional panic. + let mut report = coop::ThreadData::<_, 1>::new(); + report.write(0, 99_u32); + coop::store( + block, + &mut $reports, + &report, + coop::StoreOptions::default().offset(block_index * $b), + ); + return; + } + if $filled { + let valid = ($input.len() - offset).min($b * $n); + coop::load( + block, + $input, + &mut values, + coop::LoadOptions::default() + .offset(offset) + .valid_items(valid) + .oob_default($replacement), + ); + } else { + coop::load( + block, + $input, + &mut values, + coop::LoadOptions::default().offset(offset), + ); + } + let mut changed = values; + changed.write(0, $replacement); + let result_offset = block_index * $b * $k; + let mut first = counts(block, &values, $options); + let mut word = + observe::<_, $k, $b>(&first, $expected, result_offset + rank, None) as u32; + first.write(0, $marker); + if $manual { + thread::sync_threads(); + } + let second = counts(block, &changed, $options); + word |= (observe::<_, $k, $b>( + &second, + $expected, + $expected.len() / 2 + result_offset + rank, + None, + ) as u32) + << 1; + word |= (observe::<_, $k, $b>(&first, $expected, result_offset + rank, Some($marker)) + as u32) + << 2; + if $manual { + thread::sync_threads(); + } + coop::store( + block, + &mut $output, + &second, + coop::StoreOptions::default() + .algorithm(coop::Striped) + .offset(result_offset), + ); + coop::store( + block, + &mut $preserved, + &values, + coop::StoreOptions::default().offset(offset), + ); + coop::store( + block, + &mut $changed_output, + &changed, + coop::StoreOptions::default().offset(offset), + ); + let mut report = coop::ThreadData::<_, 1>::new(); + report.write(0, word); + coop::store( + block, + &mut $reports, + &report, + coop::StoreOptions::default().offset(block_index * $b), + ); + }}; + } + + #[allow(clippy::too_many_arguments)] + #[kernel] + #[launch_contract(domain = 1, block = (64, 1, 1), dynamic_shared = inferred)] + pub fn default_u8( + input: &[u8], + expected: &[i32], + mut output: coop::DisjointSlice>, + mut preserved: coop::DisjointSlice>, + mut changed_output: coop::DisjointSlice>, + mut reports: coop::DisjointSlice>, + invalid: u32, + ) { + exercise!( + 1, + 1, + 64, + 1, + coop::HistogramOptions::<33>::default(), + false, + true, + input, + expected, + output, + preserved, + changed_output, + reports, + 0_u8, + 101_i32, + invalid + ); + } + + #[allow(clippy::too_many_arguments)] + #[kernel] + #[launch_contract(domain = 1, block = (33, 1, 1), dynamic_shared = inferred)] + pub fn sort_i32( + input: &[i32], + expected: &[u32], + mut output: coop::DisjointSlice>, + mut preserved: coop::DisjointSlice>, + mut changed_output: coop::DisjointSlice>, + mut reports: coop::DisjointSlice>, + invalid: u32, + ) { + exercise!( + 4, + 16, + 33, + 2, + coop::HistogramOptions::<35>::default() + .counter_type::() + .bins_per_thread::<2>() + .algorithm(coop::Sort), + false, + false, + input, + expected, + output, + preserved, + changed_output, + reports, + 0_i32, + 169_u32, + invalid + ); + } + + #[allow(clippy::too_many_arguments)] + #[kernel] + #[launch_contract(domain = 1, block = (32, 1, 1), dynamic_shared = inferred)] + pub fn atomic_u32( + input: &[u32], + expected: &[i64], + mut output: coop::DisjointSlice>, + mut preserved: coop::DisjointSlice>, + mut changed_output: coop::DisjointSlice>, + mut reports: coop::DisjointSlice>, + invalid: u32, + ) { + exercise!( + 9, + 1, + 32, + 2, + coop::HistogramOptions::<17>::default() + .counter_type::() + .bins_per_thread::<2>(), + false, + false, + input, + expected, + output, + preserved, + changed_output, + reports, + 0_u32, + 325_i64, + invalid + ); + } + + #[allow(clippy::too_many_arguments)] + #[kernel] + #[launch_contract(domain = 1, block = (32, 1, 1), dynamic_shared = inferred)] + pub fn sort_i64( + input: &[i64], + expected: &[u64], + mut output: coop::DisjointSlice>, + mut preserved: coop::DisjointSlice>, + mut changed_output: coop::DisjointSlice>, + mut reports: coop::DisjointSlice>, + invalid: u32, + ) { + exercise!( + 4, + 1, + 32, + 3, + coop::HistogramOptions::<65>::default() + .counter_type::() + .bins_per_thread::<3>() + .algorithm(coop::Sort), + false, + false, + input, + expected, + output, + preserved, + changed_output, + reports, + 0_i64, + 165_u64, + invalid + ); + } + + #[allow(clippy::too_many_arguments)] + #[kernel] + #[launch_contract(domain = 1, block = (1, 1, 1), dynamic_shared = inferred)] + pub fn single_u64( + input: &[u64], + expected: &[i32], + mut output: coop::DisjointSlice>, + mut preserved: coop::DisjointSlice>, + mut changed_output: coop::DisjointSlice>, + mut reports: coop::DisjointSlice>, + invalid: u32, + ) { + exercise!( + 9, + 1, + 1, + 9, + coop::HistogramOptions::<5>::default().bins_per_thread::<9>(), + false, + false, + input, + expected, + output, + preserved, + changed_output, + reports, + 0_u64, + 46_i32, + invalid + ); + } + + #[allow(clippy::too_many_arguments)] + #[kernel] + #[launch_contract(domain = 1, block = (33, 1, 1), dynamic_shared = inferred)] + pub fn manual_i32( + input: &[i32], + expected: &[u64], + mut output: coop::DisjointSlice>, + mut preserved: coop::DisjointSlice>, + mut changed_output: coop::DisjointSlice>, + mut reports: coop::DisjointSlice>, + invalid: u32, + ) { + let mut scratch = coop::TempStorage::new(); + exercise!( + 4, + 1, + 33, + 2, + coop::HistogramOptions::<35>::default() + .counter_type::() + .bins_per_thread::<2>() + .temp_storage(&mut scratch), + true, + false, + input, + expected, + output, + preserved, + changed_output, + reports, + 0_i32, + 169_u64, + invalid + ); + } + + #[allow(clippy::too_many_arguments)] + #[kernel] + #[launch_contract(domain = 1, block = (33, 1, 1), dynamic_shared = inferred)] + pub fn auto_i32( + input: &[i32], + expected: &[u64], + mut output: coop::DisjointSlice>, + mut preserved: coop::DisjointSlice>, + mut changed_output: coop::DisjointSlice>, + mut reports: coop::DisjointSlice>, + invalid: u32, + ) { + let mut scratch = coop::TempStorage::new().auto_sync(); + exercise!( + 4, + 1, + 33, + 2, + coop::HistogramOptions::<35>::default() + .counter_type::() + .bins_per_thread::<2>() + .algorithm(coop::Sort) + .temp_storage(&mut scratch), + false, + false, + input, + expected, + output, + preserved, + changed_output, + reports, + 0_i32, + 169_u64, + invalid + ); + } +} + +// Count each complete initialized tile independently of the device's stripe +// distribution. Padding is part of the expected output, not omitted capacity. +fn oracle(input: &[usize], threads: usize, items: usize, bins: usize, k: usize) -> Vec { + let mut expected = vec![0_u64; input.len() / (threads * items) * threads * k]; + for (block, tile) in input.chunks_exact(threads * items).enumerate() { + for &sample in tile { + assert!(sample < bins); + expected[block * threads * k + sample] += 1; + } + } + expected +} + +fn samples(threads: usize, items: usize, bins: usize) -> Vec { + let tile = threads * items; + (0..3 * tile) + .map(|index| match index / tile { + 0 => bins - 1, + 1 => { + if index % 2 == 0 { + 0 + } else { + bins - 1 + } + } + _ => (7 * (index % tile) + 3) % bins, + }) + .collect() +} + +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}; + let invalid = match std::env::args().nth(1).as_deref() { + None => 0, + Some("sample-at-bins") => 1, + Some("negative-i32") => 2, + Some("negative-i64") => 3, + Some("wide-u32") => 4, + Some("wide-u64") => 5, + Some("missing-input") => 6, + Some(other) => panic!("unknown Histogram mode {other}"), + }; + init_device_contexts(0, 1)?; + let context = with_cuda_context(0, |context| context.clone())?; + let stream = context.default_stream(); + // SAFETY: typed declarations and embedded kernels are compiled together. + let module = unsafe { kernels::load_async(0)? }; + let (mut cases, mut counters, mut values, mut sources, mut reports) = (0, 0, 0, 0, 0); + macro_rules! run { + ($t:ty, $c:ty, $n:literal, $b:literal, $bins:literal, $k:literal, + $prepare:ident, $launch:ident, $blocks:expr, $input:expr) => {{ + let blocks = $blocks; + let input: Vec<$t> = $input; + let mut initialized = input.clone(); + initialized.resize(blocks * $b * $n, 0); + let mut changed = initialized.clone(); + for value in changed.iter_mut().step_by($n) { + *value = 0; + } + let count_slots = blocks * $b * $k; + let (first_expected, second_expected): (Vec<$c>, Vec<$c>) = if invalid == 0 { + let first: Vec = initialized.iter().map(|&v| v as usize).collect(); + let next: Vec = changed.iter().map(|&v| v as usize).collect(); + ( + oracle(&first, $b, $n, $bins, $k) + .into_iter() + .map(|v| v as $c) + .collect(), + oracle(&next, $b, $n, $bins, $k) + .into_iter() + .map(|v| v as $c) + .collect(), + ) + } else { + (vec![0; count_slots], vec![0; count_slots]) + }; + let expected: Vec<$c> = first_expected + .iter() + .chain(&second_expected) + .copied() + .collect(); + let source = DeviceBuffer::from_host(&stream, &input)?; + let expected_device = DeviceBuffer::from_host(&stream, &expected)?; + let output = DeviceBuffer::from_host(&stream, &vec![<$c>::MAX; count_slots])?; + let preserved = DeviceBuffer::from_host(&stream, &vec![<$t>::MAX; initialized.len()])?; + let changed_output = DeviceBuffer::from_host(&stream, &vec![<$t>::MAX; changed.len()])?; + let words = DeviceBuffer::from_host(&stream, &vec![0_u32; blocks * $b])?; + let plan = module.$prepare(LaunchConfig1D::auto_shared(blocks as u32, $b))?; + let (source, _, output, preserved, changed_output, words) = module + .$launch( + &plan, + source, + expected_device, + output, + preserved, + changed_output, + words, + invalid, + ) + .sync()?; + stream.synchronize()?; + if invalid != 0 { + panic!("invalid Histogram mode {invalid} unexpectedly completed"); + } + assert_eq!( + output.to_host_vec(&stream)?, + second_expected, + stringify!($prepare) + ); + assert_eq!( + preserved.to_host_vec(&stream)?, + initialized, + stringify!($prepare) + ); + assert_eq!( + changed_output.to_host_vec(&stream)?, + changed, + stringify!($prepare) + ); + assert_eq!(source.to_host_vec(&stream)?, input, stringify!($prepare)); + assert_eq!( + words.to_host_vec(&stream)?, + vec![CHECKS; blocks * $b], + stringify!($prepare) + ); + cases += 1; + counters += 2 * count_slots; + values += initialized.len() + changed.len(); + sources += input.len(); + reports += blocks * $b; + }}; + } + match invalid { + 1 => { + let mut v = vec![0_u8; 64]; + v[63] = 33; + run!( + u8, + i32, + 1, + 64, + 33, + 1, + prepare_default_u8, + default_u8_async_owned, + 1, + v + ); + } + 2 => { + let mut v = vec![0_i32; 132]; + v[131] = -1; + run!( + i32, + u32, + 4, + 33, + 35, + 2, + prepare_sort_i32, + sort_i32_async_owned, + 1, + v + ); + } + 3 => { + let mut v = vec![0_i64; 128]; + v[127] = i64::MIN; + run!( + i64, + u64, + 4, + 32, + 65, + 3, + prepare_sort_i64, + sort_i64_async_owned, + 1, + v + ); + } + 4 => { + let mut v = vec![0_u32; 288]; + v[287] = u32::MAX; + run!( + u32, + i64, + 9, + 32, + 17, + 2, + prepare_atomic_u32, + atomic_u32_async_owned, + 1, + v + ); + } + 5 => { + let mut v = vec![0_u64; 9]; + v[8] = (1_u64 << 32) + 1; + run!( + u64, + i32, + 9, + 1, + 5, + 9, + prepare_single_u64, + single_u64_async_owned, + 1, + v + ); + } + 6 => { + run!( + u32, + i64, + 9, + 32, + 17, + 2, + prepare_atomic_u32, + atomic_u32_async_owned, + 1, + vec![0_u32; 288] + ); + } + _ => {} + } + let mut short: Vec = samples(64, 1, 33).into_iter().map(|v| v as u8).collect(); + short.truncate(2 * 64 + 45); + run!( + u8, + i32, + 1, + 64, + 33, + 1, + prepare_default_u8, + default_u8_async_owned, + 3, + short + ); + run!( + u8, + i32, + 1, + 64, + 33, + 1, + prepare_default_u8, + default_u8_async_owned, + 1, + Vec::::new() + ); + run!( + i32, + u32, + 4, + 33, + 35, + 2, + prepare_sort_i32, + sort_i32_async_owned, + 3, + samples(33, 4, 35).into_iter().map(|v| v as i32).collect() + ); + run!( + u32, + i64, + 9, + 32, + 17, + 2, + prepare_atomic_u32, + atomic_u32_async_owned, + 3, + samples(32, 9, 17).into_iter().map(|v| v as u32).collect() + ); + run!( + i64, + u64, + 4, + 32, + 65, + 3, + prepare_sort_i64, + sort_i64_async_owned, + 3, + samples(32, 4, 65).into_iter().map(|v| v as i64).collect() + ); + run!( + u64, + i32, + 9, + 1, + 5, + 9, + prepare_single_u64, + single_u64_async_owned, + 3, + samples(1, 9, 5).into_iter().map(|v| v as u64).collect() + ); + run!( + i32, + u64, + 4, + 33, + 35, + 2, + prepare_manual_i32, + manual_i32_async_owned, + 3, + samples(33, 4, 35).into_iter().map(|v| v as i32).collect() + ); + run!( + i32, + u64, + 4, + 33, + 35, + 2, + prepare_auto_i32, + auto_i32_async_owned, + 3, + samples(33, 4, 35).into_iter().map(|v| v as i32).collect() + ); + assert_eq!(cases, 8); + println!( + "histogram: {cases} cases, {counters} counters, {values} preserved values, {sources} source values, {reports} lane reports passed" + ); + Ok(()) +} From ded365f7c7f07dc5a25359b2cf2bb971675a14d9 Mon Sep 17 00:00:00 2001 From: Trent Nelson Date: Sat, 24 Oct 2026 17:43:16 -0700 Subject: [PATCH 09/10] Add Rust Radix Sort Keys and Pairs with Rank Group complete radix key, pair and rank operations with their validity, layout, bit-range and scratch contracts. Signed-off-by: Trent Nelson --- docs/rust/index.rst | 19 +- docs/rust/radix_rank_keys.rst | 84 ++ docs/rust/radix_sort_keys.rst | 95 ++ docs/rust/radix_sort_pairs.rst | 88 ++ rust/cuda-coop-codegen/src/lib.rs | 50 +- rust/cuda-coop-codegen/src/oxide.rs | 387 ++++++- rust/cuda-coop-codegen/src/radix_sort.rs | 671 +++++++++++ rust/cuda-coop-codegen/tests/native.rs | 238 +++- rust/cuda-coop/build.rs | 47 +- .../examples/radix_rank_keys/.gitignore | 7 + .../examples/radix_rank_keys/Cargo.lock | 1019 +++++++++++++++++ .../examples/radix_rank_keys/Cargo.toml | 19 + .../radix_rank_keys/helper/Cargo.toml | 11 + .../radix_rank_keys/helper/src/lib.rs | 52 + .../examples/radix_rank_keys/src/main.rs | 116 ++ .../examples/radix_sort_keys/.gitignore | 7 + .../examples/radix_sort_keys/Cargo.lock | 1019 +++++++++++++++++ .../examples/radix_sort_keys/Cargo.toml | 19 + .../radix_sort_keys/helper/Cargo.toml | 11 + .../radix_sort_keys/helper/src/lib.rs | 49 + .../examples/radix_sort_keys/src/main.rs | 151 +++ .../examples/radix_sort_pairs/.gitignore | 7 + .../examples/radix_sort_pairs/Cargo.lock | 1019 +++++++++++++++++ .../examples/radix_sort_pairs/Cargo.toml | 19 + .../radix_sort_pairs/helper/Cargo.toml | 11 + .../radix_sort_pairs/helper/src/lib.rs | 77 ++ .../examples/radix_sort_pairs/src/main.rs | 182 +++ rust/cuda-coop/src/lib.rs | 13 + rust/cuda-coop/src/radix_sort.rs | 564 +++++++++ rust/cuda-coop/src/radix_sort/boundary.rs | 290 +++++ rust/cuda-coop/tests/check_load_store.py | 236 ++++ .../tests/radix_rank_keys/.gitignore | 7 + .../tests/radix_rank_keys/Cargo.lock | 1011 ++++++++++++++++ .../tests/radix_rank_keys/Cargo.toml | 17 + .../tests/radix_rank_keys/src/main.rs | 319 ++++++ .../tests/radix_sort_keys/.gitignore | 7 + .../tests/radix_sort_keys/Cargo.lock | 1011 ++++++++++++++++ .../tests/radix_sort_keys/Cargo.toml | 17 + .../tests/radix_sort_keys/src/main.rs | 522 +++++++++ .../tests/radix_sort_pairs/.gitignore | 7 + .../tests/radix_sort_pairs/Cargo.lock | 1011 ++++++++++++++++ .../tests/radix_sort_pairs/Cargo.toml | 17 + .../tests/radix_sort_pairs/src/main.rs | 442 +++++++ 43 files changed, 10941 insertions(+), 24 deletions(-) create mode 100644 docs/rust/radix_rank_keys.rst create mode 100644 docs/rust/radix_sort_keys.rst create mode 100644 docs/rust/radix_sort_pairs.rst create mode 100644 rust/cuda-coop-codegen/src/radix_sort.rs create mode 100644 rust/cuda-coop/examples/radix_rank_keys/.gitignore create mode 100644 rust/cuda-coop/examples/radix_rank_keys/Cargo.lock create mode 100644 rust/cuda-coop/examples/radix_rank_keys/Cargo.toml create mode 100644 rust/cuda-coop/examples/radix_rank_keys/helper/Cargo.toml create mode 100644 rust/cuda-coop/examples/radix_rank_keys/helper/src/lib.rs create mode 100644 rust/cuda-coop/examples/radix_rank_keys/src/main.rs create mode 100644 rust/cuda-coop/examples/radix_sort_keys/.gitignore create mode 100644 rust/cuda-coop/examples/radix_sort_keys/Cargo.lock create mode 100644 rust/cuda-coop/examples/radix_sort_keys/Cargo.toml create mode 100644 rust/cuda-coop/examples/radix_sort_keys/helper/Cargo.toml create mode 100644 rust/cuda-coop/examples/radix_sort_keys/helper/src/lib.rs create mode 100644 rust/cuda-coop/examples/radix_sort_keys/src/main.rs create mode 100644 rust/cuda-coop/examples/radix_sort_pairs/.gitignore create mode 100644 rust/cuda-coop/examples/radix_sort_pairs/Cargo.lock create mode 100644 rust/cuda-coop/examples/radix_sort_pairs/Cargo.toml create mode 100644 rust/cuda-coop/examples/radix_sort_pairs/helper/Cargo.toml create mode 100644 rust/cuda-coop/examples/radix_sort_pairs/helper/src/lib.rs create mode 100644 rust/cuda-coop/examples/radix_sort_pairs/src/main.rs create mode 100644 rust/cuda-coop/src/radix_sort.rs create mode 100644 rust/cuda-coop/src/radix_sort/boundary.rs create mode 100644 rust/cuda-coop/tests/radix_rank_keys/.gitignore create mode 100644 rust/cuda-coop/tests/radix_rank_keys/Cargo.lock create mode 100644 rust/cuda-coop/tests/radix_rank_keys/Cargo.toml create mode 100644 rust/cuda-coop/tests/radix_rank_keys/src/main.rs create mode 100644 rust/cuda-coop/tests/radix_sort_keys/.gitignore create mode 100644 rust/cuda-coop/tests/radix_sort_keys/Cargo.lock create mode 100644 rust/cuda-coop/tests/radix_sort_keys/Cargo.toml create mode 100644 rust/cuda-coop/tests/radix_sort_keys/src/main.rs create mode 100644 rust/cuda-coop/tests/radix_sort_pairs/.gitignore create mode 100644 rust/cuda-coop/tests/radix_sort_pairs/Cargo.lock create mode 100644 rust/cuda-coop/tests/radix_sort_pairs/Cargo.toml create mode 100644 rust/cuda-coop/tests/radix_sort_pairs/src/main.rs diff --git a/docs/rust/index.rst b/docs/rust/index.rst index 74b9e45c50d4..23e13cb5b514 100644 --- a/docs/rust/index.rst +++ b/docs/rust/index.rst @@ -45,7 +45,8 @@ group's tile, and must be uniform across that group. The public operations include :doc:`load_store`, :doc:`scan`, :doc:`reduce`, :doc:`reduce_batched`, :doc:`exchange`, :doc:`shuffle`, -:doc:`adjacent_difference`, :doc:`discontinuity` and :doc:`histogram`. +:doc:`adjacent_difference`, :doc:`discontinuity`, :doc:`histogram`, +:doc:`radix_sort_keys`, :doc:`radix_sort_pairs` and :doc:`radix_rank_keys`. Scan returns fresh per-group prefixes; Reduce returns a total only at each group's root. Batched reduction returns a fresh payload whose extent and valid slots follow its Warp width and output layout. @@ -53,11 +54,16 @@ Exchange converts between blocked and striped per-thread layouts. Shuffle shifts blocked values by one item, leaving the exposed edge invalid. Discontinuity returns head flags, tail flags or both as fresh ``i32`` payloads. Histogram returns fresh striped counters with an independent type and extent. -It requires generic native-registration support beyond the currently pinned -Oxide revision; see :doc:`histogram` for the compiler requirement. +Radix Sort Keys stably sorts integer keys into a fresh blocked payload. +Radix Sort Pairs carries associated numeric values through that stable order +and returns two fresh payloads. Radix Rank Keys returns each input's stable +selected-digit position as a fresh ``i32`` payload. Histogram, Radix Sort and +Radix Rank require generic native-registration support beyond the currently +pinned Oxide revision; their guides explain the compiler requirement. Scan, Reduce, batched reduction, Exchange, Shuffle, Adjacent Difference, -Discontinuity and Histogram preserve their input. Additional primitive families -are under development. The crate is unpublished and its interface can change. +Discontinuity, Histogram, Radix Sort and Radix Rank preserve their inputs. +Additional primitive families are under development. The crate is unpublished +and its interface can change. Build and run ------------- @@ -132,6 +138,9 @@ Read next adjacent_difference discontinuity histogram + radix_sort_keys + radix_sort_pairs + radix_rank_keys provider Generate the API reference with ``cargo doc --workspace --no-deps`` from diff --git a/docs/rust/radix_rank_keys.rst b/docs/rust/radix_rank_keys.rst new file mode 100644 index 000000000000..398f17a3992d --- /dev/null +++ b/docs/rust/radix_rank_keys.rst @@ -0,0 +1,84 @@ +.. Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. +.. SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +Ranking integer key digits +========================== + +``radix_rank_keys`` assigns each input item its zero-based position in the +stable order of selected key digits across a complete Block: + +.. code-block:: rust + + let ranks = coop::radix_rank_keys(block, &keys, Default::default()); + +The default ranks the low four bits in ascending order. Each rank stays with +its original input item; equal digits receive increasing ranks in original +blocked order. The fresh result is ``ThreadData`` with +every slot initialized and natural ``i32`` alignment, independently of the key +payload's alignment. Input values and validity remain unchanged. + +A complete program +------------------ + +This helper infers its key type from Load and forwards a generic item count and +alignment. It returns ascending and descending ranks plus the original payload. +The ascending interval derives its width from static endpoints; the descending +interval supplies an explicit width with a runtime begin: + +.. literalinclude:: ../../rust/cuda-coop/examples/radix_rank_keys/helper/src/lib.rs + :language: rust + :start-after: // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +The host instantiates one and four items per thread on two 4-by-4-by-4 Blocks. +It uses ``i64`` keys aligned to at least 32 bytes and selects the three high bits, +including the signed-key sign bit. Distinct lower-bit tags distinguish digit +ranking from whole-key ordering. Independent host sorts check stable ties by +inverting each sorted permutation, then compare both rank payloads and verify +the preserved payload and source keys: + +.. literalinclude:: ../../rust/cuda-coop/examples/radix_rank_keys/src/main.rs + :language: rust + :start-after: // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +With the pinned Oxide revision and Rust ``1.99.1`` +installed, run from a complete CCCL checkout: + +.. code-block:: shell + + cd rust/cuda-coop/examples/radix_rank_keys + 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. ``auto_shared`` uses the compiler's +complete temporary-storage plan. + +Digits and participation +------------------------ + +Keys support ``i32``, ``u32``, ``i64`` and ``u64``. Signed keys invert their sign +bit before digit extraction. ``RadixRankOptions`` defaults to ascending order, +begin zero and width four. ``order(coop::Descending)`` selects descending order +while preserving tie order. + +The static radix width must be in ``1..=8``. ``radix_bits::()`` selects it +explicitly. With a static end and no explicit width, the width is +``end - begin``; use ``begin_bit_const::()`` and ``end_bit_const::()`` to +select static endpoints. A Block-uniform ``begin_bit(i64)`` requires an explicit +width, even when passed a literal. An optional end is always static and must +equal begin plus width. Builders preserve other selections and validate the +final combined choices. + +Static choices are checked when the operation is instantiated. Runtime bounds +must satisfy ``0 <= begin`` and ``begin + width <= key_width``, together with +any selected end. Invalid bounds and unreadable input slots trap before numeric +reads. Every input slot must be initialized. + +Every member of the complete XYZ Block participates. Positive per-thread item +counts are supported; the tile contains at most 65,535 keys. Input and result +slots use blocked order: x-fastest thread rank r owns items +``r * ITEMS_PER_THREAD + i``. + +Rank uses implicit compiler-planned CUB scratch with synchronization before +dirty reuse. Internal barriers do not complete unrelated global-memory +accesses. See :doc:`payload` and :doc:`groups` for initialization and +participation, or :doc:`radix_sort_keys` to return keys in sorted order. diff --git a/docs/rust/radix_sort_keys.rst b/docs/rust/radix_sort_keys.rst new file mode 100644 index 000000000000..7cc95b8589ab --- /dev/null +++ b/docs/rust/radix_sort_keys.rst @@ -0,0 +1,95 @@ +.. Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. +.. SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +Sorting integer keys +==================== + +``radix_sort_keys`` stably sorts a complete Block's integer keys and returns a +fresh payload with the same type, item count and alignment. Default options +sort all bits in ascending order: + +.. code-block:: rust + + let sorted = coop::radix_sort_keys(block, &keys, Default::default()); + +The input remains unchanged. Every input slot must be initialized; unreadable +slots trap before native sorting reads. Results use blocked order: thread rank +r owns positions ``r * ITEMS_PER_THREAD + i``. Rank is x-fastest across the +Block's three dimensions. + +A complete program +------------------ + +This helper infers its key type from Load and forwards generic item count and +alignment. It returns a full-width sort, two selected-bit sorts and the original +payload: + +.. literalinclude:: ../../rust/cuda-coop/examples/radix_sort_keys/helper/src/lib.rs + :language: rust + :start-after: // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +The host instantiates one and four items per thread on 4-by-4-by-4 Blocks. +Signed extrema and keys sharing selected digits exercise ordering and stable +ties. Independent host sorts check all results and input preservation: + +.. literalinclude:: ../../rust/cuda-coop/examples/radix_sort_keys/src/main.rs + :language: rust + :start-after: // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +With the pinned Oxide revision and Rust ``1.99.1`` +installed, run from a complete CCCL checkout: + +.. code-block:: shell + + cd rust/cuda-coop/examples/radix_sort_keys + 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. ``auto_shared`` uses the compiler's +complete temporary-storage plan. + +Selecting bits and order +------------------------ + +Keys support ``i32``, ``u32``, ``i64`` and ``u64``. The default interval is +``[0, key_width)``. ``begin_bit`` and ``end_bit`` independently replace an +endpoint; an omitted end always remains the full key width. Use +``order(coop::Descending)`` for descending order. Both directions retain the +original blocked order of keys whose selected digits are equal. + +Signed keys invert their sign bit before selecting the interval. The result +contains the original key representation, including the bits outside that +interval. + +Value builders accept Block-uniform ``i64`` endpoints. The native operation +checks ``0 <= begin < end <= key_width`` before narrowing them or reading keys. +Empty intervals and out-of-range runtime endpoints trap. A literal passed to a +value builder follows this runtime route. To request compile-time validation, +use ``begin_bit_const::()`` or ``end_bit_const::()``: + +.. code-block:: rust + + let selected = coop::radix_sort_keys(block, &keys, + coop::RadixSortOptions::default() + .begin_bit_const::<8>().end_bit(16).order(coop::Descending)); + +Static endpoints are checked when the operation is instantiated. Mixed static +and runtime endpoints are supported. Builders preserve other selections and +validate the final options, so replacing an endpoint uses its final value. + +Participation and storage +------------------------- + +Every thread in a complete Block participates. Positive per-thread extents and +one-, two- or three-dimensional Blocks are supported; the complete tile must +contain at most 65,535 keys. This operation has no partial-count or fill option. + +Omitted storage uses compiler-planned CUB scratch with synchronization before +dirty reuse. ``temp_storage(&mut scratch)`` borrows an explicit descriptor; +manual reuse requires ``cuda_device::thread::sync_threads()`` after the call, +and ``auto_sync()`` supplies that trailing barrier. Internal CUB barriers do +not complete unrelated global-memory accesses. + +See :doc:`payload`, :doc:`groups` and :doc:`load_store` for checked access, +participation and storage. Generate the full API reference with +``cargo doc --workspace --no-deps`` from ``rust/``. diff --git a/docs/rust/radix_sort_pairs.rst b/docs/rust/radix_sort_pairs.rst new file mode 100644 index 000000000000..d663f82deab9 --- /dev/null +++ b/docs/rust/radix_sort_pairs.rst @@ -0,0 +1,88 @@ +.. Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. +.. SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +Sorting key/value pairs +======================= + +``radix_sort_pairs`` stably sorts a complete Block's integer keys and carries +each associated value with its key. It borrows two payloads and returns two +fresh payloads: + +.. code-block:: rust + + let (sorted_keys, sorted_values) = coop::radix_sort_pairs( + block, &keys, &values, Default::default()); + +The default sorts all key bits in ascending order. Keys support ``i32``, ``u32``, +``i64`` and ``u64``. Values support ``i8``, ``u8``, ``i16``, ``u16``, ``i32``, +``u32``, ``i64``, ``u64``, ``f32`` and ``f64``. Values are moved without +comparisons or arithmetic; floating-point values retain their bit patterns. + +Both inputs remain unchanged and must have the same positive item count per +thread. Their minimum alignments are independent. Each output retains its +input's type, item count and alignment; changing a result does not change either +input or another result. Every slot in both inputs must be initialized. Both +validity maps are checked before native sorting reads either numeric array. + +A complete program +------------------ + +This separate helper infers key and value types from Load and forwards a generic +item count and independent alignments. It returns ascending and descending sorts +of selected digits together with both original payloads: + +.. literalinclude:: ../../rust/cuda-coop/examples/radix_sort_pairs/helper/src/lib.rs + :language: rust + :start-after: // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +The host instantiates one and four items per thread on 4-by-4-by-4 Blocks, with +``i32`` keys and ``f64`` values aligned to at least 16 and 32 bytes respectively. +Selected digits include the signed-key sign bit. Distinct lower-bit tags and +associated values distinguish stable selected-digit ordering from sorting whole +keys. The program checks both sorts, both preserved inputs and mutations of +fresh results; floating values are compared by their bits: + +.. literalinclude:: ../../rust/cuda-coop/examples/radix_sort_pairs/src/main.rs + :language: rust + :start-after: // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +With the pinned Oxide revision and Rust ``1.99.1`` +installed, run from a complete CCCL checkout: + +.. code-block:: shell + + cd rust/cuda-coop/examples/radix_sort_pairs + 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. ``auto_shared`` uses the compiler's +complete temporary-storage plan for the paired primitive. + +Order, bounds and storage +------------------------- + +``RadixSortOptions`` is shared with :doc:`radix_sort_keys`. The default interval +is ``[0, key_width)``. ``begin_bit`` and ``end_bit`` independently replace an +endpoint; an omitted end remains the key width even with a nonzero begin. +``order(coop::Descending)`` selects descending order. Both directions preserve +the original order of pairs whose selected key digits are equal. Signed keys +invert the sign bit before selecting digits and restore their original +representation in the result. + +Value builders accept Block-uniform ``i64`` endpoints. The native operation +checks ``0 <= begin < end <= key_width`` before narrowing or reading inputs. +Invalid runtime intervals, including empty ones, trap. Const builders +``begin_bit_const::()`` and ``end_bit_const::()`` request checks when the +operation is instantiated. Mixed static and runtime bounds are supported. + +Every member of the complete XYZ Block participates. Input and output use +blocked order: x-fastest thread rank r owns items ``r * ITEMS_PER_THREAD + i``. +The complete tile contains at most 65,535 pairs. There is no partial-input or +fill option. + +Omitted scratch uses compiler-planned CUB storage with synchronization before +dirty reuse. ``temp_storage(&mut scratch)`` retains the descriptor's policy: +manual reuse requires ``cuda_device::thread::sync_threads()`` after the call, +while ``auto_sync()`` supplies that completion. Internal sorting barriers do not +complete unrelated global-memory accesses. See :doc:`payload`, :doc:`groups` +and :doc:`load_store` for initialization, participation and checked access. diff --git a/rust/cuda-coop-codegen/src/lib.rs b/rust/cuda-coop-codegen/src/lib.rs index dce4530209fa..acb62afb4f3a 100644 --- a/rust/cuda-coop-codegen/src/lib.rs +++ b/rust/cuda-coop-codegen/src/lib.rs @@ -4,8 +4,8 @@ //! Native, host-only CUB specialization for Rust cooperative algorithms. //! //! This private development crate compiles Load/Store, Scan, Reduce, batched -//! Reduce, Exchange, Shuffle, Adjacent Difference, Discontinuity and Histogram -//! with NVRTC. +//! Reduce, Exchange, Shuffle, Adjacent Difference, Discontinuity, Histogram +//! and Radix Sort Keys/Pairs and Rank Keys 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. @@ -56,7 +56,12 @@ //! and returning fresh i32 head flags, tail flags or both, with complete validity. //! [`HistogramRequest`] counts bounded integer samples in a complete 1D Block, //! with independent counter type/output extent and initialized zero padding. -//! [`Provider::compile_native`] batches these families in one native compilation. +//! [`RadixSortKeysRequest`] stably sorts a complete integer Block tile over a +//! uniform bit range, preserving input and producing fresh blocked keys. +//! [`RadixSortPairsRequest`] carries independently typed values with each key +//! into fresh paired outputs. [`RadixRankKeysRequest`] returns stable selected-digit +//! i32 ranks in original input-slot order. [`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 @@ -90,6 +95,7 @@ mod discontinuity; mod exchange; mod headers; mod histogram; +mod radix_sort; mod reduce; mod reduce_batched; mod scan; @@ -98,6 +104,10 @@ pub use adjacent_difference::{AdjacentDifferenceRecipe, AdjacentDifferenceReques pub use discontinuity::{DiscontinuityRecipe, DiscontinuityRequest}; pub use exchange::{ExchangeMode, ExchangeRecipe, ExchangeRequest}; pub use histogram::{HistogramAlgorithm, HistogramRecipe, HistogramRequest}; +pub use radix_sort::{ + RadixRankKeysRecipe, RadixRankKeysRequest, RadixSortKeysRecipe, RadixSortKeysRequest, + RadixSortPairsRecipe, RadixSortPairsRequest, +}; pub use reduce::{ReduceAlgorithm, ReduceForm, ReduceOperator, ReduceRecipe, ReduceRequest}; pub use reduce_batched::{ReduceBatchedLayout, ReduceBatchedRecipe, ReduceBatchedRequest}; pub use scan::{ScanAlgorithm, ScanOperation, ScanOperator, ScanRecipe, ScanRequest}; @@ -498,6 +508,12 @@ pub enum NativeRequest { Discontinuity(DiscontinuityRequest), /// Bounded integer samples with independent counter type and output extent. Histogram(HistogramRequest), + /// Stable Block radix sort with preserved input and fresh blocked keys. + RadixSortKeys(RadixSortKeysRequest), + /// Stable Block radix sort preserving independently typed key/value inputs. + RadixSortPairs(RadixSortPairsRequest), + /// Stable selected-digit Block ranks with preserved keys and fresh i32 output. + RadixRankKeys(RadixRankKeysRequest), } impl NativeRequest { @@ -512,6 +528,9 @@ impl NativeRequest { Self::AdjacentDifference(request) => request.block_shape(), Self::Discontinuity(request) => request.block_shape(), Self::Histogram(request) => request.block_shape(), + Self::RadixSortKeys(request) => request.block_shape(), + Self::RadixSortPairs(request) => request.block_shape(), + Self::RadixRankKeys(request) => request.block_shape(), } } @@ -526,6 +545,9 @@ impl NativeRequest { Self::AdjacentDifference(request) => request.source(symbol), Self::Discontinuity(request) => request.source(symbol), Self::Histogram(request) => request.source(symbol), + Self::RadixSortKeys(request) => request.source(symbol), + Self::RadixSortPairs(request) => request.source(symbol), + Self::RadixRankKeys(request) => request.source(symbol), } } @@ -540,6 +562,9 @@ impl NativeRequest { Self::AdjacentDifference(request) => request.recipe.token().into(), Self::Discontinuity(request) => request.recipe.token().into(), Self::Histogram(request) => request.recipe.token(), + Self::RadixSortKeys(request) => request.recipe.token().into(), + Self::RadixSortPairs(request) => request.recipe.token(request.value_dtype), + Self::RadixRankKeys(request) => request.recipe.token(), } } @@ -553,7 +578,10 @@ impl NativeRequest { | Self::Shuffle(_) | Self::AdjacentDifference(_) | Self::Discontinuity(_) - | Self::Histogram(_) => false, + | Self::Histogram(_) + | Self::RadixSortKeys(_) + | Self::RadixSortPairs(_) + | Self::RadixRankKeys(_) => false, } } @@ -570,7 +598,10 @@ impl NativeRequest { | Self::Shuffle(_) | Self::AdjacentDifference(_) | Self::Discontinuity(_) - | Self::Histogram(_) => false, + | Self::Histogram(_) + | Self::RadixSortKeys(_) + | Self::RadixSortPairs(_) + | Self::RadixRankKeys(_) => false, } } } @@ -590,10 +621,13 @@ pub struct Layout { /// /// The memory-tile and payload-layout requirements below describe Load/Store. /// Scan, Reduce, batched Reduce, Exchange, Shuffle, Adjacent Difference, -/// Discontinuity and Histogram use lane-private input and fresh output. Their +/// Discontinuity, Histogram and Radix Sort Keys use lane-private input and fresh output. Their /// ABI and validity contracts are documented by [`ScanRequest`], [`ReduceRequest`], /// [`ReduceBatchedRequest`], [`ExchangeRequest`], [`ShuffleRequest`], -/// [`AdjacentDifferenceRequest`], [`DiscontinuityRequest`] and [`HistogramRequest`]. +/// [`AdjacentDifferenceRequest`], [`DiscontinuityRequest`], [`HistogramRequest`] +/// and [`RadixSortKeysRequest`]. [`RadixSortPairsRequest`] returns fresh keys and +/// independently typed associated values while preserving both inputs. +/// [`RadixRankKeysRequest`] returns fresh i32 ranks independently of key type. /// /// 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. @@ -747,7 +781,7 @@ impl Provider { /// Compile an ordered mixed batch of the supported families in one NVRTC program. /// /// Per-specialization identity is independent of unrelated batch members. - /// Scan/Reduce/batched Reduce/Exchange/Shuffle/Adjacent Difference/Discontinuity + /// Scan/Reduce/batched Reduce/Exchange/Shuffle/Adjacent Difference/Discontinuity/Radix Sort Keys /// scratch describes one group and includes the actual CUB TempStorage, even /// when its C++ empty type occupies one byte. The caller plans replication. /// diff --git a/rust/cuda-coop-codegen/src/oxide.rs b/rust/cuda-coop-codegen/src/oxide.rs index c958469fabaa..54dadb026875 100644 --- a/rust/cuda-coop-codegen/src/oxide.rs +++ b/rust/cuda-coop-codegen/src/oxide.rs @@ -22,6 +22,11 @@ //! [`DiscontinuityRecipe`] selects Block head/tail flags with one or two i32 //! outputs and optional neighbors of the input type. [`HistogramRecipe`] selects //! bounded integer samples with an independent counter type and output extent. +//! [`RadixSortKeysRecipe`] selects stable Block key sorting with two signed64 +//! uniform bit bounds and fresh blocked output. [`RadixSortPairsRecipe`] also +//! carries an independent numeric value type into a fresh paired output. +//! [`RadixRankKeysRecipe`] selects stable i32 ranks with a static digit width +//! and signed64 begin/end-constraint controls, preserving input-slot order. //! 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 @@ -31,9 +36,11 @@ use crate::{ AdjacentDifferenceRecipe, AdjacentDifferenceRequest, Algorithm, Artifact, DiscontinuityRecipe, DiscontinuityRequest, Dtype, Error, ExchangeRecipe, ExchangeRequest, HistogramAlgorithm, - HistogramRecipe, HistogramRequest, Layout, NativeRequest, Operation, Provider, ReduceAlgorithm, - ReduceBatchedRecipe, ReduceBatchedRequest, ReduceForm, ReduceOperator, ReduceRecipe, - ReduceRequest, Request, Result, ScanRecipe, ScanRequest, ShuffleRecipe, ShuffleRequest, + HistogramRecipe, HistogramRequest, Layout, NativeRequest, Operation, Provider, + RadixRankKeysRecipe, RadixRankKeysRequest, RadixSortKeysRecipe, RadixSortKeysRequest, + RadixSortPairsRecipe, RadixSortPairsRequest, ReduceAlgorithm, ReduceBatchedRecipe, + ReduceBatchedRequest, ReduceForm, ReduceOperator, ReduceRecipe, ReduceRequest, Request, Result, + ScanRecipe, ScanRequest, ShuffleRecipe, ShuffleRequest, }; use serde::{Deserialize, Serialize}; use std::io::{Read, Write}; @@ -200,6 +207,47 @@ struct SpecializationRequest { impl SpecializationRequest { fn native(&self) -> Result { + if self.operation.starts_with("radix_rank_keys_") { + let recipe = RadixRankKeysRecipe::parse(&self.operation, &self.algorithm)?; + if self.warp_width.is_some() { + return Err(Error::message("Radix Rank Keys requires a Block scope")); + } + return RadixRankKeysRequest::block( + self.dtype.dtype()?, + recipe, + self.block, + self.items_per_thread, + ) + .map(NativeRequest::RadixRankKeys); + } + if self.operation.starts_with("radix_sort_pairs_") { + let (recipe, value_dtype) = + RadixSortPairsRecipe::parse(&self.operation, &self.algorithm)?; + if self.warp_width.is_some() { + return Err(Error::message("Radix Sort Pairs requires a Block scope")); + } + return RadixSortPairsRequest::block( + self.dtype.dtype()?, + value_dtype, + recipe, + self.block, + self.items_per_thread, + ) + .map(NativeRequest::RadixSortPairs); + } + if self.operation.starts_with("radix_sort_keys_") { + let recipe = RadixSortKeysRecipe::parse(&self.operation, &self.algorithm)?; + if self.warp_width.is_some() { + return Err(Error::message("Radix Sort Keys requires a Block scope")); + } + return RadixSortKeysRequest::block( + self.dtype.dtype()?, + recipe, + self.block, + self.items_per_thread, + ) + .map(NativeRequest::RadixSortKeys); + } if self.operation.starts_with("histogram_") { let recipe = HistogramRecipe::parse(&self.operation, &self.algorithm)?; if self.warp_width.is_some() { @@ -405,6 +453,43 @@ struct BatchResponse { fn arguments(request: NativeRequest) -> Vec { let request = match request { + NativeRequest::RadixRankKeys(request) => { + return vec![ + Argument::ConstPointer(Scalar::from(request.dtype)), + Argument::MutPointer(Scalar::Signed(32)), + Argument::ConstPointer(Scalar::Unsigned(8)), + Argument::MutPointer(Scalar::Unsigned(8)), + Argument::Value(Scalar::Signed(64)), + Argument::Value(Scalar::Signed(64)), + ]; + } + NativeRequest::RadixSortPairs(request) => { + let key = Scalar::from(request.key_dtype); + let value = Scalar::from(request.value_dtype); + return vec![ + Argument::ConstPointer(key), + Argument::MutPointer(key), + Argument::ConstPointer(value), + Argument::MutPointer(value), + Argument::ConstPointer(Scalar::Unsigned(8)), + Argument::MutPointer(Scalar::Unsigned(8)), + Argument::ConstPointer(Scalar::Unsigned(8)), + Argument::MutPointer(Scalar::Unsigned(8)), + Argument::Value(Scalar::Signed(64)), + Argument::Value(Scalar::Signed(64)), + ]; + } + NativeRequest::RadixSortKeys(request) => { + let dtype = Scalar::from(request.dtype); + return vec![ + Argument::ConstPointer(dtype), + Argument::MutPointer(dtype), + Argument::ConstPointer(Scalar::Unsigned(8)), + Argument::MutPointer(Scalar::Unsigned(8)), + Argument::Value(Scalar::Signed(64)), + Argument::Value(Scalar::Signed(64)), + ]; + } NativeRequest::Histogram(request) => { return vec![ Argument::ConstPointer(Scalar::from(request.dtype)), @@ -764,6 +849,92 @@ pub fn shuffle_registration(executable: &Path, recipe: ShuffleRecipe) -> Result< .map_err(|error| Error::message(format!("cannot encode provider registration: {error}"))) } +/// Encode one of sixteen Block Radix Rank Keys order/width registrations. +/// +/// `recipe` selects ascending/descending order and 1..=8 radix bits. Key type, +/// exact XYZ and per-lane extent arrive with specialization. Four pointers +/// describe preserved keys, fresh i32 ranks and independent validity maps. +/// Two signed64 uniform controls follow: begin bit and an explicit static end +/// constraint, or -1 for omitted end. Scratch is inserted at slot 2, giving +/// seven native arguments. See [`RadixRankKeysRequest`] for the native contract. +/// +/// # Errors +/// Rejects widths outside 1..=8, a nonabsolute executable or an unencodable path. +pub fn radix_rank_keys_registration( + executable: &Path, + recipe: RadixRankKeysRecipe, +) -> Result { + recipe.validate()?; + 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: "block", + }) + .map_err(|error| Error::message(format!("cannot encode provider registration: {error}"))) +} + +/// Encode a Block Radix Sort Keys recipe for the existing provider executable. +/// +/// Key dtype, exact XYZ geometry and per-lane extent arrive with specialization. +/// Four typed pointers describe preserved keys and fresh same-extent output +/// with independent packed validity maps. Two signed64 uniform controls follow: +/// beginning bit and exclusive ending bit. The existing ABI inserts scratch +/// at slot 2. Bounds are runtime values, never part of specialization identity. +/// See [`RadixSortKeysRequest`] for bounds, layout and participation obligations. +/// +/// # Errors +/// Rejects a nonabsolute executable or a path that cannot be encoded as JSON. +pub fn radix_sort_keys_registration( + executable: &Path, + recipe: RadixSortKeysRecipe, +) -> 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: "block", + }) + .map_err(|error| Error::message(format!("cannot encode provider registration: {error}"))) +} + +/// Encode one of the twenty typed Block Radix Sort Pairs registrations. +/// +/// `recipe` selects order and `value_dtype` selects the independent value type. +/// Key dtype, common per-lane extent and exact XYZ arrive with specialization. +/// Eight typed pointers describe key/value inputs, fresh outputs and four +/// independent validity maps. Two signed64 uniform bit bounds follow; the +/// existing native ABI inserts scratch at slot 2, for eleven arguments total. +/// See [`RadixSortPairsRequest`] for the complete native contract. +/// +/// # Errors +/// Rejects a nonabsolute executable or a path that cannot be encoded as JSON. +pub fn radix_sort_pairs_registration( + executable: &Path, + recipe: RadixSortPairsRecipe, + value_dtype: Dtype, +) -> 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(value_dtype), + algorithm: "block", + }) + .map_err(|error| Error::message(format!("cannot encode provider registration: {error}"))) +} + /// Encode a Block Adjacent Difference recipe for the existing provider executable. /// /// Dtype, exact Block geometry and payload extent arrive with specialization. @@ -1176,6 +1347,216 @@ mod tests { ScanOperator, }; + #[test] + fn radix_sort_pairs_registration_preserves_independent_value_abi() { + for value in [ + Dtype::I8, + Dtype::U8, + Dtype::I16, + Dtype::U16, + Dtype::I32, + Dtype::U32, + Dtype::I64, + Dtype::U64, + Dtype::F32, + Dtype::F64, + ] { + for recipe in [ + RadixSortPairsRecipe::Ascending, + RadixSortPairsRecipe::Descending, + ] { + let registration: serde_json::Value = serde_json::from_str( + &radix_sort_pairs_registration(Path::new("/provider"), recipe, value).unwrap(), + ) + .unwrap(); + let mut request = SpecializationRequest { + operation: registration["operation"].as_str().unwrap().into(), + algorithm: registration["algorithm"].as_str().unwrap().into(), + dtype: Scalar::Signed(32), + block: [4, 4, 4], + items_per_thread: 4, + warp_width: None, + }; + let native = request.native().unwrap(); + assert_eq!( + native, + NativeRequest::RadixSortPairs( + RadixSortPairsRequest::block(Dtype::I32, value, recipe, [4, 4, 4], 4) + .unwrap() + ) + ); + let value = Scalar::from(value); + assert_eq!( + arguments(native), + vec![ + Argument::ConstPointer(Scalar::Signed(32)), + Argument::MutPointer(Scalar::Signed(32)), + Argument::ConstPointer(value), + Argument::MutPointer(value), + Argument::ConstPointer(Scalar::Unsigned(8)), + Argument::MutPointer(Scalar::Unsigned(8)), + Argument::ConstPointer(Scalar::Unsigned(8)), + Argument::MutPointer(Scalar::Unsigned(8)), + Argument::Value(Scalar::Signed(64)), + Argument::Value(Scalar::Signed(64)), + ] + ); + assert_eq!(native.block_shape(), [4, 4, 4]); + assert!(!native.mutates_payload() && !native.scratch_free()); + request.warp_width = Some(32); + assert!(request.native().is_err()); + request.warp_width = None; + request.dtype = Scalar::Float(32); + assert!(request.native().is_err()); + request.dtype = Scalar::Signed(64); + request.block = [1, 1, 1]; + request.items_per_thread = 65535; + assert!(request.native().is_ok()); + request.block = [2, 1, 1]; + assert!(request.native().is_err()); + request.items_per_thread = 4; + request.algorithm = "warp".into(); + assert!(request.native().is_err()); + request.algorithm = "block".into(); + request.operation = "radix_sort_pairs_ascending_f16_validity_bits_v1".into(); + assert!(request.native().is_err()); + } + } + } + + #[test] + fn radix_rank_registration_preserves_width_and_independent_output_abi() { + for bits in 1..=8 { + for recipe in [ + RadixRankKeysRecipe::Ascending(bits), + RadixRankKeysRecipe::Descending(bits), + ] { + let registration: serde_json::Value = serde_json::from_str( + &radix_rank_keys_registration(Path::new("/provider"), recipe).unwrap(), + ) + .unwrap(); + let mut request = SpecializationRequest { + operation: registration["operation"].as_str().unwrap().into(), + algorithm: registration["algorithm"].as_str().unwrap().into(), + dtype: Scalar::Signed(64), + block: [4, 4, 4], + items_per_thread: 9, + warp_width: None, + }; + let native = request.native().unwrap(); + assert_eq!( + native, + NativeRequest::RadixRankKeys( + RadixRankKeysRequest::block(Dtype::I64, recipe, [4, 4, 4], 9).unwrap() + ) + ); + assert_eq!( + arguments(native), + vec![ + Argument::ConstPointer(Scalar::Signed(64)), + Argument::MutPointer(Scalar::Signed(32)), + Argument::ConstPointer(Scalar::Unsigned(8)), + Argument::MutPointer(Scalar::Unsigned(8)), + Argument::Value(Scalar::Signed(64)), + Argument::Value(Scalar::Signed(64)), + ] + ); + assert!(!native.mutates_payload() && !native.scratch_free()); + request.warp_width = Some(32); + assert!(request.native().is_err()); + request.warp_width = None; + request.dtype = Scalar::Float(64); + assert!(request.native().is_err()); + } + } + for bits in [0, 9, u32::MAX] { + for recipe in [ + RadixRankKeysRecipe::Ascending(bits), + RadixRankKeysRecipe::Descending(bits), + ] { + assert!(radix_rank_keys_registration(Path::new("/provider"), recipe).is_err()); + assert!(RadixRankKeysRequest::block(Dtype::I32, recipe, [32, 1, 1], 1).is_err()); + } + } + assert!( + radix_rank_keys_registration(Path::new("relative"), RadixRankKeysRecipe::Ascending(4)) + .is_err() + ); + for (operation, algorithm) in [ + ("radix_rank_keys_ascending_0_validity_bits_v1", "block"), + ("radix_rank_keys_descending_9_validity_bits_v1", "block"), + ("radix_rank_keys_ascending_04_validity_bits_v1", "block"), + ("radix_rank_keys_ascending_4_validity_bits_v2", "block"), + ("radix_rank_keys_ascending_4_validity_bits_v1", "warp"), + ] { + assert!(RadixRankKeysRecipe::parse(operation, algorithm).is_err()); + } + } + + #[test] + fn radix_sort_keys_registration_preserves_wide_controls_and_scope() { + for (recipe, dtype) in [ + (RadixSortKeysRecipe::Ascending, Scalar::Signed(32)), + (RadixSortKeysRecipe::Descending, Scalar::Unsigned(32)), + (RadixSortKeysRecipe::Ascending, Scalar::Signed(64)), + (RadixSortKeysRecipe::Descending, Scalar::Unsigned(64)), + ] { + let registration: serde_json::Value = serde_json::from_str( + &radix_sort_keys_registration(Path::new("/provider"), recipe).unwrap(), + ) + .unwrap(); + let mut request = SpecializationRequest { + operation: registration["operation"].as_str().unwrap().into(), + algorithm: registration["algorithm"].as_str().unwrap().into(), + dtype, + block: [4, 4, 4], + items_per_thread: 9, + warp_width: None, + }; + let native = request.native().unwrap(); + assert_eq!( + native, + NativeRequest::RadixSortKeys( + RadixSortKeysRequest::block(dtype.dtype().unwrap(), recipe, [4, 4, 4], 9) + .unwrap() + ) + ); + assert_eq!( + arguments(native), + vec![ + Argument::ConstPointer(dtype), + Argument::MutPointer(dtype), + Argument::ConstPointer(Scalar::Unsigned(8)), + Argument::MutPointer(Scalar::Unsigned(8)), + Argument::Value(Scalar::Signed(64)), + Argument::Value(Scalar::Signed(64)), + ] + ); + assert!(!native.mutates_payload() && !native.scratch_free()); + request.warp_width = Some(32); + assert!(request.native().is_err()); + request.warp_width = None; + request.dtype = Scalar::Float(64); + assert!(request.native().is_err()); + } + assert!( + radix_sort_keys_registration(Path::new("relative"), RadixSortKeysRecipe::Ascending) + .is_err() + ); + assert!( + RadixSortKeysRecipe::parse("radix_sort_keys_ascending_validity_bits_v1", "warp") + .is_err() + ); + assert!( + RadixSortKeysRecipe::parse("radix_sort_keys_ascending_validity_bits_v2", "block") + .is_err() + ); + assert!( + RadixSortKeysRecipe::parse("radix_sort_pairs_ascending_validity_bits_v1", "block") + .is_err() + ); + } + #[test] fn histogram_template_decoding_and_independent_counter_abi() { for (algorithm, counter, bins, output_items) in [ diff --git a/rust/cuda-coop-codegen/src/radix_sort.rs b/rust/cuda-coop-codegen/src/radix_sort.rs new file mode 100644 index 000000000000..45b35900333b --- /dev/null +++ b/rust/cuda-coop-codegen/src/radix_sort.rs @@ -0,0 +1,671 @@ +// Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +//! Validated Block radix sorting and ranking with preserved input keys. + +use crate::{Dtype, Error, Result}; + +/// Stable ordering of the selected radix bits in a complete Block tile. +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub enum RadixSortKeysRecipe { + /// Sort selected digits in ascending order. + Ascending, + /// Sort selected digits in descending order. + Descending, +} + +impl RadixSortKeysRecipe { + pub(crate) fn token(self) -> &'static str { + match self { + Self::Ascending => "radix_sort_keys_ascending_validity_bits_v1", + Self::Descending => "radix_sort_keys_descending_validity_bits_v1", + } + } + + pub(crate) fn parse(operation: &str, algorithm: &str) -> Result { + if algorithm != "block" { + return Err(Error::message("Radix Sort Keys requires a Block recipe")); + } + match operation { + "radix_sort_keys_ascending_validity_bits_v1" => Ok(Self::Ascending), + "radix_sort_keys_descending_validity_bits_v1" => Ok(Self::Descending), + _ => Err(Error::message("unsupported Radix Sort Keys operation")), + } + } +} + +/// Validated integer Block radix sort with fresh, fully initialized output. +/// +/// Native calls use `(const T* input, T* output, u8* scratch, +/// const u8* input_bits, u8* output_bits, i64 begin_bit, i64 end_bit)`. +/// The two scalar controls are uniform across the Block and must satisfy +/// `0 <= begin_bit < end_bit <= sizeof(T)*8`. Invalid bounds trap while still +/// signed 64-bit values, before narrowing or calling CUB. Bounds do not enter +/// specialization identity; the caller resolves defaults to zero and key width. +/// +/// Each lane supplies `items` initialized keys and `ceil(items/8)` packed, +/// low-bit-first validity bytes. Every input bit is checked before any numeric +/// read. Input values and validity bytes are preserved. Output has the same +/// capacities, is disjoint from input and scratch, and receives every key and +/// validity bit; unused high bits are zero. +/// +/// Input and output are blocked: item `i` of x-fastest rank `r` has tile +/// position `r*items+i`. Equal selected digits retain their input order in +/// either direction. CUB transforms signed keys before selecting the bit range +/// and restores their original representation after sorting. +/// +/// Every member of the exact Block participates in converged control flow. +/// All pointers remain live with their natural alignment and stated capacities. +/// Scratch uses the complete queried CUB TempStorage. The caller owns its +/// alignment and lifetime and synchronizes Block members before reuse; the +/// provider adds no trailing completion barrier. +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub struct RadixSortKeysRequest { + pub(crate) recipe: RadixSortKeysRecipe, + pub(crate) dtype: Dtype, + block: [u32; 3], + items: u32, +} + +impl RadixSortKeysRequest { + /// Validate radix sorting for an exact complete XYZ Block. + /// + /// `dtype` is `i32`, `u32`, `i64` or `u64`; `recipe` selects ordering. + /// `block` contains the three positive axes and `items` is the fixed + /// per-lane input and output extent. Warp and partial tiles are unsupported. + /// + /// # Errors + /// Rejects unsupported key types, zero/overflowing axes, more than 1024 + /// threads, zero items, or a tile exceeding CUB's 65,535-key counter bound. + pub fn block( + dtype: Dtype, + recipe: RadixSortKeysRecipe, + block: [u32; 3], + items: u32, + ) -> Result { + if !matches!(dtype, Dtype::I32 | Dtype::U32 | Dtype::I64 | Dtype::U64) { + return Err(Error::message( + "Radix Sort Keys requires i32/u32/i64/u64 keys", + )); + } + let threads = block + .into_iter() + .try_fold(1_u32, u32::checked_mul) + .filter(|n| (1..=1024).contains(n)) + .ok_or_else(|| { + Error::message("Radix Sort Keys requires 1..=1024 block threads with positive axes") + })?; + if items == 0 || items > u16::MAX as u32 / threads { + return Err(Error::message( + "Radix Sort Keys requires a positive tile of at most 65535 keys", + )); + } + Ok(Self { + recipe, + dtype, + block, + items, + }) + } + + pub(crate) fn block_shape(self) -> [u32; 3] { + self.block + } + + pub(crate) fn source(self, symbol: &str) -> String { + let (dtype, width) = self.dtype.native(); + let bits = width * 8; + let [x, y, z] = self.block; + let items = self.items; + let bytes = items / 8 + u32::from(items % 8 != 0); + let method = match self.recipe { + RadixSortKeysRecipe::Ascending => "Sort", + RadixSortKeysRecipe::Descending => "SortDescending", + }; + format!( + r#" +#include +#include +#include +#include +namespace {symbol}_detail {{ +using value_type = {dtype}; +using primitive = cub::BlockRadixSort; +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, + long long begin_bit, long long end_bit) {{ + using namespace {symbol}_detail; + if (begin_bit < 0 || begin_bit >= end_bit || end_bit > {bits}) __trap(); +#pragma unroll + for (int i = 0; i < {bytes}; ++i) output_bits[i] = 0; +#pragma unroll + for (int i = 0; i < {items}; ++i) + if ((input_bits[i / 8] & (1u << (i % 8))) == 0) __trap(); + value_type local_keys[{items}]; +#pragma unroll + for (int i = 0; i < {items}; ++i) local_keys[i] = input[i]; + primitive operation{{*reinterpret_cast(scratch)}}; + operation.{method}(local_keys, static_cast(begin_bit), static_cast(end_bit)); +#pragma unroll + for (int i = 0; i < {items}; ++i) {{ + output[i] = local_keys[i]; + output_bits[i / 8] |= static_cast(1u << (i % 8)); + }} +}} +"# + ) + } +} + +/// Stable ordering of paired keys and values over the selected key bits. +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub enum RadixSortPairsRecipe { + /// Sort selected key digits in ascending order, carrying their values. + Ascending, + /// Sort selected key digits in descending order, carrying their values. + Descending, +} + +impl RadixSortPairsRecipe { + pub(crate) fn token(self, value: Dtype) -> String { + let order = match self { + Self::Ascending => "ascending", + Self::Descending => "descending", + }; + let value = match value { + Dtype::I8 => "i8", + Dtype::U8 => "u8", + Dtype::I16 => "i16", + Dtype::U16 => "u16", + Dtype::I32 => "i32", + Dtype::U32 => "u32", + Dtype::I64 => "i64", + Dtype::U64 => "u64", + Dtype::F32 => "f32", + Dtype::F64 => "f64", + }; + format!("radix_sort_pairs_{order}_{value}_validity_bits_v1") + } + + pub(crate) fn parse(operation: &str, algorithm: &str) -> Result<(Self, Dtype)> { + if algorithm != "block" { + return Err(Error::message("Radix Sort Pairs requires a Block recipe")); + } + let unsupported = || Error::message("unsupported Radix Sort Pairs operation"); + let operation = operation + .strip_prefix("radix_sort_pairs_") + .and_then(|token| token.strip_suffix("_validity_bits_v1")) + .ok_or_else(unsupported)?; + let (recipe, value) = if let Some(value) = operation.strip_prefix("ascending_") { + (Self::Ascending, value) + } else if let Some(value) = operation.strip_prefix("descending_") { + (Self::Descending, value) + } else { + return Err(unsupported()); + }; + let value = match value { + "i8" => Dtype::I8, + "u8" => Dtype::U8, + "i16" => Dtype::I16, + "u16" => Dtype::U16, + "i32" => Dtype::I32, + "u32" => Dtype::U32, + "i64" => Dtype::I64, + "u64" => Dtype::U64, + "f32" => Dtype::F32, + "f64" => Dtype::F64, + _ => return Err(unsupported()), + }; + Ok((recipe, value)) + } +} + +/// Validated Block radix sort with preserved inputs and fresh paired outputs. +/// +/// Native calls use `(const K* keys_in, K* keys_out, u8* scratch, +/// const V* values_in, V* values_out, const u8* keys_in_bits, u8* keys_out_bits, +/// const u8* values_in_bits, u8* values_out_bits, i64 begin_bit, i64 end_bit)`. +/// Keys are `i32/u32/i64/u64`; values independently use any [`Dtype`]. Both +/// arrays have the same positive per-lane extent. The two uniform signed64 +/// bounds satisfy `0 <= begin_bit < end_bit <= sizeof(K)*8` before narrowing. +/// Defaults are resolved by the caller, and bounds do not enter identity. +/// +/// Each lane supplies initialized, naturally aligned key and value arrays +/// with separate packed low-bit-first validity maps of `ceil(items/8)` bytes. +/// Every bit in both input maps is checked before reading either numeric +/// array. Sorting uses private copies and preserves both inputs and maps. +/// Output arrays, output validity maps and scratch are mutually disjoint and +/// disjoint from the input arrays and maps. Read-only inputs may alias one another. +/// Every output slot is written and marked valid; unused high bits are zero. +/// +/// Input and output are blocked: item `i` at x-fastest rank `r` has tile +/// position `r*items+i`. CUB transforms signed keys before selecting digits +/// and restores them afterward. Equal selected digits retain input order in +/// both directions, and each value follows its associated key unchanged. +/// +/// All exact Block members participate in converged control flow. Pointers +/// remain live with their declared capacities and natural alignment. Scratch +/// is the complete queried pair primitive TempStorage, including the value +/// exchange; it cannot be inferred from the keys-only layout. The caller +/// synchronizes members before scratch reuse; no trailing barrier is added. +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub struct RadixSortPairsRequest { + pub(crate) recipe: RadixSortPairsRecipe, + pub(crate) key_dtype: Dtype, + pub(crate) value_dtype: Dtype, + block: [u32; 3], + items: u32, +} + +impl RadixSortPairsRequest { + /// Validate independently typed key/value sorting for an exact XYZ Block. + /// + /// `key_dtype` is one of the four integer key types; `value_dtype` is any + /// supported numeric type. `recipe` selects stable ordering. `items` is + /// the common positive per-lane extent of both inputs and both outputs. + /// + /// # Errors + /// Rejects key types, shapes and tile capacities rejected by + /// [`RadixSortKeysRequest::block`], including `Block threads * items > 65535`. + pub fn block( + key_dtype: Dtype, + value_dtype: Dtype, + recipe: RadixSortPairsRecipe, + block: [u32; 3], + items: u32, + ) -> Result { + let order = match recipe { + RadixSortPairsRecipe::Ascending => RadixSortKeysRecipe::Ascending, + RadixSortPairsRecipe::Descending => RadixSortKeysRecipe::Descending, + }; + RadixSortKeysRequest::block(key_dtype, order, block, items)?; + Ok(Self { + recipe, + key_dtype, + value_dtype, + block, + items, + }) + } + + pub(crate) fn block_shape(self) -> [u32; 3] { + self.block + } + + pub(crate) fn source(self, symbol: &str) -> String { + let (key, width) = self.key_dtype.native(); + let (value, _) = self.value_dtype.native(); + let bits = width * 8; + let [x, y, z] = self.block; + let items = self.items; + let bytes = items / 8 + u32::from(items % 8 != 0); + let method = match self.recipe { + RadixSortPairsRecipe::Ascending => "Sort", + RadixSortPairsRecipe::Descending => "SortDescending", + }; + format!( + r#" +#include +#include +#include +namespace {symbol}_detail {{ +using key_type = {key}; +using value_type = {value}; +using primitive = cub::BlockRadixSort; +using scratch_storage = primitive::TempStorage; +}} +extern "C" __device__ void {symbol}(const {key}* keys_in, {key}* keys_out, + unsigned char* scratch, const {value}* values_in, {value}* values_out, + const unsigned char* keys_in_bits, unsigned char* keys_out_bits, + const unsigned char* values_in_bits, unsigned char* values_out_bits, + long long begin_bit, long long end_bit) {{ + using namespace {symbol}_detail; + if (begin_bit < 0 || begin_bit >= end_bit || end_bit > {bits}) __trap(); +#pragma unroll + for (int i = 0; i < {bytes}; ++i) {{ + keys_out_bits[i] = 0; + values_out_bits[i] = 0; + }} +#pragma unroll + for (int i = 0; i < {items}; ++i) {{ + if ((keys_in_bits[i / 8] & (1u << (i % 8))) == 0) __trap(); + if ((values_in_bits[i / 8] & (1u << (i % 8))) == 0) __trap(); + }} + key_type local_keys[{items}]; + value_type local_values[{items}]; +#pragma unroll + for (int i = 0; i < {items}; ++i) {{ + local_keys[i] = keys_in[i]; + local_values[i] = values_in[i]; + }} + primitive operation{{*reinterpret_cast(scratch)}}; + operation.{method}(local_keys, local_values, static_cast(begin_bit), static_cast(end_bit)); +#pragma unroll + for (int i = 0; i < {items}; ++i) {{ + keys_out[i] = local_keys[i]; + values_out[i] = local_values[i]; + keys_out_bits[i / 8] |= static_cast(1u << (i % 8)); + values_out_bits[i / 8] |= static_cast(1u << (i % 8)); + }} +}} +"# + ) + } +} + +/// Stable ordering and static digit width for Block radix ranks. +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub enum RadixRankKeysRecipe { + /// Rank selected digits in ascending order with a width of 1 through 8 bits. + Ascending(u32), + /// Rank selected digits in descending order with a width of 1 through 8 bits. + Descending(u32), +} + +impl RadixRankKeysRecipe { + fn bits(self) -> u32 { + match self { + Self::Ascending(bits) | Self::Descending(bits) => bits, + } + } + + pub(crate) fn validate(self) -> Result<()> { + if !(1..=8).contains(&self.bits()) { + return Err(Error::message("Radix Rank Keys requires 1..=8 radix bits")); + } + Ok(()) + } + + pub(crate) fn token(self) -> String { + let order = match self { + Self::Ascending(_) => "ascending", + Self::Descending(_) => "descending", + }; + format!("radix_rank_keys_{order}_{}_validity_bits_v1", self.bits()) + } + + pub(crate) fn parse(operation: &str, algorithm: &str) -> Result { + if algorithm != "block" { + return Err(Error::message("Radix Rank Keys requires a Block recipe")); + } + let unsupported = || Error::message("unsupported Radix Rank Keys operation"); + let operation = operation + .strip_prefix("radix_rank_keys_") + .and_then(|token| token.strip_suffix("_validity_bits_v1")) + .ok_or_else(unsupported)?; + let (descending, bits) = if let Some(bits) = operation.strip_prefix("ascending_") { + (false, bits) + } else if let Some(bits) = operation.strip_prefix("descending_") { + (true, bits) + } else { + return Err(unsupported()); + }; + let bits = match bits { + "1" => 1, + "2" => 2, + "3" => 3, + "4" => 4, + "5" => 5, + "6" => 6, + "7" => 7, + "8" => 8, + _ => return Err(unsupported()), + }; + Ok(if descending { + Self::Descending(bits) + } else { + Self::Ascending(bits) + }) + } +} + +/// Validated integer Block ranking with fresh i32 ranks in input-slot order. +/// +/// Native calls use `(const T* keys, i32* ranks, u8* scratch, +/// const u8* key_bits, u8* rank_bits, i64 begin_bit, i64 end_constraint)`. +/// Both scalar controls are uniform across the Block. The static digit width +/// is 1 through 8; `begin_bit` must be in `0..=sizeof(T)*8-width`. +/// `end_constraint` is -1 for an omitted end, or an explicit static endpoint +/// equal to `begin_bit+width` and at most the key width. Checks use signed64 +/// arithmetic before narrowing, reading keys or calling CUB. +/// +/// Each lane supplies `items` initialized keys and `ceil(items/8)` packed, +/// low-bit-first validity bytes. Every input bit is checked before any numeric +/// read. Keys and their validity map remain unchanged. Output contains `items` +/// naturally aligned i32 ranks with a separate packed validity map, independent +/// of key width/alignment. All output slots become valid; padding bits are zero. +/// Output arrays, output maps and scratch are mutually disjoint and disjoint +/// from the input arrays and maps, which remain live throughout the call. +/// +/// Ranks describe stable selected-digit order in the complete blocked tile: +/// input slot `i` at x-fastest rank `r` has position `r*items+i`. Signed keys +/// are transformed by flipping their sign bit before digit extraction. Equal +/// digits retain input order in either direction. The ranks form a permutation +/// of `0..Block_threads*items`, stored in original input-slot order. +/// +/// Every member of the exact Block participates in converged control flow. +/// Scratch is the queried BlockRadixRank TempStorage, not sort/exchange storage. +/// The caller supplies its alignment and lifetime and synchronizes members +/// before reuse. No trailing completion barrier is added. +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub struct RadixRankKeysRequest { + pub(crate) recipe: RadixRankKeysRecipe, + pub(crate) dtype: Dtype, + block: [u32; 3], + items: u32, +} + +impl RadixRankKeysRequest { + /// Validate stable ranking for an exact complete XYZ Block. + /// + /// `dtype` is i32/u32/i64/u64; `recipe` selects ordering and static width. + /// `block` contains positive XYZ axes and `items` is the positive per-lane + /// extent of keys and ranks. Runtime begin/end controls do not enter identity. + /// + /// # Errors + /// Rejects widths outside 1..=8 and the key types, shapes and tile capacities + /// rejected by [`RadixSortKeysRequest::block`], including tiles above 65535 keys. + pub fn block( + dtype: Dtype, + recipe: RadixRankKeysRecipe, + block: [u32; 3], + items: u32, + ) -> Result { + recipe.validate()?; + RadixSortKeysRequest::block(dtype, RadixSortKeysRecipe::Ascending, block, items)?; + Ok(Self { + recipe, + dtype, + block, + items, + }) + } + + pub(crate) fn block_shape(self) -> [u32; 3] { + self.block + } + + pub(crate) fn source(self, symbol: &str) -> String { + let (dtype, width) = self.dtype.native(); + let bits = width * 8; + let unsigned = if width == 4 { + "unsigned int" + } else { + "unsigned long long" + }; + let transform = if matches!(self.dtype, Dtype::I32 | Dtype::I64) { + format!(" ^ (unsigned_type{{1}} << {})", bits - 1) + } else { + String::new() + }; + let [x, y, z] = self.block; + let items = self.items; + let bytes = items / 8 + u32::from(items % 8 != 0); + let radix_bits = self.recipe.bits(); + let descending = matches!(self.recipe, RadixRankKeysRecipe::Descending(_)); + format!( + r#" +#include +#include +#include +#include +namespace {symbol}_detail {{ +using unsigned_type = {unsigned}; +using primitive = cub::BlockRadixRank<{x}, {radix_bits}, {descending}, true, + cub::BLOCK_SCAN_WARP_SCANS, cudaSharedMemBankSizeFourByte, {y}, {z}>; +using scratch_storage = primitive::TempStorage; +}} +extern "C" __device__ void {symbol}(const {dtype}* keys, int* ranks, + unsigned char* scratch, const unsigned char* key_bits, unsigned char* rank_bits, + long long begin_bit, long long end_constraint) {{ + using namespace {symbol}_detail; + if (begin_bit < 0 || begin_bit > {bits} - {radix_bits} || + (end_constraint != -1 && + (end_constraint < {radix_bits} || end_constraint > {bits} || + begin_bit != end_constraint - {radix_bits}))) __trap(); +#pragma unroll + for (int i = 0; i < {items}; ++i) + if ((key_bits[i / 8] & (1u << (i % 8))) == 0) __trap(); + unsigned_type local_keys[{items}]; + int local_ranks[{items}]; +#pragma unroll + for (int i = 0; i < {items}; ++i) + local_keys[i] = static_cast(keys[i]){transform}; + primitive operation{{*reinterpret_cast(scratch)}}; + operation.RankKeys(local_keys, local_ranks, + cub::BFEDigitExtractor(static_cast(begin_bit), {radix_bits})); +#pragma unroll + for (int i = 0; i < {bytes}; ++i) rank_bits[i] = 0; +#pragma unroll + for (int i = 0; i < {items}; ++i) {{ + ranks[i] = local_ranks[i]; + rank_bits[i / 8] |= static_cast(1u << (i % 8)); + }} +}} +"# + ) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn rank_key_types_and_tile_capacity() { + for dtype in [Dtype::I32, Dtype::U32, Dtype::I64, Dtype::U64] { + assert!( + RadixRankKeysRequest::block( + dtype, + RadixRankKeysRecipe::Ascending(1), + [3, 5, 17], + 257 + ) + .is_ok() + ); + assert!( + RadixRankKeysRequest::block( + dtype, + RadixRankKeysRecipe::Descending(8), + [256, 1, 1], + 256 + ) + .is_err() + ); + } + for dtype in [ + Dtype::I8, + Dtype::U8, + Dtype::I16, + Dtype::U16, + Dtype::F32, + Dtype::F64, + ] { + assert!( + RadixRankKeysRequest::block( + dtype, + RadixRankKeysRecipe::Ascending(4), + [64, 1, 1], + 4 + ) + .is_err() + ); + } + for (block, items) in [ + ([0, 1, 1], 1), + ([u32::MAX, 2, 2], 1), + ([1025, 1, 1], 1), + ([64, 1, 1], 0), + ] { + assert!( + RadixRankKeysRequest::block( + Dtype::I64, + RadixRankKeysRecipe::Ascending(4), + block, + items + ) + .is_err() + ); + } + } + + #[test] + fn key_types_and_complete_tile_counter_bounds() { + for dtype in [Dtype::I32, Dtype::U32, Dtype::I64, Dtype::U64] { + for (block, items) in [ + ([1, 1, 1], 65535), + ([3, 5, 17], 257), + ([4, 4, 4], 4), + ([33, 1, 1], 9), + ] { + assert!( + RadixSortKeysRequest::block( + dtype, + RadixSortKeysRecipe::Ascending, + block, + items + ) + .is_ok() + ); + } + } + for dtype in [ + Dtype::I8, + Dtype::U8, + Dtype::I16, + Dtype::U16, + Dtype::F32, + Dtype::F64, + ] { + assert!( + RadixSortKeysRequest::block(dtype, RadixSortKeysRecipe::Ascending, [64, 1, 1], 4) + .is_err() + ); + } + for (block, items) in [ + ([0, 1, 1], 1), + ([1, 0, 1], 1), + ([1, 1, 0], 1), + ([u32::MAX, 2, 2], 1), + ([1025, 1, 1], 1), + ([1, 1, 1], 0), + ([1, 1, 1], 65536), + ([256, 1, 1], 256), + ([64, 1, 1], u32::MAX), + ] { + assert!( + RadixSortKeysRequest::block( + Dtype::I64, + RadixSortKeysRecipe::Descending, + block, + items + ) + .is_err() + ); + } + } +} diff --git a/rust/cuda-coop-codegen/tests/native.rs b/rust/cuda-coop-codegen/tests/native.rs index 3cb7cb988141..e10f80c5bbe1 100644 --- a/rust/cuda-coop-codegen/tests/native.rs +++ b/rust/cuda-coop-codegen/tests/native.rs @@ -6,13 +6,243 @@ use cuda_coop_codegen::{ AdjacentDifferenceRecipe, AdjacentDifferenceRequest, Algorithm, DiscontinuityRecipe, DiscontinuityRequest, Dtype, ExchangeMode, ExchangeRecipe, ExchangeRequest, HistogramAlgorithm, - HistogramRecipe, HistogramRequest, Layout, NativeRequest, Operation, Provider, ReduceAlgorithm, - ReduceBatchedLayout, ReduceBatchedRecipe, ReduceBatchedRequest, ReduceForm, ReduceOperator, - ReduceRequest, Request, ScanAlgorithm, ScanOperation, ScanOperator, ScanRequest, ShuffleRecipe, - ShuffleRequest, + HistogramRecipe, HistogramRequest, Layout, NativeRequest, Operation, Provider, + RadixRankKeysRecipe, RadixRankKeysRequest, RadixSortKeysRecipe, RadixSortKeysRequest, + RadixSortPairsRecipe, RadixSortPairsRequest, ReduceAlgorithm, ReduceBatchedLayout, + ReduceBatchedRecipe, ReduceBatchedRequest, ReduceForm, ReduceOperator, ReduceRequest, Request, + ScanAlgorithm, ScanOperation, ScanOperator, ScanRequest, ShuffleRecipe, ShuffleRequest, }; use std::{env, path::Path, time::Instant}; +#[test] +#[ignore = "requires CUDA_HOME and a Linux CUDA toolkit with NVRTC LTO support"] +fn native_mixed_radix_rank_keys_recipes() -> Result<(), Box> { + use RadixRankKeysRecipe::{Ascending, Descending}; + 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)?; + // The public XYZ example, then width limits, bitmap tails and unsigned keys. + let rows = [ + (Dtype::I64, [4, 4, 4], 1, Ascending(3)), + (Dtype::I64, [4, 4, 4], 1, Descending(3)), + (Dtype::I64, [4, 4, 4], 4, Ascending(3)), + (Dtype::I64, [4, 4, 4], 4, Descending(3)), + (Dtype::I32, [33, 1, 1], 9, Ascending(1)), + (Dtype::U32, [1, 1, 1], 1, Descending(8)), + (Dtype::U64, [64, 1, 1], 4, Ascending(4)), + ]; + let mut requests = Vec::new(); + for (dtype, block, items, recipe) in rows { + requests.push(NativeRequest::RadixRankKeys(RadixRankKeysRequest::block( + dtype, recipe, block, items, + )?)); + } + requests.extend([ + NativeRequest::RadixSortKeys(RadixSortKeysRequest::block( + Dtype::I64, + RadixSortKeysRecipe::Ascending, + [4, 4, 4], + 4, + )?), + NativeRequest::RadixSortPairs(RadixSortPairsRequest::block( + Dtype::I32, + Dtype::F64, + RadixSortPairsRecipe::Descending, + [4, 4, 4], + 4, + )?), + ]); + let artifact = provider.compile_native(&requests, &target)?; + assert_eq!(artifact.specializations.len(), 9); + 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)); + assert!(!entry.mutates_payload && !entry.scratch_free); + assert_eq!(entry.scratch, entry.cub_layout); + assert!(entry.scratch.size > 0 && entry.scratch.alignment.is_power_of_two()); + eprintln!( + "{request:?}: CUB {:?}, scratch {:?}", + entry.cub_layout, entry.scratch + ); + } + // Rank TempStorage depends on geometry/width, not key count or direction. + for entry in &artifact.specializations[1..4] { + assert_eq!(entry.scratch, artifact.specializations[0].scratch); + } + eprintln!( + "{} LTO bytes, {} specializations", + artifact.ltoir.len(), + artifact.specializations.len() + ); + Ok(()) +} + +#[test] +#[ignore = "requires CUDA_HOME and a Linux CUDA toolkit with NVRTC LTO support"] +fn native_mixed_radix_sort_pairs_recipes() -> Result<(), Box> { + use RadixSortPairsRecipe::{Ascending, Descending}; + 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)?; + // Independent key/value widths, both directions, XYZ and bitmap tails. + // Only the first pair repeats types and geometry to distinguish order. + let rows = [ + (Dtype::I32, Dtype::F64, [4, 4, 4], 4, Ascending), + (Dtype::I32, Dtype::F64, [4, 4, 4], 4, Descending), + (Dtype::U32, Dtype::I8, [33, 1, 1], 9, Descending), + (Dtype::I64, Dtype::U8, [32, 1, 1], 1, Ascending), + (Dtype::U64, Dtype::I16, [64, 1, 1], 4, Descending), + (Dtype::I32, Dtype::U16, [4, 4, 4], 4, Ascending), + (Dtype::U32, Dtype::I32, [1, 1, 1], 1, Descending), + (Dtype::I64, Dtype::U32, [33, 1, 1], 4, Ascending), + (Dtype::U64, Dtype::I64, [32, 1, 1], 4, Descending), + (Dtype::I32, Dtype::U64, [64, 1, 1], 4, Ascending), + (Dtype::U32, Dtype::F32, [4, 4, 4], 4, Descending), + ]; + let mut requests = Vec::new(); + for (key, value, block, items, recipe) in rows { + requests.push(NativeRequest::RadixSortPairs(RadixSortPairsRequest::block( + key, value, recipe, block, items, + )?)); + } + requests.extend([ + NativeRequest::RadixSortKeys(RadixSortKeysRequest::block( + Dtype::I32, + RadixSortKeysRecipe::Ascending, + [4, 4, 4], + 4, + )?), + NativeRequest::Discontinuity(DiscontinuityRequest::block( + Dtype::F64, + DiscontinuityRecipe::HeadsAndTailsBoth, + [4, 4, 4], + 4, + )?), + ]); + let artifact = provider.compile_native(&requests, &target)?; + assert_eq!(artifact.specializations.len(), 13); + 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)); + assert!(!entry.mutates_payload && !entry.scratch_free); + assert_eq!(entry.scratch, entry.cub_layout); + assert!(entry.scratch.size > 0 && entry.scratch.alignment.is_power_of_two()); + eprintln!( + "{request:?}: CUB {:?}, scratch {:?}", + entry.cub_layout, entry.scratch + ); + } + assert_eq!( + artifact.specializations[0].scratch, + artifact.specializations[1].scratch + ); + // Pair scratch includes the value exchange, even if the ranker dominates. + assert!(artifact.specializations[0].scratch.size >= artifact.specializations[11].scratch.size); + assert!(artifact.specializations[0].scratch.alignment >= std::mem::align_of::()); + eprintln!( + "{} LTO bytes, {} specializations", + artifact.ltoir.len(), + artifact.specializations.len() + ); + Ok(()) +} + +#[test] +#[ignore = "requires CUDA_HOME and a Linux CUDA toolkit with NVRTC LTO support"] +fn native_mixed_radix_sort_keys_recipes() -> Result<(), Box> { + use RadixSortKeysRecipe::{Ascending, Descending}; + 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)?; + // Four key types, both directions, XYZ, non-warp geometry and bitmap tails. + // The last two rows distinguish order and exact shape at the same capacity. + let rows = [ + (Dtype::I64, [4, 4, 4], 1, Ascending), + (Dtype::I64, [4, 4, 4], 4, Descending), + (Dtype::I32, [33, 1, 1], 9, Ascending), + (Dtype::U32, [1, 1, 1], 1, Descending), + (Dtype::U64, [64, 1, 1], 4, Ascending), + (Dtype::I64, [4, 4, 4], 4, Ascending), + (Dtype::I64, [64, 1, 1], 4, Ascending), + ]; + let mut requests = Vec::new(); + for (dtype, block, items, recipe) in rows { + requests.push(NativeRequest::RadixSortKeys(RadixSortKeysRequest::block( + dtype, recipe, block, items, + )?)); + } + requests.extend([ + NativeRequest::LoadStore(Request::new_typed( + Dtype::I32, + Operation::LoadFullValidity, + Algorithm::Direct, + 64, + 4, + )?), + NativeRequest::AdjacentDifference(AdjacentDifferenceRequest::block( + Dtype::I32, + AdjacentDifferenceRecipe::LeftPartialPredecessor, + [64, 1, 1], + 4, + )?), + NativeRequest::Histogram(HistogramRequest::block( + Dtype::U8, + HistogramRecipe::new(HistogramAlgorithm::Sort, Dtype::U64, 65, 2)?, + [64, 1, 1], + 4, + )?), + ]); + let artifact = provider.compile_native(&requests, &target)?; + assert_eq!(artifact.specializations.len(), 10); + assert!(!artifact.ltoir.is_empty()); + let mut identities = std::collections::HashSet::new(); + let mut symbols = std::collections::HashSet::new(); + for (index, (request, entry)) in requests.iter().zip(&artifact.specializations).enumerate() { + assert!(identities.insert(entry.identity)); + assert!(symbols.insert(&entry.symbol)); + assert!(!entry.mutates_payload); + if index < rows.len() { + assert_eq!(entry.scratch, entry.cub_layout); + assert!(entry.scratch.size > 0 && entry.scratch.alignment.is_power_of_two()); + assert!(!entry.scratch_free); + } + eprintln!( + "{request:?}: CUB {:?}, scratch {:?}", + entry.cub_layout, entry.scratch + ); + } + assert_eq!( + artifact.specializations[1].scratch, + artifact.specializations[5].scratch + ); + assert_eq!( + artifact.specializations[5].scratch, + artifact.specializations[6].scratch + ); + assert!(artifact.specializations[7].scratch_free); + assert_eq!(artifact.specializations[7].scratch.size, 0); + assert_eq!( + artifact.specializations[8].scratch, + artifact.specializations[8].cub_layout + ); + assert!(artifact.specializations[9].scratch.size > artifact.specializations[9].cub_layout.size); + eprintln!( + "{} LTO bytes, {} specializations", + artifact.ltoir.len(), + artifact.specializations.len() + ); + Ok(()) +} + #[test] #[ignore = "requires CUDA_HOME and a Linux CUDA toolkit with NVRTC LTO support"] fn native_mixed_histogram_recipes() -> Result<(), Box> { diff --git a/rust/cuda-coop/build.rs b/rust/cuda-coop/build.rs index 650ef7dfe1f7..ed72ea69ee37 100644 --- a/rust/cuda-coop/build.rs +++ b/rust/cuda-coop/build.rs @@ -3,10 +3,10 @@ //! Register the owning host provider for native cooperative operations. use cuda_coop_codegen::{ - AdjacentDifferenceRecipe, Algorithm, DiscontinuityRecipe, ExchangeMode, ExchangeRecipe, - HistogramAlgorithm, Operation, ReduceAlgorithm, ReduceBatchedLayout, ReduceBatchedRecipe, - ReduceForm, ReduceOperator, ReduceRecipe, ScanAlgorithm, ScanOperation, ScanOperator, - ScanRecipe, ShuffleRecipe, oxide, + AdjacentDifferenceRecipe, Algorithm, DiscontinuityRecipe, Dtype, ExchangeMode, ExchangeRecipe, + HistogramAlgorithm, Operation, RadixRankKeysRecipe, RadixSortKeysRecipe, RadixSortPairsRecipe, + ReduceAlgorithm, ReduceBatchedLayout, ReduceBatchedRecipe, ReduceForm, ReduceOperator, + ReduceRecipe, ScanAlgorithm, ScanOperation, ScanOperator, ScanRecipe, ShuffleRecipe, oxide, }; fn main() { @@ -25,6 +25,45 @@ fn main() { let callable = oxide::reduce_callable_registration(&executable) .expect("cuda.coop callable Reduce registration"); println!("cargo::rustc-env=COOP_REDUCE_RAKING_CALLABLE_PAYLOAD={callable}"); + for bits in 1..=8 { + for (recipe, order) in [ + (RadixRankKeysRecipe::Ascending(bits), "ASCENDING"), + (RadixRankKeysRecipe::Descending(bits), "DESCENDING"), + ] { + let value = oxide::radix_rank_keys_registration(&executable, recipe) + .expect("cuda.coop radix rank keys registration"); + println!("cargo::rustc-env=COOP_RADIX_RANK_KEYS_{order}_{bits}={value}"); + } + } + for (recipe, name) in [ + (RadixSortKeysRecipe::Ascending, "ASCENDING"), + (RadixSortKeysRecipe::Descending, "DESCENDING"), + ] { + let value = oxide::radix_sort_keys_registration(&executable, recipe) + .expect("cuda.coop radix sort keys registration"); + println!("cargo::rustc-env=COOP_RADIX_SORT_KEYS_{name}={value}"); + } + for (recipe, order) in [ + (RadixSortPairsRecipe::Ascending, "ASCENDING"), + (RadixSortPairsRecipe::Descending, "DESCENDING"), + ] { + for (dtype, name) in [ + (Dtype::I8, "I8"), + (Dtype::U8, "U8"), + (Dtype::I16, "I16"), + (Dtype::U16, "U16"), + (Dtype::I32, "I32"), + (Dtype::U32, "U32"), + (Dtype::I64, "I64"), + (Dtype::U64, "U64"), + (Dtype::F32, "F32"), + (Dtype::F64, "F64"), + ] { + let value = oxide::radix_sort_pairs_registration(&executable, recipe, dtype) + .expect("cuda.coop radix sort pairs registration"); + println!("cargo::rustc-env=COOP_RADIX_SORT_PAIRS_{order}_{name}={value}"); + } + } for (algorithm, name) in [ (HistogramAlgorithm::Atomic, "ATOMIC"), (HistogramAlgorithm::Sort, "SORT"), diff --git a/rust/cuda-coop/examples/radix_rank_keys/.gitignore b/rust/cuda-coop/examples/radix_rank_keys/.gitignore new file mode 100644 index 000000000000..74bffa8af4a4 --- /dev/null +++ b/rust/cuda-coop/examples/radix_rank_keys/.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/radix_rank_keys/Cargo.lock b/rust/cuda-coop/examples/radix_rank_keys/Cargo.lock new file mode 100644 index 000000000000..eec3d97249e5 --- /dev/null +++ b/rust/cuda-coop/examples/radix_rank_keys/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-radix-rank-keys" +version = "0.0.0" +dependencies = [ + "coop-radix-rank-keys-helper", + "cuda-async", + "cuda-coop", + "cuda-core", + "cuda-device", + "cuda-host", +] + +[[package]] +name = "coop-radix-rank-keys-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/radix_rank_keys/Cargo.toml b/rust/cuda-coop/examples/radix_rank_keys/Cargo.toml new file mode 100644 index 000000000000..f1c265d74e9c --- /dev/null +++ b/rust/cuda-coop/examples/radix_rank_keys/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-radix-rank-keys" +version = "0.0.0" +edition = "2024" +publish = false + +[workspace] +members = ["helper"] + +[dependencies] +cuda-coop = { path = "../.." } +coop-radix-rank-keys-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/radix_rank_keys/helper/Cargo.toml b/rust/cuda-coop/examples/radix_rank_keys/helper/Cargo.toml new file mode 100644 index 000000000000..a078c76c0b44 --- /dev/null +++ b/rust/cuda-coop/examples/radix_rank_keys/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-radix-rank-keys-helper" +version = "0.0.0" +edition = "2024" +publish = false + +[dependencies] +cuda-coop = { path = "../../.." } diff --git a/rust/cuda-coop/examples/radix_rank_keys/helper/src/lib.rs b/rust/cuda-coop/examples/radix_rank_keys/helper/src/lib.rs new file mode 100644 index 000000000000..489029251dee --- /dev/null +++ b/rust/cuda-coop/examples/radix_rank_keys/helper/src/lib.rs @@ -0,0 +1,52 @@ +// Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +//! Infer the key type from Load and return independent i32 rank payloads. +#![no_std] +use cuda_coop as coop; + +/// Ascending ranks, descending ranks, then the preserved key payload. +pub type RankedTile = ( + coop::ThreadData, + coop::ThreadData, + coop::ThreadData, +); + +/// Rank a sign-touching three-bit digit with static and runtime begin options. +/// +/// The static interval derives its width. The runtime interval supplies width +/// explicitly. Both results have natural i32 alignment, independent of input. +pub fn rank_tile( + block: coop::ThreadGroup, + keys: &[T], + offset: usize, + begin: i64, +) -> RankedTile +where + coop::Alignment: coop::SupportedAlignment, +{ + let mut input = coop::ThreadData::<_, ITEMS_PER_THREAD, ALIGN>::new(); + coop::load( + block, + keys, + &mut input, + coop::LoadOptions::default().offset(offset), + ); + let ascending = coop::radix_rank_keys( + block, + &input, + coop::RadixRankOptions::default() + .begin_bit_const::<61>() + .end_bit_const::<64>(), + ); + let descending = coop::radix_rank_keys( + block, + &input, + coop::RadixRankOptions::default() + .begin_bit(begin) + .radix_bits::<3>() + .end_bit_const::<64>() + .order(coop::Descending), + ); + (ascending, descending, input) +} diff --git a/rust/cuda-coop/examples/radix_rank_keys/src/main.rs b/rust/cuda-coop/examples/radix_rank_keys/src/main.rs new file mode 100644 index 000000000000..aa18c93801fd --- /dev/null +++ b/rust/cuda-coop/examples/radix_rank_keys/src/main.rs @@ -0,0 +1,116 @@ +// 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 = 3, block = (4, 4, 4), dynamic_shared = inferred)] + pub fn rank_tile( + keys: &[i64], + mut ascending: coop::DisjointSlice>, + mut descending: coop::DisjointSlice>, + mut preserved: coop::DisjointSlice>, + begin: i64, + ) { + let block = coop::this_block(); + let offset = + block.rank_at::() as usize * block.count() as usize * ITEMS_PER_THREAD; + let (ascending_ranks, descending_ranks, input) = + coop_radix_rank_keys_helper::rank_tile::<_, ITEMS_PER_THREAD, 32>( + block, keys, offset, begin, + ); + coop::store( + block, + &mut ascending, + &ascending_ranks, + coop::StoreOptions::default().offset(offset), + ); + coop::store( + block, + &mut descending, + &descending_ranks, + coop::StoreOptions::default().offset(offset), + ); + coop::store( + block, + &mut preserved, + &input, + coop::StoreOptions::default().offset(offset), + ); + } +} + +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, LaunchConfig3D}; + use std::cmp::Reverse; + + 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 = 64 * $items; + let keys: Vec = (0..2 * tile).map(|i| { + let local = i % tile; + let block = i / tile; + let digit = ((local * 5 + block * 3) % 8) as u64; + // Distinct increasing low tags make stable digit ranking differ + // from whole-key descending ranking on repeated high digits. + ((digit << 61 | i as u64) ^ (1_u64 << 63)) as i64 + }).collect(); + let mut expected_ascending = vec![0_i32; keys.len()]; + let mut expected_descending = vec![0_i32; keys.len()]; + for (block, input) in keys.chunks_exact(tile).enumerate() { + let digit = |index: &usize| ((input[*index] as u64) ^ (1_u64 << 63)) >> 61; + let mut ascending: Vec = (0..tile).collect(); + let mut descending = ascending.clone(); + ascending.sort_by_key(digit); + descending.sort_by_key(|index| Reverse(digit(index))); + let mut full_key = (0..tile).collect::>(); + full_key.sort_by_key(|index| Reverse(input[*index])); + assert_ne!(full_key, descending); + // Rank belongs to the original item, so invert each sorted + // permutation rather than storing the sorted index sequence. + for (rank, index) in ascending.into_iter().enumerate() { + expected_ascending[block * tile + index] = rank as i32; + } + for (rank, index) in descending.into_iter().enumerate() { + expected_descending[block * tile + index] = rank as i32; + } + } + let source = DeviceBuffer::from_host(&stream, &keys)?; + let ascending_target = DeviceBuffer::from_host(&stream, &vec![-1_i32; keys.len()])?; + let descending_target = DeviceBuffer::from_host(&stream, &vec![-1_i32; keys.len()])?; + let preserved_target = DeviceBuffer::from_host(&stream, &vec![17_i64; keys.len()])?; + let plan = module.prepare_rank_tile::<$items>(LaunchConfig3D::auto_shared((2, 1, 1), (4, 4, 4)))?; + let (source, ascending_target, descending_target, preserved_target) = module.rank_tile_async_owned::<$items, _, _, _, _>(&plan, source, ascending_target, descending_target, preserved_target, 61).sync()?; + let mut ascending = vec![0_i32; keys.len()]; + let mut descending = vec![0_i32; keys.len()]; + let mut preserved = vec![0_i64; keys.len()]; + let mut source_after = vec![0_i64; keys.len()]; + ascending_target.copy_to_host(&stream, &mut ascending)?; + descending_target.copy_to_host(&stream, &mut descending)?; + preserved_target.copy_to_host(&stream, &mut preserved)?; + source.copy_to_host(&stream, &mut source_after)?; + stream.synchronize()?; + assert_eq!(ascending, expected_ascending); + assert_eq!(descending, expected_descending); + assert_eq!(preserved, keys); + assert_eq!(source_after, keys); + println!("items/thread={}: {} rank observations, {} preserved keys and {} source keys checked", $items, 2 * keys.len(), keys.len(), keys.len()); + }}; + } + run!(1); + run!(4); + Ok(()) +} diff --git a/rust/cuda-coop/examples/radix_sort_keys/.gitignore b/rust/cuda-coop/examples/radix_sort_keys/.gitignore new file mode 100644 index 000000000000..74bffa8af4a4 --- /dev/null +++ b/rust/cuda-coop/examples/radix_sort_keys/.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/radix_sort_keys/Cargo.lock b/rust/cuda-coop/examples/radix_sort_keys/Cargo.lock new file mode 100644 index 000000000000..6190019cce64 --- /dev/null +++ b/rust/cuda-coop/examples/radix_sort_keys/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-radix-sort-keys" +version = "0.0.0" +dependencies = [ + "coop-radix-sort-keys-helper", + "cuda-async", + "cuda-coop", + "cuda-core", + "cuda-device", + "cuda-host", +] + +[[package]] +name = "coop-radix-sort-keys-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/radix_sort_keys/Cargo.toml b/rust/cuda-coop/examples/radix_sort_keys/Cargo.toml new file mode 100644 index 000000000000..231ae966e378 --- /dev/null +++ b/rust/cuda-coop/examples/radix_sort_keys/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-radix-sort-keys" +version = "0.0.0" +edition = "2024" +publish = false + +[workspace] +members = ["helper"] + +[dependencies] +cuda-coop = { path = "../.." } +coop-radix-sort-keys-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/radix_sort_keys/helper/Cargo.toml b/rust/cuda-coop/examples/radix_sort_keys/helper/Cargo.toml new file mode 100644 index 000000000000..79a6c548e184 --- /dev/null +++ b/rust/cuda-coop/examples/radix_sort_keys/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-radix-sort-keys-helper" +version = "0.0.0" +edition = "2024" +publish = false + +[dependencies] +cuda-coop = { path = "../../.." } diff --git a/rust/cuda-coop/examples/radix_sort_keys/helper/src/lib.rs b/rust/cuda-coop/examples/radix_sort_keys/helper/src/lib.rs new file mode 100644 index 000000000000..876e2b8cdf38 --- /dev/null +++ b/rust/cuda-coop/examples/radix_sort_keys/helper/src/lib.rs @@ -0,0 +1,49 @@ +// Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +//! Infer generic keys from Load, then sort full and selected bit intervals. +#![no_std] +use cuda_coop as coop; + +/// Return three independent sorts and the preserved input payload. +/// +/// The begin-only ascending sort uses the full key width as its end. The +/// descending sort uses both runtime endpoints. Every input slot is loaded. +pub fn sort_tile( + block: coop::ThreadGroup, + input: &[T], + offset: usize, + begin: i64, + end: i64, +) -> ( + coop::ThreadData, + coop::ThreadData, + coop::ThreadData, + coop::ThreadData, +) +where + coop::Alignment: coop::SupportedAlignment, +{ + let mut values = coop::ThreadData::<_, ITEMS_PER_THREAD, ALIGN>::new(); + coop::load( + block, + input, + &mut values, + coop::LoadOptions::default().offset(offset), + ); + let full = coop::radix_sort_keys(block, &values, Default::default()); + let selected = coop::radix_sort_keys( + block, + &values, + coop::RadixSortOptions::default().begin_bit(begin), + ); + let descending = coop::radix_sort_keys( + block, + &values, + coop::RadixSortOptions::default() + .begin_bit(begin) + .end_bit(end) + .order(coop::Descending), + ); + (full, selected, descending, values) +} diff --git a/rust/cuda-coop/examples/radix_sort_keys/src/main.rs b/rust/cuda-coop/examples/radix_sort_keys/src/main.rs new file mode 100644 index 000000000000..35884fb1f905 --- /dev/null +++ b/rust/cuda-coop/examples/radix_sort_keys/src/main.rs @@ -0,0 +1,151 @@ +// 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 = 3, block = (4, 4, 4), dynamic_shared = inferred)] + pub fn sort_tile( + input: &[i64], + mut full: coop::DisjointSlice>, + mut selected: coop::DisjointSlice>, + mut descending: coop::DisjointSlice>, + mut preserved: coop::DisjointSlice>, + begin: i64, + end: i64, + ) { + let block = coop::this_block(); + let offset = + block.rank_at::() as usize * block.count() as usize * ITEMS_PER_THREAD; + let (full_keys, selected_keys, descending_keys, values) = + coop_radix_sort_keys_helper::sort_tile::<_, ITEMS_PER_THREAD, 16>( + block, input, offset, begin, end, + ); + coop::store( + block, + &mut full, + &full_keys, + coop::StoreOptions::default().offset(offset), + ); + coop::store( + block, + &mut selected, + &selected_keys, + coop::StoreOptions::default().offset(offset), + ); + coop::store( + block, + &mut descending, + &descending_keys, + coop::StoreOptions::default().offset(offset), + ); + coop::store( + block, + &mut preserved, + &values, + coop::StoreOptions::default().offset(offset), + ); + } +} + +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, LaunchConfig3D}; + use std::cmp::Reverse; + + 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 = 64 * $items; + let input: Vec = (0..2 * tile) + .map(|i| { + let local = i % tile; + let block = i / tile; + match local % 23 { + 0 => i64::MIN, + 1 => i64::MAX, + 2 => -1, + 3 => 0, + _ => { + let digit = ((local * 5 + block * 3) % 8) as u64; + let tag = (tile - local + block * tile) as u64; + ((digit << 61 | tag) ^ (1_u64 << 63)) as i64 + } + } + }) + .collect(); + let mut expected_full = input.clone(); + let mut expected_selected = input.clone(); + let mut expected_descending = input.clone(); + for values in expected_full.chunks_exact_mut(tile) { + values.sort(); + } + // Stable host sorting retains the original order of distinct keys + // sharing selected digits. Signed normalization precedes selection. + let digit = |key: &i64| ((*key as u64) ^ (1_u64 << 63)) >> 61; + for values in expected_selected.chunks_exact_mut(tile) { + values.sort_by_key(digit); + } + for values in expected_descending.chunks_exact_mut(tile) { + values.sort_by_key(|key| Reverse(digit(key))); + } + let source = DeviceBuffer::from_host(&stream, &input)?; + let full_target = DeviceBuffer::from_host(&stream, &vec![17_i64; input.len()])?; + let selected_target = DeviceBuffer::from_host(&stream, &vec![17_i64; input.len()])?; + let descending_target = DeviceBuffer::from_host(&stream, &vec![17_i64; input.len()])?; + let preserved_target = DeviceBuffer::from_host(&stream, &vec![17_i64; input.len()])?; + let plan = module + .prepare_sort_tile::<$items>(LaunchConfig3D::auto_shared((2, 1, 1), (4, 4, 4)))?; + let (source, full_target, selected_target, descending_target, preserved_target) = + module + .sort_tile_async_owned::<$items, _, _, _, _, _>( + &plan, + source, + full_target, + selected_target, + descending_target, + preserved_target, + 61, + 64, + ) + .sync()?; + let mut full = vec![0_i64; input.len()]; + let mut selected = vec![0_i64; input.len()]; + let mut descending = vec![0_i64; input.len()]; + let mut preserved = vec![0_i64; input.len()]; + let mut source_after = vec![0_i64; input.len()]; + full_target.copy_to_host(&stream, &mut full)?; + selected_target.copy_to_host(&stream, &mut selected)?; + descending_target.copy_to_host(&stream, &mut descending)?; + preserved_target.copy_to_host(&stream, &mut preserved)?; + source.copy_to_host(&stream, &mut source_after)?; + stream.synchronize()?; + assert_eq!(full, expected_full); + assert_eq!(selected, expected_selected); + assert_eq!(descending, expected_descending); + assert_eq!(preserved, input); + assert_eq!(source_after, input); + println!( + "items/thread={}: {} sorted keys, {} preserved inputs and {} source values checked", + $items, + 3 * input.len(), + preserved.len(), + source_after.len() + ); + }}; + } + run!(1); + run!(4); + Ok(()) +} diff --git a/rust/cuda-coop/examples/radix_sort_pairs/.gitignore b/rust/cuda-coop/examples/radix_sort_pairs/.gitignore new file mode 100644 index 000000000000..74bffa8af4a4 --- /dev/null +++ b/rust/cuda-coop/examples/radix_sort_pairs/.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/radix_sort_pairs/Cargo.lock b/rust/cuda-coop/examples/radix_sort_pairs/Cargo.lock new file mode 100644 index 000000000000..65790854221e --- /dev/null +++ b/rust/cuda-coop/examples/radix_sort_pairs/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-radix-sort-pairs" +version = "0.0.0" +dependencies = [ + "coop-radix-sort-pairs-helper", + "cuda-async", + "cuda-coop", + "cuda-core", + "cuda-device", + "cuda-host", +] + +[[package]] +name = "coop-radix-sort-pairs-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/radix_sort_pairs/Cargo.toml b/rust/cuda-coop/examples/radix_sort_pairs/Cargo.toml new file mode 100644 index 000000000000..d75869b9ef6f --- /dev/null +++ b/rust/cuda-coop/examples/radix_sort_pairs/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-radix-sort-pairs" +version = "0.0.0" +edition = "2024" +publish = false + +[workspace] +members = ["helper"] + +[dependencies] +cuda-coop = { path = "../.." } +coop-radix-sort-pairs-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/radix_sort_pairs/helper/Cargo.toml b/rust/cuda-coop/examples/radix_sort_pairs/helper/Cargo.toml new file mode 100644 index 000000000000..1a094793ba8e --- /dev/null +++ b/rust/cuda-coop/examples/radix_sort_pairs/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-radix-sort-pairs-helper" +version = "0.0.0" +edition = "2024" +publish = false + +[dependencies] +cuda-coop = { path = "../../.." } diff --git a/rust/cuda-coop/examples/radix_sort_pairs/helper/src/lib.rs b/rust/cuda-coop/examples/radix_sort_pairs/helper/src/lib.rs new file mode 100644 index 000000000000..e4153cf4b64f --- /dev/null +++ b/rust/cuda-coop/examples/radix_sort_pairs/helper/src/lib.rs @@ -0,0 +1,77 @@ +// Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +//! Infer both payload types from Load, then stably sort their associated pairs. +#![no_std] +use cuda_coop as coop; + +/// Ascending pairs, descending pairs, then the two preserved input payloads. +pub type SortedTiles = ( + coop::ThreadData, + coop::ThreadData, + coop::ThreadData, + coop::ThreadData, + coop::ThreadData, + coop::ThreadData, +); + +/// Return ascending and descending selected-digit pairs and both intact inputs. +/// +/// Begin-only ascending uses the key width as its end. Descending supplies +/// both runtime bounds. Keys and values have independent minimum alignments. +pub fn sort_tile< + K: coop::RadixKey, + V: coop::RadixValue, + const ITEMS_PER_THREAD: usize, + const KEY_ALIGN: usize, + const VALUE_ALIGN: usize, +>( + block: coop::ThreadGroup, + keys: &[K], + values: &[V], + offset: usize, + begin: i64, + end: i64, +) -> SortedTiles +where + coop::Alignment: coop::SupportedAlignment, + coop::Alignment: coop::SupportedAlignment, +{ + let mut input_keys = coop::ThreadData::<_, ITEMS_PER_THREAD, KEY_ALIGN>::new(); + let mut input_values = coop::ThreadData::<_, ITEMS_PER_THREAD, VALUE_ALIGN>::new(); + coop::load( + block, + keys, + &mut input_keys, + coop::LoadOptions::default().offset(offset), + ); + coop::load( + block, + values, + &mut input_values, + coop::LoadOptions::default().offset(offset), + ); + let (ascending_keys, ascending_values) = coop::radix_sort_pairs( + block, + &input_keys, + &input_values, + coop::RadixSortOptions::default().begin_bit(begin), + ); + let (descending_keys, descending_values) = coop::radix_sort_pairs( + block, + &input_keys, + &input_values, + coop::RadixSortOptions::default() + .begin_bit(begin) + .end_bit(end) + .order(coop::Descending), + ); + ( + ascending_keys, + ascending_values, + descending_keys, + descending_values, + input_keys, + input_values, + ) +} diff --git a/rust/cuda-coop/examples/radix_sort_pairs/src/main.rs b/rust/cuda-coop/examples/radix_sort_pairs/src/main.rs new file mode 100644 index 000000000000..3810d76284d2 --- /dev/null +++ b/rust/cuda-coop/examples/radix_sort_pairs/src/main.rs @@ -0,0 +1,182 @@ +// Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +// The kernel and generated launch wrapper expose separate typed result buffers. +#![allow(clippy::too_many_arguments)] + +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 = 3, block = (4, 4, 4), dynamic_shared = inferred)] + pub fn sort_tile( + keys: &[i32], + values: &[f64], + mut ascending_keys: coop::DisjointSlice>, + mut ascending_values: coop::DisjointSlice>, + mut descending_keys: coop::DisjointSlice>, + mut descending_values: coop::DisjointSlice>, + mut preserved_keys: coop::DisjointSlice>, + mut preserved_values: coop::DisjointSlice>, + mut mutated_keys: coop::DisjointSlice>, + mut mutated_values: coop::DisjointSlice>, + begin: i64, + end: i64, + ) { + let block = coop::this_block(); + let offset = + block.rank_at::() as usize * block.count() as usize * ITEMS_PER_THREAD; + let (mut ascending_k, mut ascending_v, descending_k, descending_v, input_k, input_v) = + coop_radix_sort_pairs_helper::sort_tile::<_, _, ITEMS_PER_THREAD, 16, 32>( + block, keys, values, offset, begin, end, + ); + coop::store( + block, + &mut ascending_keys, + &ascending_k, + coop::StoreOptions::default().offset(offset), + ); + coop::store( + block, + &mut ascending_values, + &ascending_v, + coop::StoreOptions::default().offset(offset), + ); + // Modify fresh outputs after observing the complete original sort. + // The next Stores observe both independent results and both inputs. + ascending_k.write(0, 123_456_789); + ascending_v.write(0, -4096.5); + coop::store( + block, + &mut descending_keys, + &descending_k, + coop::StoreOptions::default().offset(offset), + ); + coop::store( + block, + &mut descending_values, + &descending_v, + coop::StoreOptions::default().offset(offset), + ); + coop::store( + block, + &mut preserved_keys, + &input_k, + coop::StoreOptions::default().offset(offset), + ); + coop::store( + block, + &mut preserved_values, + &input_v, + coop::StoreOptions::default().offset(offset), + ); + coop::store( + block, + &mut mutated_keys, + &ascending_k, + coop::StoreOptions::default().offset(offset), + ); + coop::store( + block, + &mut mutated_values, + &ascending_v, + coop::StoreOptions::default().offset(offset), + ); + } +} + +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, LaunchConfig3D}; + use std::cmp::Reverse; + + 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 = 64 * $items; + let keys: Vec = (0..2 * tile).map(|i| { + let local = i % tile; + let block = i / tile; + let digit = ((local * 5 + block * 3) % 8) as u32; + // Increasing distinct tags make stable selected-digit sorting + // distinguishable from whole-key sorting in descending order. + ((digit << 29 | i as u32) ^ (1_u32 << 31)) as i32 + }).collect(); + let values: Vec = (0..keys.len()).map(|i| i as f64 + 0.25).collect(); + let pairs: Vec<(i32, f64)> = keys.iter().copied().zip(values.iter().copied()).collect(); + let mut expected_ascending = pairs.clone(); + let mut expected_descending = pairs.clone(); + let digit = |pair: &(i32, f64)| ((pair.0 as u32) ^ (1_u32 << 31)) >> 29; + for block in expected_ascending.chunks_exact_mut(tile) { block.sort_by_key(digit); } + for block in expected_descending.chunks_exact_mut(tile) { block.sort_by_key(|pair| Reverse(digit(pair))); } + let mut whole_key_descending = pairs; + for block in whole_key_descending.chunks_exact_mut(tile) { block.sort_by_key(|pair| Reverse(pair.0)); } + assert_ne!(whole_key_descending, expected_descending); + let (expected_ascending_keys, expected_ascending_values): (Vec<_>, Vec<_>) = expected_ascending.into_iter().unzip(); + let (expected_descending_keys, expected_descending_values): (Vec<_>, Vec<_>) = expected_descending.into_iter().unzip(); + let mut expected_mutated_keys = expected_ascending_keys.clone(); + let mut expected_mutated_values = expected_ascending_values.clone(); + for index in (0..keys.len()).step_by($items) { + expected_mutated_keys[index] = 123_456_789; + expected_mutated_values[index] = -4096.5; + } + let source_keys = DeviceBuffer::from_host(&stream, &keys)?; + let source_values = DeviceBuffer::from_host(&stream, &values)?; + let ascending_keys_target = DeviceBuffer::from_host(&stream, &vec![17 as i32; keys.len()])?; + let ascending_values_target = DeviceBuffer::from_host(&stream, &vec![17 as f64; keys.len()])?; + let descending_keys_target = DeviceBuffer::from_host(&stream, &vec![17 as i32; keys.len()])?; + let descending_values_target = DeviceBuffer::from_host(&stream, &vec![17 as f64; keys.len()])?; + let preserved_keys_target = DeviceBuffer::from_host(&stream, &vec![17 as i32; keys.len()])?; + let preserved_values_target = DeviceBuffer::from_host(&stream, &vec![17 as f64; keys.len()])?; + let mutated_keys_target = DeviceBuffer::from_host(&stream, &vec![17 as i32; keys.len()])?; + let mutated_values_target = DeviceBuffer::from_host(&stream, &vec![17 as f64; keys.len()])?; + let plan = module.prepare_sort_tile::<$items>(LaunchConfig3D::auto_shared((2, 1, 1), (4, 4, 4)))?; + let (source_keys, source_values, ascending_keys_target, ascending_values_target, descending_keys_target, descending_values_target, preserved_keys_target, preserved_values_target, mutated_keys_target, mutated_values_target) = module.sort_tile_async_owned::<$items, _, _, _, _, _, _, _, _, _, _>(&plan, source_keys, source_values, ascending_keys_target, ascending_values_target, descending_keys_target, descending_values_target, preserved_keys_target, preserved_values_target, mutated_keys_target, mutated_values_target, 29, 32).sync()?; + let mut ascending_keys = vec![0 as i32; keys.len()]; + ascending_keys_target.copy_to_host(&stream, &mut ascending_keys)?; + let mut ascending_values = vec![0 as f64; keys.len()]; + ascending_values_target.copy_to_host(&stream, &mut ascending_values)?; + let mut descending_keys = vec![0 as i32; keys.len()]; + descending_keys_target.copy_to_host(&stream, &mut descending_keys)?; + let mut descending_values = vec![0 as f64; keys.len()]; + descending_values_target.copy_to_host(&stream, &mut descending_values)?; + let mut preserved_keys = vec![0 as i32; keys.len()]; + preserved_keys_target.copy_to_host(&stream, &mut preserved_keys)?; + let mut preserved_values = vec![0 as f64; keys.len()]; + preserved_values_target.copy_to_host(&stream, &mut preserved_values)?; + let mut mutated_keys = vec![0 as i32; keys.len()]; + mutated_keys_target.copy_to_host(&stream, &mut mutated_keys)?; + let mut mutated_values = vec![0 as f64; keys.len()]; + mutated_values_target.copy_to_host(&stream, &mut mutated_values)?; + let mut source_keys_after = vec![0_i32; keys.len()]; + let mut source_values_after = vec![0_f64; values.len()]; + source_keys.copy_to_host(&stream, &mut source_keys_after)?; + source_values.copy_to_host(&stream, &mut source_values_after)?; + stream.synchronize()?; + assert_eq!(ascending_keys, expected_ascending_keys); + assert_eq!(ascending_values.iter().map(|x| x.to_bits()).collect::>(), expected_ascending_values.iter().map(|x| x.to_bits()).collect::>()); + assert_eq!(descending_keys, expected_descending_keys); + assert_eq!(descending_values.iter().map(|x| x.to_bits()).collect::>(), expected_descending_values.iter().map(|x| x.to_bits()).collect::>()); + assert_eq!(preserved_keys, keys); + assert_eq!(preserved_values.iter().map(|x| x.to_bits()).collect::>(), values.iter().map(|x| x.to_bits()).collect::>()); + assert_eq!(mutated_keys, expected_mutated_keys); + assert_eq!(mutated_values.iter().map(|x| x.to_bits()).collect::>(), expected_mutated_values.iter().map(|x| x.to_bits()).collect::>()); + assert_eq!(source_keys_after, keys); + assert_eq!(source_values_after.iter().map(|x| x.to_bits()).collect::>(), values.iter().map(|x| x.to_bits()).collect::>()); + println!("items/thread={}: {} sorted key/value observations, {} preserved inputs, {} source values and {} mutated outputs checked", $items, 4 * keys.len(), 2 * keys.len(), 2 * keys.len(), 2 * keys.len()); + }}; + } + run!(1); + run!(4); + Ok(()) +} diff --git a/rust/cuda-coop/src/lib.rs b/rust/cuda-coop/src/lib.rs index 69ee57255a18..a015aa5c07fd 100644 --- a/rust/cuda-coop/src/lib.rs +++ b/rust/cuda-coop/src/lib.rs @@ -68,6 +68,13 @@ //! [`histogram`] counts a complete 1D Block's samples into fresh striped counters. //! [`HistogramOptions`] selects the required bin count, independent counter type //! and extent, Atomic or Sort algorithm, and existing scratch policy. +//! [`radix_sort_keys`] stably sorts a complete XYZ Block's integer keys into a +//! fresh blocked payload, preserving input. [`RadixSortOptions`] selects static +//! order, independent static/runtime bit endpoints and existing scratch policy. +//! [`radix_sort_pairs`] carries associated numeric values through the same stable +//! ordering, returning fresh key/value payloads with independent alignments. +//! [`radix_rank_keys`] returns fresh naturally aligned i32 ranks for selected +//! key digits. [`RadixRankOptions`] selects static width/order and begin/end. //! Host payload tests need no GPU; //! executing hardware queries requires an Oxide-compiled CUDA kernel. @@ -82,6 +89,7 @@ mod exchange; mod group; mod histogram; mod load_store; +mod radix_sort; mod reduce; mod reduce_batched; mod scan; @@ -108,6 +116,11 @@ pub use load_store::{ Algorithm, Direct, LoadOptions, LoadSource, LoadStoreGroup, StoreOptions, StoreValue, Striped, Transpose, Vectorize, WarpTranspose, WarpTransposeTimesliced, load, store, }; +pub use radix_sort::{ + Ascending, Descending, RadixBit, RadixKey, RadixRankBits, RadixRankOptions, RadixRankWidth, + RadixSortBit, RadixSortOptions, RadixSortOrder, RadixValue, radix_rank_keys, radix_sort_keys, + radix_sort_pairs, +}; pub use reduce::{ RakingCommutativeOnly, ReduceAlgorithm, ReduceGroup, ReduceOperator, ReduceOptions, ReduceValue, SumOptions, WarpReductions, reduce, sum, diff --git a/rust/cuda-coop/src/radix_sort.rs b/rust/cuda-coop/src/radix_sort.rs new file mode 100644 index 000000000000..0eb95c97a112 --- /dev/null +++ b/rust/cuda-coop/src/radix_sort.rs @@ -0,0 +1,564 @@ +// 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, Numeric, SupportedAlignment, TempStorage, TempStorageLike, ThreadData, + ThreadGroup, +}; + +mod boundary; + +/// Sort selected digits in ascending order (the default Radix Sort order). +#[derive(Clone, Copy, Debug, Default)] +pub struct Ascending; + +/// Sort selected digits in descending order, preserving the order of ties. +#[derive(Clone, Copy, Debug, Default)] +pub struct Descending; + +/// Sealed Radix Sort and Rank key types: i32, u32, i64 and u64. +pub trait RadixKey: Numeric + sealed::Key {} + +/// Sealed associated-value types for [`radix_sort_pairs`]. +/// +/// Supports i8, u8, i16, u16, i32, u32, i64, u64, f32 and f64. Values move +/// with their keys without comparisons or arithmetic, preserving their bits. +pub trait RadixValue: Numeric + sealed::Value {} + +/// A sealed static Radix Sort or Rank order: [`Ascending`] or [`Descending`]. +pub trait RadixSortOrder: sealed::Order {} + +/// A compile-time bit endpoint selected by the const option builders. +/// +/// A [`RadixSortOptions::begin_bit_const`] or +/// [`RadixSortOptions::end_bit_const`] endpoint is checked when the operation +/// is instantiated with its key type and other endpoint. +#[derive(Clone, Copy, Debug, Default)] +pub struct RadixBit; + +/// A sealed Radix Sort endpoint: omitted `()`, runtime `i64` or [`RadixBit`]. +/// +/// This trait records option types; it does not establish uniformity. Runtime +/// endpoints must be Block-uniform, as checked by the device compiler. +pub trait RadixSortBit: sealed::Bit {} + +/// Order, independent bit endpoints and scratch for Radix Sort keys or pairs. +/// +/// Canonical [`Default`] selects [`Ascending`], begin zero, omitted end and +/// implicit scratch. The omitted end is always the key width, even with a +/// nonzero begin. `T` is inferred from input; builders consume options and +/// preserve every other choice, including an explicit storage borrow. +/// +/// Value builders accept `i64` and validate the final interval at runtime. +/// Const builders request compilation-time checks at monomorphization. A +/// literal passed to a value builder still uses the runtime route. Mixed +/// static/runtime bounds are supported. Every interval must satisfy +/// `0 <= begin < end <= key_width`; empty intervals are invalid. +/// +/// ``` +/// use cuda_coop as coop; +/// let _: coop::RadixSortOptions = Default::default(); +/// let _ = coop::RadixSortOptions::::default() +/// .begin_bit_const::<8>().end_bit(32).order(coop::Descending); +/// ``` +pub struct RadixSortOptions { + begin: Begin, + end: End, + storage: S, + selection: PhantomData<(T, O)>, +} + +impl Default for RadixSortOptions { + fn default() -> Self { + Self { + begin: (), + end: (), + storage: (), + selection: PhantomData, + } + } +} + +impl RadixSortOptions { + /// Select ascending or descending order, preserving both endpoints. + pub fn order(self, _order: O2) -> RadixSortOptions { + RadixSortOptions { + begin: self.begin, + end: self.end, + storage: self.storage, + selection: PhantomData, + } + } + + /// Set a Block-uniform runtime begin, without changing the end. + /// + /// An omitted end remains the full key width. Widen smaller integer values + /// losslessly to `i64`; bounds are checked before conversion to CUB's int. + pub fn begin_bit(self, begin: i64) -> RadixSortOptions { + RadixSortOptions { + begin, + end: self.end, + storage: self.storage, + selection: PhantomData, + } + } + + /// Set a Block-uniform runtime end, without changing the begin. + pub fn end_bit(self, end: i64) -> RadixSortOptions { + RadixSortOptions { + begin: self.begin, + end, + storage: self.storage, + selection: PhantomData, + } + } + + /// Set a static begin, checked when the sort operation is instantiated. + /// + /// An omitted end remains the full key width. This replaces any prior begin + /// selection; validation uses the final options, not intermediate builders. + pub fn begin_bit_const(self) -> RadixSortOptions, End, S> { + RadixSortOptions { + begin: RadixBit, + end: self.end, + storage: self.storage, + selection: PhantomData, + } + } + + /// Set a static end, checked when the sort operation is instantiated. + pub fn end_bit_const(self) -> RadixSortOptions, S> { + RadixSortOptions { + begin: self.begin, + end: RadixBit, + storage: self.storage, + selection: PhantomData, + } + } + + /// Borrow scratch with its existing sharing and completion policy. + /// + /// Manual reuse requires `cuda_device::thread::sync_threads()` after the + /// call. Auto-sync storage supplies its trailing barrier. Consuming options + /// permit ordinary mutable reborrows on successive calls. + pub fn temp_storage( + self, + storage: &mut TempStorage, + ) -> RadixSortOptions> { + RadixSortOptions { + begin: self.begin, + end: self.end, + storage, + selection: PhantomData, + } + } +} + +/// Stably sort a complete Block's keys into a fresh blocked payload. +/// +/// Every member participates with positive `ITEMS_PER_THREAD` initialized keys. +/// Rank `r`, item `i` has input/output index `r * ITEMS_PER_THREAD + i`, with +/// x-fastest rank across all three Block dimensions. The checked product of +/// Block size and extent must not exceed 65,535. The result has the same key +/// type, extent and minimum alignment, and every output slot is initialized. +/// Input values and validity bits are preserved, including through a mutable +/// reborrow. Generic T/N/ALIGN consumers need no `generic_const_exprs` feature. +/// +/// Keys support the four [`RadixKey`] types. Equal selected digits retain their +/// original blocked order in both directions. Signed keys invert the sign bit +/// before selecting `[begin_bit, end_bit)` and restore the original key +/// representation in the result. Default options sort the full key ascending. +/// An omitted end always means full key width, independently of begin. +/// +/// Static endpoints are checked at monomorphization; runtime endpoints are +/// Block-uniform `i64` values. Invalid runtime intervals or unreadable inputs +/// trap before native sorting reads. There is no partial count, fill, scalar, +/// Warp, floating-key or striped-output form. +/// +/// Omitting storage requests compiler-planned CUB scratch and synchronization +/// before dirty reuse. Explicit [`TempStorage`] retains its manual/auto-sync +/// policy. Internal CUB synchronization grants no completion of unrelated +/// global accesses; use `cuda_device::thread::sync_threads()` when needed. +/// Execution requires an Oxide-compiled CUDA kernel. +/// +/// ```no_run +/// use cuda_coop as coop; +/// fn sorted( +/// block: coop::ThreadGroup, +/// keys: &coop::ThreadData, +/// ) -> coop::ThreadData +/// where coop::Alignment: coop::SupportedAlignment { +/// coop::radix_sort_keys(block, keys, Default::default()) +/// } +/// ``` +#[inline(always)] +pub fn radix_sort_keys( + _group: ThreadGroup, + keys: &ThreadData, + options: RadixSortOptions, +) -> ThreadData +where + T: RadixKey, + O: RadixSortOrder, + Begin: RadixSortBit, + End: RadixSortBit, + S: TempStorageLike, + Alignment: SupportedAlignment, +{ + boundary::radix_sort_keys(keys, options) +} + +/// Stably sort a complete Block's key/value pairs into two fresh payloads. +/// +/// Keys use the four [`RadixKey`] types; associated values use [`RadixValue`]. +/// Both borrowed inputs have the same positive `ITEMS_PER_THREAD`, with +/// independent minimum alignments. Each result preserves its input's type, +/// extent and alignment, and every result slot is initialized. Both inputs +/// and their validity maps remain unchanged, including through mutable +/// reborrows. Generic K/V/N/ALIGN consumers need no `generic_const_exprs`. +/// +/// Every member of the complete XYZ Block participates. Rank `r`, item `i` +/// has blocked index `r * ITEMS_PER_THREAD + i`, with x-fastest rank. The +/// checked Block-size/extent product must not exceed 65,535. Both input maps +/// must be fully readable, even when an unreadable value's key is initialized. +/// +/// Selected digits order the keys, and values follow their corresponding +/// keys. Equal selected digits retain their original blocked order in either +/// direction. Signed keys invert the sign bit before selecting the interval, +/// then restore their representation. Default options sort all key bits +/// ascending. An omitted end means the key width, independently of begin. +/// +/// [`RadixSortOptions`] shares the keys-only static/runtime endpoint and +/// storage rules: `0 <= begin < end <= key_width`, runtime endpoints are +/// Block-uniform i64, and invalid bounds or unreadable inputs trap before +/// native numeric reads. Manual scratch reuse needs an explicit Block +/// completion; auto-sync storage supplies it. Internal sorting does not +/// complete unrelated global accesses. There is no Warp, partial-input, +/// floating-key or striped-output form. Execution needs an Oxide CUDA kernel. +/// +/// ```no_run +/// use cuda_coop as coop; +/// fn sorted( +/// block: coop::ThreadGroup, +/// keys: &coop::ThreadData, +/// values: &coop::ThreadData, +/// ) -> (coop::ThreadData, coop::ThreadData) +/// where coop::Alignment: coop::SupportedAlignment, +/// coop::Alignment: coop::SupportedAlignment { +/// coop::radix_sort_pairs(block, keys, values, Default::default()) +/// } +/// ``` +#[inline(always)] +pub fn radix_sort_pairs< + K, + V, + O, + Begin, + End, + S, + const ITEMS_PER_THREAD: usize, + const KEY_ALIGN: usize, + const VALUE_ALIGN: usize, +>( + _group: ThreadGroup, + keys: &ThreadData, + values: &ThreadData, + options: RadixSortOptions, +) -> ( + ThreadData, + ThreadData, +) +where + K: RadixKey, + V: RadixValue, + O: RadixSortOrder, + Begin: RadixSortBit, + End: RadixSortBit, + S: TempStorageLike, + Alignment: SupportedAlignment, + Alignment: SupportedAlignment, +{ + boundary::radix_sort_pairs(keys, values, options) +} + +/// A compile-time digit width selected by [`RadixRankOptions::radix_bits`]. +#[derive(Clone, Copy, Debug, Default)] +pub struct RadixRankBits; + +/// A sealed Rank width: omitted `()` or explicit [`RadixRankBits`]. +pub trait RadixRankWidth: sealed::RankWidth {} + +/// Static order and digit width, with a static or Block-uniform runtime begin. +/// +/// Canonical [`Default`] selects ascending, begin zero and four radix bits. +/// With a static end and no explicit width, the width is `end - begin`. +/// Every width must be in `1..=8`, and the complete interval must fit the key. +/// A runtime begin requires an explicit width; a literal passed to `begin_bit` +/// still selects the runtime route. An optional end is always static and must +/// equal begin plus width. Builders validate their final combined choices at +/// monomorphization, or at runtime where begin is dynamic. +/// +/// `T` is inferred from the input. Rank uses implicit compiler-planned scratch. +/// +/// ``` +/// use cuda_coop as coop; +/// let _: coop::RadixRankOptions = Default::default(); +/// let _ = coop::RadixRankOptions::::default() +/// .begin_bit_const::<61>().end_bit_const::<64>(); +/// let _ = coop::RadixRankOptions::::default() +/// .begin_bit(61).radix_bits::<3>().order(coop::Descending); +/// ``` +pub struct RadixRankOptions { + begin: Begin, + selection: PhantomData<(T, O, End, Bits)>, +} + +impl Default for RadixRankOptions { + fn default() -> Self { + Self { + begin: (), + selection: PhantomData, + } + } +} + +impl RadixRankOptions { + /// Select static order, preserving the interval choices. + pub fn order( + self, + _order: O2, + ) -> RadixRankOptions { + RadixRankOptions { + begin: self.begin, + selection: PhantomData, + } + } + + /// Set a Block-uniform runtime begin; an explicit radix width is required. + pub fn begin_bit(self, begin: i64) -> RadixRankOptions { + RadixRankOptions { + begin, + selection: PhantomData, + } + } + + /// Set a static begin, checked with the final width and optional end. + pub fn begin_bit_const( + self, + ) -> RadixRankOptions, End, Bits> { + RadixRankOptions { + begin: RadixBit, + selection: PhantomData, + } + } + + /// Set a static end, deriving width when no explicit width is selected. + pub fn end_bit_const( + self, + ) -> RadixRankOptions, Bits> { + RadixRankOptions { + begin: self.begin, + selection: PhantomData, + } + } + + /// Select a static digit width in `1..=8`, preserving begin and end. + pub fn radix_bits( + self, + ) -> RadixRankOptions> { + RadixRankOptions { + begin: self.begin, + selection: PhantomData, + } + } +} + +/// Rank selected key digits across a complete XYZ Block, preserving the input. +/// +/// Each input item receives its position in the stable ascending or descending +/// selected-digit order. Ties receive increasing ranks in original blocked +/// order. The fresh result is `ThreadData` with natural +/// i32 alignment, independently of the key payload's minimum alignment. +/// Every rank slot is initialized; input values and validity bits are unchanged. +/// +/// All Block members participate with positive `ITEMS_PER_THREAD` initialized +/// keys. The x-fastest Block-size/extent product must not exceed 65,535. Signed +/// keys invert their sign bit before extracting the selected digit. Defaults +/// rank the low four transformed bits, rather than sorting the whole key. +/// Generic T/N/ALIGN callers need no `generic_const_exprs` feature. +/// +/// [`RadixRankOptions`] validates static choices at monomorphization. Dynamic +/// begin values must be Block-uniform; invalid intervals or unreadable inputs +/// trap before numeric reads, with wide bounds checked before narrowing. +/// Scratch and synchronization before dirty reuse are compiler-planned. Rank +/// does not complete unrelated global accesses. There is no explicit storage, +/// prefix-output, Warp or partial-input form. Execution needs an Oxide kernel. +/// +/// ```no_run +/// use cuda_coop as coop; +/// fn ranks( +/// block: coop::ThreadGroup, +/// keys: &coop::ThreadData, +/// ) -> coop::ThreadData +/// where coop::Alignment: coop::SupportedAlignment { +/// coop::radix_rank_keys(block, keys, Default::default()) +/// } +/// ``` +#[inline(always)] +pub fn radix_rank_keys( + _group: ThreadGroup, + keys: &ThreadData, + options: RadixRankOptions, +) -> ThreadData +where + T: RadixKey, + O: RadixSortOrder, + Begin: RadixSortBit, + End: RadixSortBit, + Bits: RadixRankWidth, + Alignment: SupportedAlignment, +{ + boundary::radix_rank_keys(keys, options) +} + +mod sealed { + pub trait Key { + const WIDTH: i64; + } + pub trait Value { + const ASCENDING_REGISTRATION: &'static str; + const DESCENDING_REGISTRATION: &'static str; + } + pub trait RankWidth { + const STATIC: Option; + } + pub trait Order { + const DESCENDING: bool; + } + pub trait Bit { + const DEFAULT: bool; + const STATIC: Option; + fn value(self, default: i64) -> i64; + } +} + +macro_rules! keys { + ($(($ty:ty, $width:literal)),+ $(,)?) => {$ ( + impl sealed::Key for $ty { const WIDTH: i64 = $width; } + impl RadixKey for $ty {} + )+}; +} +keys!((i32, 32), (u32, 32), (i64, 64), (u64, 64)); + +macro_rules! values { + ($(($ty:ty, $ascending:literal, $descending:literal)),+ $(,)?) => {$ ( + impl sealed::Value for $ty { + const ASCENDING_REGISTRATION: &'static str = + cuda_device::library::registration!($ascending); + const DESCENDING_REGISTRATION: &'static str = + cuda_device::library::registration!($descending); + } + impl RadixValue for $ty {} + )+}; +} +values!( + ( + i8, + "COOP_RADIX_SORT_PAIRS_ASCENDING_I8", + "COOP_RADIX_SORT_PAIRS_DESCENDING_I8" + ), + ( + u8, + "COOP_RADIX_SORT_PAIRS_ASCENDING_U8", + "COOP_RADIX_SORT_PAIRS_DESCENDING_U8" + ), + ( + i16, + "COOP_RADIX_SORT_PAIRS_ASCENDING_I16", + "COOP_RADIX_SORT_PAIRS_DESCENDING_I16" + ), + ( + u16, + "COOP_RADIX_SORT_PAIRS_ASCENDING_U16", + "COOP_RADIX_SORT_PAIRS_DESCENDING_U16" + ), + ( + i32, + "COOP_RADIX_SORT_PAIRS_ASCENDING_I32", + "COOP_RADIX_SORT_PAIRS_DESCENDING_I32" + ), + ( + u32, + "COOP_RADIX_SORT_PAIRS_ASCENDING_U32", + "COOP_RADIX_SORT_PAIRS_DESCENDING_U32" + ), + ( + i64, + "COOP_RADIX_SORT_PAIRS_ASCENDING_I64", + "COOP_RADIX_SORT_PAIRS_DESCENDING_I64" + ), + ( + u64, + "COOP_RADIX_SORT_PAIRS_ASCENDING_U64", + "COOP_RADIX_SORT_PAIRS_DESCENDING_U64" + ), + ( + f32, + "COOP_RADIX_SORT_PAIRS_ASCENDING_F32", + "COOP_RADIX_SORT_PAIRS_DESCENDING_F32" + ), + ( + f64, + "COOP_RADIX_SORT_PAIRS_ASCENDING_F64", + "COOP_RADIX_SORT_PAIRS_DESCENDING_F64" + ), +); + +impl sealed::Order for Ascending { + const DESCENDING: bool = false; +} +impl RadixSortOrder for Ascending {} +impl sealed::Order for Descending { + const DESCENDING: bool = true; +} +impl RadixSortOrder for Descending {} + +impl sealed::Bit for () { + const DEFAULT: bool = true; + const STATIC: Option = None; + #[inline(always)] + fn value(self, default: i64) -> i64 { + default + } +} +impl RadixSortBit for () {} +impl sealed::Bit for i64 { + const DEFAULT: bool = false; + const STATIC: Option = None; + #[inline(always)] + fn value(self, _default: i64) -> i64 { + self + } +} +impl RadixSortBit for i64 {} +impl sealed::Bit for RadixBit { + const DEFAULT: bool = false; + const STATIC: Option = Some(BIT); + #[inline(always)] + fn value(self, _default: i64) -> i64 { + BIT + } +} +impl RadixSortBit for RadixBit {} + +impl sealed::RankWidth for () { + const STATIC: Option = None; +} +impl RadixRankWidth for () {} +impl sealed::RankWidth for RadixRankBits { + const STATIC: Option = Some(BITS); +} +impl RadixRankWidth for RadixRankBits {} diff --git a/rust/cuda-coop/src/radix_sort/boundary.rs b/rust/cuda-coop/src/radix_sort/boundary.rs new file mode 100644 index 000000000000..0f4bf809b9e9 --- /dev/null +++ b/rust/cuda-coop/src/radix_sort/boundary.rs @@ -0,0 +1,290 @@ +// Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +use super::{ + RadixKey, RadixRankOptions, RadixRankWidth, RadixSortBit, RadixSortOptions, RadixSortOrder, + RadixValue, +}; +use crate::{Alignment, SupportedAlignment, TempStorageLike, ThreadData}; +use cuda_device::library; + +fn ascending() -> &'static str { + library::registration!("COOP_RADIX_SORT_KEYS_ASCENDING") +} +fn descending() -> &'static str { + library::registration!("COOP_RADIX_SORT_KEYS_DESCENDING") +} + +struct Bounds(core::marker::PhantomData<(T, Begin, End)>); +impl Bounds { + const BEGIN: Option = if Begin::DEFAULT { + Some(0) + } else { + Begin::STATIC + }; + const END: Option = if End::DEFAULT { + Some(T::WIDTH) + } else { + End::STATIC + }; + const VALID: () = { + if let Some(begin) = Self::BEGIN { + assert!( + begin >= 0 && begin < T::WIDTH, + "Radix Sort begin must be in 0..key_width" + ); + } + if let Some(end) = Self::END { + assert!( + end > 0 && end <= T::WIDTH, + "Radix Sort end must be in 1..=key_width" + ); + } + if let (Some(begin), Some(end)) = (Self::BEGIN, Self::END) { + assert!(begin < end, "Radix Sort requires begin < end"); + } + }; +} + +#[inline(always)] +pub(super) fn radix_sort_keys( + input: &ThreadData, + options: RadixSortOptions, +) -> ThreadData +where + T: RadixKey, + O: RadixSortOrder, + Begin: RadixSortBit, + End: RadixSortBit, + S: TempStorageLike, + Alignment: SupportedAlignment, +{ + let () = Bounds::::VALID; + let mut output = ThreadData::new(); + let (input_pointer, input_bits) = input.block_io_read_parts(); + let (output_pointer, output_bits) = output.block_io_parts(); + let payload = (input_pointer, output_pointer, input_bits, output_bits); + let uniform = (options.begin.value(0), options.end.value(T::WIDTH)); + let storage = options.storage.into_raw(); + // SAFETY: input and fresh output are distinct aligned live local buffers + // with independent N-bit validity maps. The exact recipe checks the wide + // bounds and every input bit before numeric reads, sorts a private copy, + // then writes every result and output bit. No pointer escapes. The genuine + // marker checks complete Block participation, uniform controls and storage + // origin/layout/reuse; the provider checks the 65,535-item tile capacity. + unsafe { + if O::DESCENDING { + library::checked_block_payload_call_v1::( + descending, storage, uniform, payload, + ); + } else { + library::checked_block_payload_call_v1::( + ascending, storage, uniform, payload, + ); + } + } + output +} + +fn pairs_ascending() -> &'static str { + V::ASCENDING_REGISTRATION +} +fn pairs_descending() -> &'static str { + V::DESCENDING_REGISTRATION +} + +#[inline(always)] +pub(super) fn radix_sort_pairs< + K, + V, + O, + Begin, + End, + S, + const N: usize, + const KA: usize, + const VA: usize, +>( + keys: &ThreadData, + values: &ThreadData, + options: RadixSortOptions, +) -> (ThreadData, ThreadData) +where + K: RadixKey, + V: RadixValue, + O: RadixSortOrder, + Begin: RadixSortBit, + End: RadixSortBit, + S: TempStorageLike, + Alignment: SupportedAlignment, + Alignment: SupportedAlignment, +{ + let () = Bounds::::VALID; + let mut output_keys = ThreadData::new(); + let mut output_values = ThreadData::new(); + let (keys_pointer, keys_bits) = keys.block_io_read_parts(); + let (output_keys_pointer, output_keys_bits) = output_keys.block_io_parts(); + let (values_pointer, values_bits) = values.block_io_read_parts(); + let (output_values_pointer, output_values_bits) = output_values.block_io_parts(); + let payload = ( + keys_pointer, + output_keys_pointer, + values_pointer, + output_values_pointer, + keys_bits, + output_keys_bits, + values_bits, + output_values_bits, + ); + let uniform = (options.begin.value(0), options.end.value(K::WIDTH)); + let storage = options.storage.into_raw(); + // SAFETY: both borrowed inputs and both fresh outputs are aligned live + // local buffers, each with its own N-bit validity map. The closed recipe + // identifies K through the marker and V through its genuine registration. + // It validates wide bounds and BOTH input maps before numeric reads, sorts + // private copies and initializes every output slot/map. No pointer escapes. + // The marker proves complete Block participation, control uniformity and + // storage policy; the provider checks the 65,535-item capacity and mixed ABI. + unsafe { + if O::DESCENDING { + library::checked_block_payload_call_v1::( + pairs_descending::, + storage, + uniform, + payload, + ); + } else { + library::checked_block_payload_call_v1::( + pairs_ascending::, + storage, + uniform, + payload, + ); + } + } + (output_keys, output_values) +} + +struct RankBounds(core::marker::PhantomData<(T, O, Begin, End, Bits)>); +impl + RankBounds +{ + const BEGIN: Option = if Begin::DEFAULT { + Some(0) + } else { + Begin::STATIC + }; + const END: Option = if End::DEFAULT { None } else { End::STATIC }; + const WIDTH: i64 = { + assert!( + End::DEFAULT || End::STATIC.is_some(), + "Radix Rank end must be static" + ); + if let Some(begin) = Self::BEGIN { + assert!( + begin >= 0 && begin < T::WIDTH, + "Radix Rank begin must be in 0..key_width" + ); + } + if let Some(end) = Self::END { + assert!( + end > 0 && end <= T::WIDTH, + "Radix Rank end must be in 1..=key_width" + ); + } + let width = match (Bits::STATIC, Self::BEGIN, Self::END) { + (Some(width), _, _) => width, + (None, Some(begin), Some(end)) => end - begin, + (None, Some(_), None) => 4, + (None, None, _) => panic!("Radix Rank runtime begin requires explicit radix_bits"), + }; + assert!( + width >= 1 && width <= 8, + "Radix Rank radix_bits must be in 1..=8" + ); + if let Some(begin) = Self::BEGIN { + assert!( + begin <= T::WIDTH - width, + "Radix Rank interval exceeds key width" + ); + } + if let Some(end) = Self::END { + assert!(end >= width, "Radix Rank end is smaller than radix_bits"); + if let Some(begin) = Self::BEGIN { + assert!( + begin == end - width, + "Radix Rank end must equal begin + radix_bits" + ); + } + } + width + }; + const REGISTRATION: &'static str = match (O::DESCENDING, Self::WIDTH) { + (false, 1) => library::registration!("COOP_RADIX_RANK_KEYS_ASCENDING_1"), + (false, 2) => library::registration!("COOP_RADIX_RANK_KEYS_ASCENDING_2"), + (false, 3) => library::registration!("COOP_RADIX_RANK_KEYS_ASCENDING_3"), + (false, 4) => library::registration!("COOP_RADIX_RANK_KEYS_ASCENDING_4"), + (false, 5) => library::registration!("COOP_RADIX_RANK_KEYS_ASCENDING_5"), + (false, 6) => library::registration!("COOP_RADIX_RANK_KEYS_ASCENDING_6"), + (false, 7) => library::registration!("COOP_RADIX_RANK_KEYS_ASCENDING_7"), + (false, 8) => library::registration!("COOP_RADIX_RANK_KEYS_ASCENDING_8"), + (true, 1) => library::registration!("COOP_RADIX_RANK_KEYS_DESCENDING_1"), + (true, 2) => library::registration!("COOP_RADIX_RANK_KEYS_DESCENDING_2"), + (true, 3) => library::registration!("COOP_RADIX_RANK_KEYS_DESCENDING_3"), + (true, 4) => library::registration!("COOP_RADIX_RANK_KEYS_DESCENDING_4"), + (true, 5) => library::registration!("COOP_RADIX_RANK_KEYS_DESCENDING_5"), + (true, 6) => library::registration!("COOP_RADIX_RANK_KEYS_DESCENDING_6"), + (true, 7) => library::registration!("COOP_RADIX_RANK_KEYS_DESCENDING_7"), + (true, 8) => library::registration!("COOP_RADIX_RANK_KEYS_DESCENDING_8"), + _ => unreachable!(), + }; +} + +fn rank_registration< + T: RadixKey, + O: RadixSortOrder, + Begin: RadixSortBit, + End: RadixSortBit, + Bits: RadixRankWidth, +>() -> &'static str { + RankBounds::::REGISTRATION +} + +#[inline(always)] +pub(super) fn radix_rank_keys( + input: &ThreadData, + options: RadixRankOptions, +) -> ThreadData +where + T: RadixKey, + O: RadixSortOrder, + Begin: RadixSortBit, + End: RadixSortBit, + Bits: RadixRankWidth, + Alignment: SupportedAlignment, +{ + let _ = RankBounds::::WIDTH; + let mut output = ThreadData::::new(); + let (input_pointer, input_bits) = input.block_io_read_parts(); + let (output_pointer, output_bits) = output.block_io_parts(); + let payload = (input_pointer, output_pointer, input_bits, output_bits); + let uniform = ( + options.begin.value(0), + RankBounds::::END.unwrap_or(-1), + ); + // SAFETY: the borrowed T input and fresh naturally aligned i32 output are + // distinct live local buffers with independent N-bit maps. The closed + // width/order recipe checks the wide interval and every input bit before + // reads, ranks a private unsigned copy and initializes every output/map. + // No pointer escapes. The genuine marker proves complete Block membership + // and uniform controls; the provider checks XYZ and 65,535-item capacity. + unsafe { + library::checked_block_payload_call_v1::( + rank_registration::, + (), + uniform, + payload, + ); + } + output +} diff --git a/rust/cuda-coop/tests/check_load_store.py b/rust/cuda-coop/tests/check_load_store.py index d85775cfc065..3f897769a8f6 100644 --- a/rust/cuda-coop/tests/check_load_store.py +++ b/rust/cuda-coop/tests/check_load_store.py @@ -1642,6 +1642,237 @@ ), ) +# Radix Sort keys retain the input, scalar uniformity and storage policy. +radix_sort = "let result=coop::radix_sort_keys(block,&values,Default::default());" +radix_manual = radix_sort.replace( + "Default::default()", "coop::RadixSortOptions::default().temp_storage(&mut scratch)" +) +cases.update( + radix_sort_generic_bounds=( + load + "fn forward(" + "block:coop::ThreadGroup,keys:&coop::ThreadData,end:i64)" + "->coop::ThreadData where coop::Alignment:coop::SupportedAlignment {" + "coop::radix_sort_keys(block,keys,coop::RadixSortOptions::default()" + ".begin_bit_const::<8>().end_bit(end).order(coop::Descending))}" + "let result=forward(block,&mut values,valid as i64);" + + store.replace("&values", "&result"), + None, + ), + radix_sort_storage_completed=( + load + new_storage + radix_manual + sync + radix_manual + store, + None, + ), + radix_sort_storage_auto=( + load + + new_storage.replace("new()", "new().auto_sync()") + + radix_manual + + radix_manual + + store, + None, + ), + radix_sort_read_completed=( + reload.replace(".valid_items(valid)", ".valid_items(valid).oob_default(0)") + + radix_sort + + sync + + store, + None, + ), + radix_sort_varying_begin=( + load + + radix_sort.replace( + "Default::default()", + "coop::RadixSortOptions::default().begin_bit(thread::threadIdx_x() as i64)", + ), + "not proved block-uniform", + ), + radix_sort_varying_end=( + load + + radix_sort.replace( + "Default::default()", + "coop::RadixSortOptions::default().end_bit(thread::threadIdx_x() as i64 + 1)", + ), + "not proved block-uniform", + ), + radix_sort_result_control=( + load + radix_sort + "if result[0]==0 {" + radix_sort + "}", + "lane-varying", + ), + radix_sort_storage_missing_barrier=( + load + new_storage + radix_manual + radix_manual, + "requires a full-block barrier", + ), + radix_sort_read_epoch=( + reload.replace(".valid_items(valid)", ".valid_items(valid).oob_default(0)") + + radix_sort + + store, + "requires a barrier covering both accesses", + ), + radix_sort_tile_capacity=( + "values.write(0,0_i32);let keys=coop::ThreadData::::new();" + "let result=coop::radix_sort_keys(block,&keys,Default::default());", + "at most 65535 keys", + ), + radix_sort_warp=( + load + radix_sort.replace("keys(block,", "keys(coop::this_warp(),"), + "mismatched types", + ), +) + +# Paired Radix Sort keeps independent value types and both returned payloads. +pair_values = ( + "let mut associated=coop::ThreadData::::new();" + "for item in 0..4 { associated.write(item,item as f64+0.25); }" +) +pair_sort = "let (result,pair_result)=coop::radix_sort_pairs(block,&values,&associated,Default::default());" +pair_manual = pair_sort.replace( + "Default::default()", "coop::RadixSortOptions::default().temp_storage(&mut scratch)" +) +cases.update( + radix_pairs_generic_defaults=( + load + + pair_values + + "let mut keys=coop::ThreadData::::new();" + + load.replace("&mut values", "&mut keys") + + "fn forward(" + "block:coop::ThreadGroup,keys:&coop::ThreadData," + "values:&coop::ThreadData,end:i64)" + "->(coop::ThreadData,coop::ThreadData) " + "where coop::Alignment:coop::SupportedAlignment,coop::Alignment:coop::SupportedAlignment {" + "let _=coop::radix_sort_pairs(block,keys,values,Default::default());" + "coop::radix_sort_pairs(block,keys,values,coop::RadixSortOptions::default()" + ".begin_bit_const::<8>().end_bit(end).order(coop::Descending))}" + "let (result,pair_result)=forward(block,&mut keys,&mut associated,valid as i64);" + + store.replace("&values", "&result") + + store.replace("&mut output", "&mut value_output").replace( + "&values", "&pair_result" + ), + None, + ), + radix_pairs_storage_completed=( + load + pair_values + new_storage + pair_manual + sync + pair_manual + store, + None, + ), + radix_pairs_storage_auto=( + load + + pair_values + + new_storage.replace("new()", "new().auto_sync()") + + pair_manual + + pair_manual + + store, + None, + ), + radix_pairs_read_completed=( + reload.replace(".valid_items(valid)", ".valid_items(valid).oob_default(0)") + + pair_values + + pair_sort + + sync + + store, + None, + ), + radix_pairs_unequal_extent=( + load + + pair_values.replace("f64,4,32", "f64,3,32").replace("0..4", "0..3") + + pair_sort, + "mismatched types", + ), + radix_pairs_empty_static=( + load + + pair_values + + pair_sort.replace( + "Default::default()", + "coop::RadixSortOptions::default().begin_bit_const::<8>().end_bit_const::<8>()", + ), + "error[E0080]: evaluation panicked: Radix Sort requires begin < end", + ), + radix_pairs_varying_end=( + load + + pair_values + + pair_sort.replace( + "Default::default()", + "coop::RadixSortOptions::default().end_bit(thread::threadIdx_x() as i64+1)", + ), + "not proved block-uniform", + ), + radix_pairs_value_control=( + load + pair_values + pair_sort + "if pair_result[0]==0.0 {" + pair_sort + "}", + "lane-varying", + ), + radix_pairs_storage_missing_barrier=( + load + pair_values + new_storage + pair_manual + pair_manual, + "requires a full-block barrier", + ), + radix_pairs_read_epoch=( + reload.replace(".valid_items(valid)", ".valid_items(valid).oob_default(0)") + + pair_values + + pair_sort + + store, + "requires a barrier covering both accesses", + ), + radix_pairs_tile_capacity=( + "values.write(0,0_i32);let keys=coop::ThreadData::::new();" + "let associated=coop::ThreadData::::new();" + "let _=coop::radix_sort_pairs(block,&keys,&associated,Default::default());", + "at most 65535 keys", + ), +) + +# Rank returns natural-aligned i32 positions while preserving its key payload. +radix_rank = "let result=coop::radix_rank_keys(block,&values,Default::default());" +cases.update( + radix_rank_generic_defaults=( + load + "fn forward(" + "block:coop::ThreadGroup,keys:&coop::ThreadData,begin:i64)" + "->coop::ThreadData where coop::Alignment:coop::SupportedAlignment {" + "let _:coop::ThreadData=coop::radix_rank_keys(block,keys,Default::default());" + "coop::radix_rank_keys(block,keys,coop::RadixRankOptions::default()" + ".radix_bits::<4>().begin_bit(begin).order(coop::Descending))}" + "let mut keys=coop::ThreadData::::new();" + + load.replace("&mut values", "&mut keys") + + "let result=forward(block,&keys,valid as i64);" + + store.replace("&values", "&result"), + None, + ), + radix_rank_implicit_reuse=( + load + + radix_rank.replace("result", "first") + + radix_rank + + "values.write(0,first[0]+result[0]);" + + store, + None, + ), + radix_rank_read_completed=( + reload.replace(".valid_items(valid)", ".valid_items(valid).oob_default(0)") + + radix_rank + + sync + + store, + None, + ), + radix_rank_varying_begin=( + load + + radix_rank.replace( + "Default::default()", + "coop::RadixRankOptions::default().radix_bits::<4>()" + ".begin_bit(thread::threadIdx_x() as i64)", + ), + "not proved block-uniform", + ), + radix_rank_result_control=( + load + radix_rank + "if result[0]==0 {" + radix_rank + "}", + "lane-varying", + ), + radix_rank_read_epoch=( + reload.replace(".valid_items(valid)", ".valid_items(valid).oob_default(0)") + + radix_rank + + store, + "requires a barrier covering both accesses", + ), + radix_rank_tile_capacity=( + "values.write(0,0_i32);let keys=coop::ThreadData::::new();" + "let _=coop::radix_rank_keys(block,&keys,Default::default());", + "at most 65535 keys", + ), +) + if args.case: unknown = set(args.case).difference(cases) if unknown: @@ -1681,6 +1912,11 @@ "use cuda_coop as coop;", "#![feature(generic_const_exprs)]\n#![allow(incomplete_features)]\nuse cuda_coop as coop;", ) + if name == "radix_pairs_generic_defaults": + case_prefix = case_prefix.replace( + "valid:usize)", + "valid:usize,mut value_output:DisjointSlice>)", + ) if name == "logical_width1_origin": case_prefix = case_prefix.replace( "valid:usize)", diff --git a/rust/cuda-coop/tests/radix_rank_keys/.gitignore b/rust/cuda-coop/tests/radix_rank_keys/.gitignore new file mode 100644 index 000000000000..74bffa8af4a4 --- /dev/null +++ b/rust/cuda-coop/tests/radix_rank_keys/.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/radix_rank_keys/Cargo.lock b/rust/cuda-coop/tests/radix_rank_keys/Cargo.lock new file mode 100644 index 000000000000..4f950d2e7016 --- /dev/null +++ b/rust/cuda-coop/tests/radix_rank_keys/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-radix-rank-keys-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/radix_rank_keys/Cargo.toml b/rust/cuda-coop/tests/radix_rank_keys/Cargo.toml new file mode 100644 index 000000000000..0322286ab408 --- /dev/null +++ b/rust/cuda-coop/tests/radix_rank_keys/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-radix-rank-keys-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/radix_rank_keys/src/main.rs b/rust/cuda-coop/tests/radix_rank_keys/src/main.rs new file mode 100644 index 000000000000..5a8a2b552e02 --- /dev/null +++ b/rust/cuda-coop/tests/radix_rank_keys/src/main.rs @@ -0,0 +1,319 @@ +// Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +// Kernels expose separate ranks, preserved inputs and a retained fresh result. +#![allow(clippy::too_many_arguments)] + +use cuda_coop as coop; +use cuda_device::{kernel, launch_contract, thread}; +use cuda_host::cuda_module; + +#[cuda_module] +mod kernels { + use super::*; + + macro_rules! exercise { + ($n:literal, $a:literal, $b:literal, $first:expr, $second:expr, + $keys:ident, $first_output:ident, $second_output:ident, $preserved:ident + $(, $mutated:ident)?) => {{ + let block = coop::this_block(); + let offset = block.rank_at::() as usize * $b * $n; + let mut keys = coop::ThreadData::<_, $n, $a>::new(); + coop::load(block, $keys, &mut keys, coop::LoadOptions::default().offset(offset)); + #[allow(unused_mut)] + let mut first: coop::ThreadData = + coop::radix_rank_keys(block, &keys, $first); + coop::store(block, &mut $first_output, &first, coop::StoreOptions::default().offset(offset)); + $( + first.write(0, -7); + let _ = &$mutated; + )? + let second: coop::ThreadData = + coop::radix_rank_keys(block, &keys, $second); + coop::store(block, &mut $second_output, &second, coop::StoreOptions::default().offset(offset)); + coop::store(block, &mut $preserved, &keys, coop::StoreOptions::default().offset(offset)); + $(coop::store(block, &mut $mutated, &first, coop::StoreOptions::default().offset(offset));)? + }}; + } + + #[kernel] + #[launch_contract(domain = 1, block = (33, 1, 1), dynamic_shared = inferred)] + pub fn i32_nonwarp( + keys: &[i32], + mut first: coop::DisjointSlice>, + mut second: coop::DisjointSlice>, + mut preserved: coop::DisjointSlice>, + mut mutated: coop::DisjointSlice>, + begin: i64, + invalid: u32, + ) { + if invalid != 0 { + let block = coop::this_block(); + let rank = thread::threadIdx_x() as usize; + let mut input = coop::ThreadData::<_, 9, 32>::new(); + for item in 0..9 { + if !(invalid == 1 && rank == 32 && item == 8) { + input.write(item, keys[rank * 9 + item]); + } + } + if invalid == 2 || invalid == 3 { + let _ = coop::radix_rank_keys( + block, + &input, + coop::RadixRankOptions::default() + .radix_bits::<8>() + .begin_bit(begin) + .order(coop::Descending), + ); + } else { + let _ = coop::radix_rank_keys( + block, + &input, + coop::RadixRankOptions::default() + .radix_bits::<8>() + .begin_bit(begin) + .end_bit_const::<32>() + .order(coop::Descending), + ); + } + // A missing native preflight must fail on host completion. The + // invalid Rank result is discarded, never read to force a trap. + let mut report = coop::ThreadData::<_, 9>::new(); + for item in 0..9 { + report.write(item, 99_i32); + } + coop::store(block, &mut mutated, &report, Default::default()); + return; + } + exercise!( + 9, + 32, + 33, + coop::RadixRankOptions::default() + .begin_bit_const::<24>() + .end_bit_const::<32>(), + coop::RadixRankOptions::default() + .radix_bits::<8>() + .begin_bit(begin) + .end_bit_const::<32>() + .order(coop::Descending), + keys, + first, + second, + preserved, + mutated + ); + } + + #[kernel] + #[launch_contract(domain = 3, block = (3, 5, 2), dynamic_shared = inferred)] + pub fn u32_xyz( + keys: &[u32], + mut first: coop::DisjointSlice>, + mut second: coop::DisjointSlice>, + mut preserved: coop::DisjointSlice>, + begin: i64, + ) { + exercise!( + 4, + 16, + 30, + coop::RadixRankOptions::default(), + coop::RadixRankOptions::default() + .radix_bits::<4>() + .begin_bit(begin) + .order(coop::Descending), + keys, + first, + second, + preserved + ); + } + + #[kernel] + #[launch_contract(domain = 1, block = (64, 1, 1), dynamic_shared = inferred)] + pub fn u64_one_bit( + keys: &[u64], + mut first: coop::DisjointSlice>, + mut second: coop::DisjointSlice>, + mut preserved: coop::DisjointSlice>, + ) { + exercise!( + 1, + 1, + 64, + coop::RadixRankOptions::default() + .begin_bit_const::<63>() + .radix_bits::<1>(), + coop::RadixRankOptions::default() + .begin_bit_const::<63>() + .end_bit_const::<64>() + .order(coop::Descending), + keys, + first, + second, + preserved + ); + } +} + +// Stable sorting of original indices, followed by permutation inversion, gives +// ranks in input-slot order and distinguishes stable ties from sorted digits. +fn oracle( + keys: &[K], + tile: usize, + begin: u32, + end: u32, + descending: bool, + ordered: impl Fn(K) -> u128, +) -> Vec { + let mask = (1_u128 << (end - begin)) - 1; + let mut output = vec![-1; keys.len()]; + for (block, values) in keys.chunks_exact(tile).enumerate() { + let digit = |index: usize| (ordered(values[index]) >> begin) & mask; + let mut indices: Vec = (0..tile).collect(); + indices.sort_by(|&left, &right| { + if descending { + digit(right).cmp(&digit(left)) + } else { + digit(left).cmp(&digit(right)) + } + }); + for (rank, index) in indices.into_iter().enumerate() { + output[block * tile + index] = rank as i32; + } + let mut permutation = output[block * tile..(block + 1) * tile].to_vec(); + permutation.sort_unstable(); + assert_eq!(permutation, (0..tile as i32).collect::>()); + } + output +} + +fn signed_keys(blocks: usize) -> Vec { + (0..blocks * 297) + .map(|index| { + let j = index % 297; + let digit = (((5 * j + 3 * (index / 297)) % 16) * 17) as u32; + match j { + 0 => i32::MIN, + 1 => i32::MAX, + 2 => -1, + 3 => 0, + _ => ((digit << 24 | (index as u32 + 1)) ^ (1_u32 << 31)) as i32, + } + }) + .collect() +} + +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}; + let (invalid, begin) = match std::env::args().nth(1).as_deref() { + None => (0, 24_i64), + Some("missing-input") => (1, 24), + Some("negative-begin") => (2, i64::MIN), + Some("wide-begin") => (3, (1_i64 << 32) + 24), + Some("end-mismatch") => (4, 23), + Some(other) => panic!("unknown Radix Rank mode {other}"), + }; + init_device_contexts(0, 1)?; + let context = with_cuda_context(0, |context| context.clone())?; + let stream = context.default_stream(); + // SAFETY: typed declarations and embedded kernels are compiled together. + let module = unsafe { kernels::load_async(0)? }; + let (mut cases, mut ranks, mut preserved, mut sources, mut mutated) = (0, 0, 0, 0, 0); + macro_rules! run { + ($t:ty, $n:literal, $b:literal, $prepare:ident, $launch:ident, + $config:expr, $keys:expr, $first:expr, $second:expr, $ordered:expr, + [$($argument:expr),*] $(, $mutation:ident)?) => {{ + let keys: Vec<$t> = $keys; + let (b1, e1, d1) = $first; + let (b2, e2, d2) = $second; + let first = oracle(&keys, $b * $n, b1, e1, d1, $ordered); + let second = oracle(&keys, $b * $n, b2, e2, d2, $ordered); + let whole = oracle(&keys, $b * $n, 0, <$t>::BITS, d2, $ordered); + assert_ne!(second, whole, "selected-digit ranks must distinguish whole-key order"); + let input = DeviceBuffer::from_host(&stream, &keys)?; + let output1 = DeviceBuffer::from_host(&stream, &vec![-1_i32; keys.len()])?; + let output2 = DeviceBuffer::from_host(&stream, &vec![-1_i32; keys.len()])?; + let keep = DeviceBuffer::from_host(&stream, &vec![0 as $t; keys.len()])?; + $(let $mutation = DeviceBuffer::from_host(&stream, &vec![-1_i32; keys.len()])?;)? + stream.synchronize()?; + let plan = module.$prepare($config)?; + let (input, output1, output2, keep $(, $mutation)?) = module.$launch( + &plan, input, output1, output2, keep $(, $mutation)? $(, $argument)*).sync()?; + stream.synchronize()?; + if invalid != 0 { panic!("invalid Radix Rank mode {invalid} unexpectedly completed"); } + assert_eq!(output1.to_host_vec(&stream)?, first, stringify!($prepare)); + assert_eq!(output2.to_host_vec(&stream)?, second, stringify!($prepare)); + assert_eq!(keep.to_host_vec(&stream)?, keys, stringify!($prepare)); + assert_eq!(input.to_host_vec(&stream)?, keys, stringify!($prepare)); + $(let mut expected = first; + for value in expected.iter_mut().step_by($n) { *value = -7; } + assert_eq!($mutation.to_host_vec(&stream)?, expected); + mutated += keys.len();)? + cases += 1; ranks += 2 * keys.len(); preserved += keys.len(); sources += keys.len(); + }}; + } + run!( + i32, + 9, + 33, + prepare_i32_nonwarp, + i32_nonwarp_async_owned, + LaunchConfig1D::auto_shared(if invalid == 0 { 2 } else { 1 }, 33), + signed_keys(if invalid == 0 { 2 } else { 1 }), + (24, 32, false), + (24, 32, true), + |v: i32| ((v as u32) ^ (1_u32 << 31)) as u128, + [begin, invalid], + changed + ); + run!( + u32, + 4, + 30, + prepare_u32_xyz, + u32_xyz_async_owned, + LaunchConfig3D::auto_shared((2, 1, 1), (3, 5, 2)), + (0..240) + .map(|i| match i % 120 { + 0 => 0, + 1 => u32::MAX, + _ => (1_u32 << 31) | ((i + 1) << 16) | (((5 * i + 3) % 16) << 12) | ((3 * i) % 16), + }) + .collect(), + (0, 4, false), + (12, 16, true), + |v: u32| v as u128, + [12_i64] + ); + run!( + u64, + 1, + 64, + prepare_u64_one_bit, + u64_one_bit_async_owned, + LaunchConfig1D::auto_shared(2, 64), + (0..128) + .map(|i| match i % 64 { + 0 => 0, + 1 => u64::MAX, + _ => (((i & 1) as u64) << 63) | (i as u64 + 1), + }) + .collect(), + (63, 64, false), + (63, 64, true), + |v: u64| v as u128, + [] + ); + assert_eq!( + (cases, ranks, preserved, sources, mutated), + (3, 1924, 962, 962, 594) + ); + println!( + "radix_rank_keys: {cases} cases, {ranks} ranks, {preserved} preserved keys, {sources} source keys, {mutated} retained mutated ranks passed" + ); + Ok(()) +} diff --git a/rust/cuda-coop/tests/radix_sort_keys/.gitignore b/rust/cuda-coop/tests/radix_sort_keys/.gitignore new file mode 100644 index 000000000000..74bffa8af4a4 --- /dev/null +++ b/rust/cuda-coop/tests/radix_sort_keys/.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/radix_sort_keys/Cargo.lock b/rust/cuda-coop/tests/radix_sort_keys/Cargo.lock new file mode 100644 index 000000000000..e1c5905bc2c5 --- /dev/null +++ b/rust/cuda-coop/tests/radix_sort_keys/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-radix-sort-keys-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/radix_sort_keys/Cargo.toml b/rust/cuda-coop/tests/radix_sort_keys/Cargo.toml new file mode 100644 index 000000000000..7a3572a3143d --- /dev/null +++ b/rust/cuda-coop/tests/radix_sort_keys/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-radix-sort-keys-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/radix_sort_keys/src/main.rs b/rust/cuda-coop/tests/radix_sort_keys/src/main.rs new file mode 100644 index 000000000000..ec85faccfd01 --- /dev/null +++ b/rust/cuda-coop/tests/radix_sort_keys/src/main.rs @@ -0,0 +1,522 @@ +// 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; + +const CHECKS: u32 = 15; + +#[cuda_module] +mod kernels { + use super::*; + + fn same( + left: &coop::ThreadData, + right: &coop::ThreadData, + ) -> bool + where + coop::Alignment: coop::SupportedAlignment, + { + let mut equal = true; + let mut i = 0; + while i < N { + equal &= left.get(i) == right.get(i); + i += 1; + } + equal + } + + fn observe( + result: &coop::ThreadData, + expected: &[T], + offset: usize, + ) -> bool + where + coop::Alignment: coop::SupportedAlignment, + { + let mut valid = true; + let mut i = 0; + while i < N { + valid &= result.get(i) == Some(expected[offset + i]); + i += 1; + } + valid + } + + // Expressions retain the concrete bound and storage types of each root. + macro_rules! exercise { + ($n:literal, $align:literal, $b:literal, $rank:expr, + $first_options:expr, $second_options:expr, $manual:literal, + $input:ident, $expected:ident, $output:ident, $preserved:ident, + $changed_output:ident, $reports:ident, $replacement:expr, $marker:expr) => {{ + let block = coop::this_block(); + let block_index = block.rank_at::() as usize; + let rank = $rank as usize; + let offset = block_index * $b * $n; + let mut values = coop::ThreadData::<_, $n, $align>::new(); + coop::load( + block, + $input, + &mut values, + coop::LoadOptions::default().offset(offset), + ); + let before = values; + let mut first = coop::radix_sort_keys(block, &values, $first_options); + let first_matches = observe(&first, $expected, offset + rank * $n); + let mut inputs_preserved = same(&values, &before); + first.write(0, $marker); + let first_modified = first; + let mut changed = values; + changed.write(0, $replacement); + let changed_before = changed; + if $manual { + thread::sync_threads(); + } + let second = coop::radix_sort_keys(block, &changed, $second_options); + let second_matches = + observe(&second, $expected, $expected.len() / 2 + offset + rank * $n); + inputs_preserved &= same(&values, &before) && same(&changed, &changed_before); + let independent = same(&first, &first_modified) && first.get(0) == Some($marker); + if $manual { + thread::sync_threads(); + } + coop::store( + block, + &mut $output, + &second, + coop::StoreOptions::default().offset(offset), + ); + coop::store( + block, + &mut $preserved, + &values, + coop::StoreOptions::default().offset(offset), + ); + coop::store( + block, + &mut $changed_output, + &changed, + coop::StoreOptions::default().offset(offset), + ); + let word = first_matches as u32 + | ((second_matches as u32) << 1) + | ((inputs_preserved as u32) << 2) + | ((independent as u32) << 3); + let mut report = coop::ThreadData::<_, 1>::new(); + report.write(0, word); + coop::store( + block, + &mut $reports, + &report, + coop::StoreOptions::default().offset(block_index * $b), + ); + }}; + } + + #[allow(clippy::too_many_arguments)] + #[kernel] + #[launch_contract(domain = 1, block = (33, 1, 1), dynamic_shared = inferred)] + pub fn i32_nonwarp( + input: &[i32], + expected: &[i32], + mut output: coop::DisjointSlice>, + mut preserved: coop::DisjointSlice>, + mut changed_output: coop::DisjointSlice>, + mut reports: coop::DisjointSlice>, + end: i64, + ) { + exercise!( + 1, + 1, + 33, + thread::threadIdx_x(), + coop::RadixSortOptions::default(), + coop::RadixSortOptions::default() + .order(coop::Descending) + .begin_bit_const::<29>() + .end_bit(end), + false, + input, + expected, + output, + preserved, + changed_output, + reports, + -7_i32, + 0x1357_9bdf_i32 + ); + } + + #[allow(clippy::too_many_arguments)] + #[kernel] + #[launch_contract(domain = 3, block = (3, 5, 2), dynamic_shared = inferred)] + pub fn u32_xyz( + input: &[u32], + expected: &[u32], + mut output: coop::DisjointSlice>, + mut preserved: coop::DisjointSlice>, + mut changed_output: coop::DisjointSlice>, + mut reports: coop::DisjointSlice>, + begin: i64, + end: i64, + ) { + exercise!( + 4, + 16, + 30, + thread::threadIdx_x() + 3 * (thread::threadIdx_y() + 5 * thread::threadIdx_z()), + coop::RadixSortOptions::default() + .begin_bit_const::<4>() + .end_bit_const::<7>(), + coop::RadixSortOptions::default() + .order(coop::Descending) + .begin_bit(begin) + .end_bit(end), + false, + input, + expected, + output, + preserved, + changed_output, + reports, + 7_u32, + 0x1357_9bdf_u32 + ); + } + + #[allow(clippy::too_many_arguments)] + #[kernel] + #[launch_contract(domain = 1, block = (64, 1, 1), dynamic_shared = inferred)] + pub fn i64_manual( + input: &[i64], + expected: &[i64], + mut output: coop::DisjointSlice>, + mut preserved: coop::DisjointSlice>, + mut changed_output: coop::DisjointSlice>, + mut reports: coop::DisjointSlice>, + begin: i64, + end: i64, + invalid: u32, + ) { + let mut scratch = coop::TempStorage::new(); + if invalid != 0 { + let block = coop::this_block(); + let rank = thread::threadIdx_x() as usize; + let mut values = coop::ThreadData::<_, 9, 16>::new(); + let mut i = 0; + while i < 9 { + if !(invalid == 1 && rank == 63 && i == 8) { + values.write(i, input[rank * 9 + i]); + } + i += 1; + } + let _ = coop::radix_sort_keys( + block, + &values, + coop::RadixSortOptions::default() + .begin_bit(begin) + .end_bit(end) + .temp_storage(&mut scratch), + ); + thread::sync_threads(); + // A lost preflight must complete normally and fail on the host. + // No result read or assertion may counterfeit the expected trap. + let mut report = coop::ThreadData::<_, 1>::new(); + report.write(0, 99_u32); + coop::store(block, &mut reports, &report, Default::default()); + return; + } + exercise!( + 9, + 16, + 64, + thread::threadIdx_x(), + coop::RadixSortOptions::default() + .begin_bit(begin) + .end_bit_const::<64>() + .temp_storage(&mut scratch), + coop::RadixSortOptions::default() + .order(coop::Descending) + .begin_bit(begin) + .end_bit(end) + .temp_storage(&mut scratch), + true, + input, + expected, + output, + preserved, + changed_output, + reports, + -7_i64, + 0x1357_9bdf_2468_ace0_i64 + ); + } + + #[allow(clippy::too_many_arguments)] + #[kernel] + #[launch_contract(domain = 1, block = (33, 1, 1), dynamic_shared = inferred)] + pub fn u64_auto( + input: &[u64], + expected: &[u64], + mut output: coop::DisjointSlice>, + mut preserved: coop::DisjointSlice>, + mut changed_output: coop::DisjointSlice>, + mut reports: coop::DisjointSlice>, + ) { + let mut scratch = coop::TempStorage::new().auto_sync(); + exercise!( + 4, + 1, + 33, + thread::threadIdx_x(), + coop::RadixSortOptions::default() + .order(coop::Descending) + .temp_storage(&mut scratch), + coop::RadixSortOptions::default() + .begin_bit_const::<0>() + .end_bit_const::<3>() + .temp_storage(&mut scratch), + false, + input, + expected, + output, + preserved, + changed_output, + reports, + 7_u64, + 0x1357_9bdf_2468_ace0_u64 + ); + } +} + +// Sort original keys by selected ordered digits. Stable descending compares +// digits in reverse; reversing an ascending result would reverse equal digits. +fn oracle( + input: &[T], + tile: usize, + begin: u32, + end: u32, + descending: bool, + ordered: impl Fn(T) -> u128, +) -> Vec { + let mask = (1_u128 << (end - begin)) - 1; + let digit = |key| (ordered(key) >> begin) & mask; + let mut result = input.to_vec(); + for block in result.chunks_exact_mut(tile) { + block.sort_by(|&left, &right| { + if descending { + digit(right).cmp(&digit(left)) + } else { + digit(left).cmp(&digit(right)) + } + }); + } + result +} + +fn signed_keys(width: u32, tile: usize, blocks: usize) -> Vec { + (0..blocks * tile) + .map(|index| { + let j = index % tile; + let b = index / tile; + let digit = ((5 * j + 3 * b) % 8) as u64; + let bits = ((digit << (width - 3)) | (index as u64 + 1)) ^ (1_u64 << (width - 1)); + let key = if width == 32 { + bits as i32 as i64 + } else { + bits as i64 + }; + match j { + 0 => { + if width == 32 { + i32::MIN as i64 + } else { + i64::MIN + } + } + 1 => { + if width == 32 { + i32::MAX as i64 + } else { + i64::MAX + } + } + 2 => -1, + 3 => 0, + _ => key, + } + }) + .collect() +} + +fn unsigned_keys(width: u32, tile: usize, blocks: usize) -> Vec { + (0..blocks * tile) + .map(|index| { + let j = index % tile; + let b = index / tile; + let digit = ((5 * j + 3 * b) % 8) as u64; + let tag = index as u64 + 1; + match j { + 0 => 0, + 1 => { + if width == 32 { + u32::MAX as u64 + } else { + u64::MAX + } + } + _ if width == 32 => (tag << 7) | (digit << 4) | (tag & 15), + _ => (((j & 1) as u64) << 63) | (tag << 3) | digit, + } + }) + .collect() +} + +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}; + let (invalid, begin, end) = match std::env::args().nth(1).as_deref() { + None => (0, 61_i64, 64_i64), + Some("missing-input") => (1, 61, 64), + Some("negative-begin") => (2, i64::MIN, 64), + Some("empty") => (3, 61, 61), + Some("reversed") => (4, 62, 61), + Some("end-beyond-width") => (5, 0, 65), + Some("wide-end") => (6, 0, (1_i64 << 32) + 64), + Some(other) => panic!("unknown Radix Sort Keys mode {other}"), + }; + init_device_contexts(0, 1)?; + let context = with_cuda_context(0, |context| context.clone())?; + let stream = context.default_stream(); + // SAFETY: typed declarations and embedded kernels are compiled together. + let module = unsafe { kernels::load_async(0)? }; + let (mut cases, mut sorted, mut values, mut sources, mut reports) = (0, 0, 0, 0, 0); + macro_rules! run { + ($t:ty, $n:literal, $b:literal, $prepare:ident, $launch:ident, + $config:expr, $input:expr, $replacement:expr, $first:expr, $second:expr, + $ordered:expr $(, $argument:expr)*) => {{ + let input: Vec<$t> = $input; + let blocks = input.len() / ($b * $n); + let mut changed = input.clone(); + for value in changed.iter_mut().step_by($n) { *value = $replacement; } + let (b1, e1, d1) = $first; + let (b2, e2, d2) = $second; + let first_expected = oracle(&input, $b * $n, b1, e1, d1, $ordered); + let second_expected = oracle(&changed, $b * $n, b2, e2, d2, $ordered); + let expected: Vec<$t> = first_expected.iter().chain(&second_expected).copied().collect(); + let source = DeviceBuffer::from_host(&stream, &input)?; + let expected_device = DeviceBuffer::from_host(&stream, &expected)?; + let output = DeviceBuffer::from_host(&stream, &vec![<$t>::MAX; input.len()])?; + let preserved = DeviceBuffer::from_host(&stream, &vec![<$t>::MAX; input.len()])?; + let changed_output = DeviceBuffer::from_host(&stream, &vec![<$t>::MAX; input.len()])?; + let words = DeviceBuffer::from_host(&stream, &vec![0_u32; blocks * $b])?; + stream.synchronize()?; + let plan = module.$prepare($config)?; + let (source, _, output, preserved, changed_output, words) = module.$launch( + &plan, source, expected_device, output, preserved, changed_output, words + $(, $argument)*).sync()?; + stream.synchronize()?; + if invalid != 0 { panic!("invalid Radix Sort Keys mode {invalid} unexpectedly completed"); } + assert_eq!(output.to_host_vec(&stream)?, second_expected, stringify!($prepare)); + assert_eq!(preserved.to_host_vec(&stream)?, input, stringify!($prepare)); + assert_eq!(changed_output.to_host_vec(&stream)?, changed, stringify!($prepare)); + assert_eq!(source.to_host_vec(&stream)?, input, stringify!($prepare)); + assert_eq!(words.to_host_vec(&stream)?, vec![CHECKS; blocks * $b], stringify!($prepare)); + cases += 1; sorted += 2 * input.len(); values += 2 * input.len(); + sources += input.len(); reports += blocks * $b; + }}; + } + if invalid != 0 { + run!( + i64, + 9, + 64, + prepare_i64_manual, + i64_manual_async_owned, + LaunchConfig1D::auto_shared(1, 64), + signed_keys(64, 576, 1), + -7, + (61, 64, false), + (61, 64, true), + |v: i64| ((v as u64) ^ (1_u64 << 63)) as u128, + begin, + end, + invalid + ); + } else { + run!( + i32, + 1, + 33, + prepare_i32_nonwarp, + i32_nonwarp_async_owned, + LaunchConfig1D::auto_shared(2, 33), + signed_keys(32, 33, 2) + .into_iter() + .map(|v| v as i32) + .collect(), + -7, + (0, 32, false), + (29, 32, true), + |v: i32| ((v as u32) ^ (1_u32 << 31)) as u128, + 32_i64 + ); + run!( + u32, + 4, + 30, + prepare_u32_xyz, + u32_xyz_async_owned, + LaunchConfig3D::auto_shared((2, 1, 1), (3, 5, 2)), + unsigned_keys(32, 120, 2) + .into_iter() + .map(|v| v as u32) + .collect(), + 7, + (4, 7, false), + (0, 32, true), + |v: u32| v as u128, + 0_i64, + 32_i64 + ); + run!( + i64, + 9, + 64, + prepare_i64_manual, + i64_manual_async_owned, + LaunchConfig1D::auto_shared(2, 64), + signed_keys(64, 576, 2), + -7, + (61, 64, false), + (61, 64, true), + |v: i64| ((v as u64) ^ (1_u64 << 63)) as u128, + begin, + end, + invalid + ); + run!( + u64, + 4, + 33, + prepare_u64_auto, + u64_auto_async_owned, + LaunchConfig1D::auto_shared(2, 33), + unsigned_keys(64, 132, 2), + 7, + (0, 64, true), + (0, 3, false), + |v: u64| v as u128 + ); + } + assert_eq!( + (cases, sorted, values, sources, reports), + (4, 3444, 3444, 1722, 320) + ); + println!( + "radix_sort_keys: {cases} cases, {sorted} sorted keys, {values} preserved values, {sources} source values, {reports} lane reports passed" + ); + Ok(()) +} diff --git a/rust/cuda-coop/tests/radix_sort_pairs/.gitignore b/rust/cuda-coop/tests/radix_sort_pairs/.gitignore new file mode 100644 index 000000000000..74bffa8af4a4 --- /dev/null +++ b/rust/cuda-coop/tests/radix_sort_pairs/.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/radix_sort_pairs/Cargo.lock b/rust/cuda-coop/tests/radix_sort_pairs/Cargo.lock new file mode 100644 index 000000000000..020847158adb --- /dev/null +++ b/rust/cuda-coop/tests/radix_sort_pairs/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-radix-sort-pairs-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/radix_sort_pairs/Cargo.toml b/rust/cuda-coop/tests/radix_sort_pairs/Cargo.toml new file mode 100644 index 000000000000..e644ae90aef2 --- /dev/null +++ b/rust/cuda-coop/tests/radix_sort_pairs/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-radix-sort-pairs-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/radix_sort_pairs/src/main.rs b/rust/cuda-coop/tests/radix_sort_pairs/src/main.rs new file mode 100644 index 000000000000..ad3633c98c0e --- /dev/null +++ b/rust/cuda-coop/tests/radix_sort_pairs/src/main.rs @@ -0,0 +1,442 @@ +// Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +// Kernels expose independent typed key/value outputs and preserved inputs. +#![allow(clippy::too_many_arguments)] + +use cuda_coop as coop; +use cuda_device::{kernel, launch_contract, thread}; +use cuda_host::cuda_module; + +#[cuda_module] +mod kernels { + use super::*; + + macro_rules! exercise { + ($n:literal, $ka:literal, $va:literal, $b:literal, + $first_options:expr, $second_options:expr, $manual:literal, + $keys:ident, $values:ident, $first_keys:ident, $first_values:ident, + $second_keys:ident, $second_values:ident, $preserved_keys:ident, + $preserved_values:ident $(, $mutated_keys:ident)?) => {{ + let block = coop::this_block(); + let offset = block.rank_at::() as usize * $b * $n; + let mut keys = coop::ThreadData::<_, $n, $ka>::new(); + let mut values = coop::ThreadData::<_, $n, $va>::new(); + coop::load(block, $keys, &mut keys, coop::LoadOptions::default().offset(offset)); + coop::load(block, $values, &mut values, coop::LoadOptions::default().offset(offset)); + #[allow(unused_mut)] + let (mut first_k, first_v) = + coop::radix_sort_pairs(block, &keys, &values, $first_options); + if $manual { + thread::sync_threads(); + coop::store(block, &mut $first_keys, &first_k, coop::StoreOptions::default().offset(offset)); + } + $( + // Observe the original result before changing one fresh slot. + // Both this payload and first_v remain live across scratch reuse. + first_k.write(0, -7_i64); + let _ = &$mutated_keys; + )? + let (second_k, second_v) = + coop::radix_sort_pairs(block, &keys, &values, $second_options); + if $manual { + thread::sync_threads(); + } + if !$manual { + coop::store(block, &mut $first_keys, &first_k, coop::StoreOptions::default().offset(offset)); + } + coop::store(block, &mut $first_values, &first_v, coop::StoreOptions::default().offset(offset)); + coop::store(block, &mut $second_keys, &second_k, coop::StoreOptions::default().offset(offset)); + coop::store(block, &mut $second_values, &second_v, coop::StoreOptions::default().offset(offset)); + coop::store(block, &mut $preserved_keys, &keys, coop::StoreOptions::default().offset(offset)); + coop::store(block, &mut $preserved_values, &values, coop::StoreOptions::default().offset(offset)); + $(coop::store(block, &mut $mutated_keys, &first_k, coop::StoreOptions::default().offset(offset));)? + }}; + } + + #[kernel] + #[launch_contract(domain = 3, block = (3, 5, 2), dynamic_shared = inferred)] + pub fn u32_f32_xyz( + keys: &[u32], + values: &[f32], + mut first_keys: coop::DisjointSlice>, + mut first_values: coop::DisjointSlice>, + mut second_keys: coop::DisjointSlice>, + mut second_values: coop::DisjointSlice>, + mut preserved_keys: coop::DisjointSlice>, + mut preserved_values: coop::DisjointSlice>, + begin: i64, + end: i64, + ) { + exercise!( + 4, + 16, + 1, + 30, + coop::RadixSortOptions::default(), + coop::RadixSortOptions::default() + .order(coop::Descending) + .begin_bit(begin) + .end_bit(end), + false, + keys, + values, + first_keys, + first_values, + second_keys, + second_values, + preserved_keys, + preserved_values + ); + } + + #[kernel] + #[launch_contract(domain = 1, block = (64, 1, 1), dynamic_shared = inferred)] + pub fn i64_f64_manual( + keys: &[i64], + values: &[f64], + mut first_keys: coop::DisjointSlice>, + mut first_values: coop::DisjointSlice>, + mut second_keys: coop::DisjointSlice>, + mut second_values: coop::DisjointSlice>, + mut preserved_keys: coop::DisjointSlice>, + mut preserved_values: coop::DisjointSlice>, + mut mutated_keys: coop::DisjointSlice>, + begin: i64, + end: i64, + invalid: u32, + ) { + let mut scratch = coop::TempStorage::new(); + if invalid != 0 { + let block = coop::this_block(); + let rank = thread::threadIdx_x() as usize; + let mut input_keys = coop::ThreadData::<_, 9, 1>::new(); + let mut input_values = coop::ThreadData::<_, 9, 32>::new(); + let mut i = 0; + while i < 9 { + if !(invalid == 1 && rank == 63 && i == 8) { + input_keys.write(i, keys[rank * 9 + i]); + } + if !(invalid == 2 && rank == 63 && i == 8) { + input_values.write(i, values[rank * 9 + i]); + } + i += 1; + } + let _ = coop::radix_sort_pairs( + block, + &input_keys, + &input_values, + coop::RadixSortOptions::default() + .begin_bit(begin) + .end_bit(end) + .temp_storage(&mut scratch), + ); + thread::sync_threads(); + // If preflight is lost, return an independent report and fail on + // the host. Never observe either invalid result to force a trap. + let mut report = coop::ThreadData::<_, 9>::new(); + let mut i = 0; + while i < 9 { + report.write(i, 99_i64); + i += 1; + } + coop::store(block, &mut mutated_keys, &report, Default::default()); + return; + } + exercise!( + 9, + 1, + 32, + 64, + coop::RadixSortOptions::default() + .begin_bit(begin) + .temp_storage(&mut scratch), + coop::RadixSortOptions::default() + .order(coop::Descending) + .begin_bit_const::<61>() + .end_bit(end) + .temp_storage(&mut scratch), + true, + keys, + values, + first_keys, + first_values, + second_keys, + second_values, + preserved_keys, + preserved_values, + mutated_keys + ); + } + + #[kernel] + #[launch_contract(domain = 1, block = (33, 1, 1), dynamic_shared = inferred)] + pub fn u64_i8_auto( + keys: &[u64], + values: &[i8], + mut first_keys: coop::DisjointSlice>, + mut first_values: coop::DisjointSlice>, + mut second_keys: coop::DisjointSlice>, + mut second_values: coop::DisjointSlice>, + mut preserved_keys: coop::DisjointSlice>, + mut preserved_values: coop::DisjointSlice>, + ) { + let mut scratch = coop::TempStorage::new().auto_sync(); + exercise!( + 1, + 32, + 1, + 33, + coop::RadixSortOptions::default() + .order(coop::Descending) + .temp_storage(&mut scratch), + coop::RadixSortOptions::default() + .begin_bit_const::<0>() + .end_bit_const::<3>() + .temp_storage(&mut scratch), + false, + keys, + values, + first_keys, + first_values, + second_keys, + second_values, + preserved_keys, + preserved_values + ); + } +} + +// Sorting original indices independently checks both stability and association. +fn oracle( + keys: &[K], + values: &[V], + tile: usize, + begin: u32, + end: u32, + descending: bool, + ordered: impl Fn(K) -> u128, +) -> (Vec, Vec) { + let mask = (1_u128 << (end - begin)) - 1; + let digit = |index: usize| (ordered(keys[index]) >> begin) & mask; + let mut indices: Vec = (0..keys.len()).collect(); + for block in indices.chunks_exact_mut(tile) { + block.sort_by(|&left, &right| { + if descending { + digit(right).cmp(&digit(left)) + } else { + digit(left).cmp(&digit(right)) + } + }); + } + ( + indices.iter().map(|&i| keys[i]).collect(), + indices.iter().map(|&i| values[i]).collect(), + ) +} + +fn signed_keys(tile: usize, blocks: usize) -> Vec { + (0..blocks * tile) + .map(|index| { + let j = index % tile; + let digit = ((5 * j + 3 * (index / tile)) % 8) as u64; + match j { + 0 => i64::MIN, + 1 => i64::MAX, + 2 => -1, + 3 => 0, + _ => ((digit << 61 | (index as u64 + 1)) ^ (1_u64 << 63)) as i64, + } + }) + .collect() +} + +fn unsigned_keys(width: u32, tile: usize, blocks: usize) -> Vec { + (0..blocks * tile) + .map(|index| { + let j = index % tile; + let digit = ((5 * j + 3 * (index / tile)) % 8) as u64; + let tag = index as u64 + 1; + match j { + 0 => 0, + 1 if width == 32 => u32::MAX as u64, + 1 => u64::MAX, + _ if width == 32 => (tag << 7) | (digit << 4) | (tag & 15), + _ => (((j & 1) as u64) << 63) | (tag << 3) | digit, + } + }) + .collect() +} + +fn f32_values(count: usize) -> Vec { + (0..count) + .map(|i| match i % 8 { + 0 => f32::from_bits(0x7fc0_0000 | (i as u32 + 1)), + 1 => f32::from_bits(0xffc0_0000 | (i as u32 + 1)), + 2 => 0.0, + 3 => -0.0, + 4 => f32::INFINITY, + 5 => f32::NEG_INFINITY, + _ => i as f32 + 0.25, + }) + .collect() +} + +fn f64_values(count: usize) -> Vec { + (0..count) + .map(|i| match i % 8 { + 0 => f64::from_bits(0x7ff8_0000_0000_0000 | (i as u64 + 1)), + 1 => f64::from_bits(0xfff8_0000_0000_0000 | (i as u64 + 1)), + 2 => 0.0, + 3 => -0.0, + 4 => f64::INFINITY, + 5 => f64::NEG_INFINITY, + _ => i as f64 + 0.25, + }) + .collect() +} + +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}; + let (invalid, begin, end) = match std::env::args().nth(1).as_deref() { + None => (0, 61_i64, 64_i64), + Some("missing-key") => (1, 61, 64), + Some("missing-value") => (2, 61, 64), + Some("negative-begin") => (3, i64::MIN, 64), + Some("empty") => (4, 61, 61), + Some("wide-end") => (5, 0, (1_i64 << 32) + 64), + Some(other) => panic!("unknown Radix Sort Pairs mode {other}"), + }; + init_device_contexts(0, 1)?; + let context = with_cuda_context(0, |context| context.clone())?; + let stream = context.default_stream(); + // SAFETY: typed declarations and embedded kernels are compiled together. + let module = unsafe { kernels::load_async(0)? }; + let (mut cases, mut sorted, mut preserved, mut sources, mut mutated) = (0, 0, 0, 0, 0); + macro_rules! run { + ($kt:ty, $vt:ty, $n:literal, $b:literal, $prepare:ident, $launch:ident, + $config:expr, $keys:expr, $values:expr, $first:expr, $second:expr, + $ordered:expr, $bits:expr, [$($argument:expr),*] $(, $mutation:ident)?) => {{ + let keys: Vec<$kt> = $keys; + let values: Vec<$vt> = $values; + let (b1, e1, d1) = $first; + let (b2, e2, d2) = $second; + let (first_k, first_v) = oracle(&keys, &values, $b * $n, b1, e1, d1, $ordered); + let (second_k, second_v) = oracle(&keys, &values, $b * $n, b2, e2, d2, $ordered); + let (whole_k, whole_v) = oracle(&keys, &values, $b * $n, 0, <$kt>::BITS, d2, $ordered); + let bits = $bits; + let value_bits = |v: &[$vt]| v.iter().copied().map(bits).collect::>(); + assert!(second_k != whole_k || value_bits(&second_v) != value_bits(&whole_v), "selected-digit oracle must distinguish whole-key order"); + let source_k = DeviceBuffer::from_host(&stream, &keys)?; + let source_v = DeviceBuffer::from_host(&stream, &values)?; + let out_k1 = DeviceBuffer::from_host(&stream, &vec![0 as $kt; keys.len()])?; + let out_v1 = DeviceBuffer::from_host(&stream, &vec![0 as $vt; keys.len()])?; + let out_k2 = DeviceBuffer::from_host(&stream, &vec![0 as $kt; keys.len()])?; + let out_v2 = DeviceBuffer::from_host(&stream, &vec![0 as $vt; keys.len()])?; + let keep_k = DeviceBuffer::from_host(&stream, &vec![0 as $kt; keys.len()])?; + let keep_v = DeviceBuffer::from_host(&stream, &vec![0 as $vt; keys.len()])?; + $(let $mutation = DeviceBuffer::from_host(&stream, &vec![0 as $kt; keys.len()])?;)? + stream.synchronize()?; + let plan = module.$prepare($config)?; + let (source_k, source_v, out_k1, out_v1, out_k2, out_v2, keep_k, keep_v $(, $mutation)?) = module.$launch( + &plan, source_k, source_v, out_k1, out_v1, out_k2, out_v2, keep_k, keep_v $(, $mutation)? $(, $argument)*).sync()?; + stream.synchronize()?; + if invalid != 0 { panic!("invalid Radix Sort Pairs mode {invalid} unexpectedly completed"); } + assert_eq!(out_k1.to_host_vec(&stream)?, first_k, stringify!($prepare)); + assert_eq!(value_bits(&out_v1.to_host_vec(&stream)?), value_bits(&first_v), stringify!($prepare)); + assert_eq!(out_k2.to_host_vec(&stream)?, second_k, stringify!($prepare)); + assert_eq!(value_bits(&out_v2.to_host_vec(&stream)?), value_bits(&second_v), stringify!($prepare)); + assert_eq!(keep_k.to_host_vec(&stream)?, keys, stringify!($prepare)); + assert_eq!(value_bits(&keep_v.to_host_vec(&stream)?), value_bits(&values), stringify!($prepare)); + assert_eq!(source_k.to_host_vec(&stream)?, keys, stringify!($prepare)); + assert_eq!(value_bits(&source_v.to_host_vec(&stream)?), value_bits(&values), stringify!($prepare)); + $(let mut expected_mutated = first_k; + for key in expected_mutated.iter_mut().step_by($n) { *key = -7_i64; } + assert_eq!($mutation.to_host_vec(&stream)?, expected_mutated); + mutated += keys.len();)? + cases += 1; sorted += 4 * keys.len(); preserved += 2 * keys.len(); sources += 2 * keys.len(); + }}; + } + if invalid != 0 { + run!( + i64, + f64, + 9, + 64, + prepare_i64_f64_manual, + i64_f64_manual_async_owned, + LaunchConfig1D::auto_shared(1, 64), + signed_keys(576, 1), + f64_values(576), + (61, 64, false), + (61, 64, true), + |v: i64| ((v as u64) ^ (1_u64 << 63)) as u128, + f64::to_bits, + [begin, end, invalid], + changed + ); + } else { + run!( + u32, + f32, + 4, + 30, + prepare_u32_f32_xyz, + u32_f32_xyz_async_owned, + LaunchConfig3D::auto_shared((2, 1, 1), (3, 5, 2)), + unsigned_keys(32, 120, 2) + .into_iter() + .map(|v| v as u32) + .collect(), + f32_values(240), + (0, 32, false), + (4, 7, true), + |v: u32| v as u128, + f32::to_bits, + [4_i64, 7_i64] + ); + run!( + i64, + f64, + 9, + 64, + prepare_i64_f64_manual, + i64_f64_manual_async_owned, + LaunchConfig1D::auto_shared(2, 64), + signed_keys(576, 2), + f64_values(1152), + (61, 64, false), + (61, 64, true), + |v: i64| ((v as u64) ^ (1_u64 << 63)) as u128, + f64::to_bits, + [begin, end, invalid], + changed + ); + run!( + u64, + i8, + 1, + 33, + prepare_u64_i8_auto, + u64_i8_auto_async_owned, + LaunchConfig1D::auto_shared(2, 33), + unsigned_keys(64, 33, 2), + (0..66).map(|i| i as i8 - 33).collect(), + (0, 64, true), + (0, 3, false), + |v: u64| v as u128, + |v: i8| v, + [] + ); + } + assert_eq!( + (cases, sorted, preserved, sources, mutated), + (3, 5832, 2916, 2916, 1152) + ); + println!( + "radix_sort_pairs: {cases} cases, {sorted} sorted key/value observations, {preserved} preserved values, {sources} source values, {mutated} retained mutated keys passed" + ); + Ok(()) +} From b4b7061df433bf9dc10dea1e8ea4456394c7d72d Mon Sep 17 00:00:00 2001 From: Trent Nelson Date: Sat, 24 Oct 2026 17:44:15 -0700 Subject: [PATCH 10/10] Add Rust Merge Sort Keys and Pairs with supported comparators Provide Block and Warp key and pair sorting with partial inputs. Include the supported stateless full-input i32 Block comparator at its owning family and retain the checked callback ABI. Signed-off-by: Trent Nelson --- docs/rust/index.rst | 16 +- docs/rust/merge_sort_keys.rst | 116 ++ docs/rust/merge_sort_pairs.rst | 84 ++ rust/cuda-coop-codegen/src/lib.rs | 42 +- rust/cuda-coop-codegen/src/merge_sort.rs | 637 +++++++++++ rust/cuda-coop-codegen/src/oxide.rs | 545 ++++++++- rust/cuda-coop-codegen/tests/native.rs | 275 ++++- rust/cuda-coop/build.rs | 52 +- .../examples/merge_sort_keys/.gitignore | 7 + .../examples/merge_sort_keys/Cargo.lock | 1019 +++++++++++++++++ .../examples/merge_sort_keys/Cargo.toml | 19 + .../merge_sort_keys/helper/Cargo.toml | 11 + .../merge_sort_keys/helper/src/lib.rs | 61 + .../examples/merge_sort_keys/src/main.rs | 290 +++++ .../examples/merge_sort_pairs/.gitignore | 7 + .../examples/merge_sort_pairs/Cargo.lock | 1019 +++++++++++++++++ .../examples/merge_sort_pairs/Cargo.toml | 19 + .../merge_sort_pairs/helper/Cargo.toml | 11 + .../merge_sort_pairs/helper/src/lib.rs | 72 ++ .../examples/merge_sort_pairs/src/main.rs | 248 ++++ rust/cuda-coop/src/lib.rs | 10 + rust/cuda-coop/src/merge_sort.rs | 571 +++++++++ rust/cuda-coop/src/merge_sort/boundary.rs | 217 ++++ rust/cuda-coop/src/radix_sort.rs | 4 +- rust/cuda-coop/tests/check_load_store.py | 216 ++++ .../tests/merge_sort_keys/.gitignore | 7 + .../tests/merge_sort_keys/Cargo.lock | 1011 ++++++++++++++++ .../tests/merge_sort_keys/Cargo.toml | 17 + .../tests/merge_sort_keys/src/main.rs | 645 +++++++++++ .../tests/merge_sort_pairs/.gitignore | 7 + .../tests/merge_sort_pairs/Cargo.lock | 1011 ++++++++++++++++ .../tests/merge_sort_pairs/Cargo.toml | 17 + .../tests/merge_sort_pairs/src/main.rs | 560 +++++++++ 33 files changed, 8804 insertions(+), 39 deletions(-) create mode 100644 docs/rust/merge_sort_keys.rst create mode 100644 docs/rust/merge_sort_pairs.rst create mode 100644 rust/cuda-coop-codegen/src/merge_sort.rs create mode 100644 rust/cuda-coop/examples/merge_sort_keys/.gitignore create mode 100644 rust/cuda-coop/examples/merge_sort_keys/Cargo.lock create mode 100644 rust/cuda-coop/examples/merge_sort_keys/Cargo.toml create mode 100644 rust/cuda-coop/examples/merge_sort_keys/helper/Cargo.toml create mode 100644 rust/cuda-coop/examples/merge_sort_keys/helper/src/lib.rs create mode 100644 rust/cuda-coop/examples/merge_sort_keys/src/main.rs create mode 100644 rust/cuda-coop/examples/merge_sort_pairs/.gitignore create mode 100644 rust/cuda-coop/examples/merge_sort_pairs/Cargo.lock create mode 100644 rust/cuda-coop/examples/merge_sort_pairs/Cargo.toml create mode 100644 rust/cuda-coop/examples/merge_sort_pairs/helper/Cargo.toml create mode 100644 rust/cuda-coop/examples/merge_sort_pairs/helper/src/lib.rs create mode 100644 rust/cuda-coop/examples/merge_sort_pairs/src/main.rs create mode 100644 rust/cuda-coop/src/merge_sort.rs create mode 100644 rust/cuda-coop/src/merge_sort/boundary.rs create mode 100644 rust/cuda-coop/tests/merge_sort_keys/.gitignore create mode 100644 rust/cuda-coop/tests/merge_sort_keys/Cargo.lock create mode 100644 rust/cuda-coop/tests/merge_sort_keys/Cargo.toml create mode 100644 rust/cuda-coop/tests/merge_sort_keys/src/main.rs create mode 100644 rust/cuda-coop/tests/merge_sort_pairs/.gitignore create mode 100644 rust/cuda-coop/tests/merge_sort_pairs/Cargo.lock create mode 100644 rust/cuda-coop/tests/merge_sort_pairs/Cargo.toml create mode 100644 rust/cuda-coop/tests/merge_sort_pairs/src/main.rs diff --git a/docs/rust/index.rst b/docs/rust/index.rst index 23e13cb5b514..e1fe66597cb7 100644 --- a/docs/rust/index.rst +++ b/docs/rust/index.rst @@ -46,7 +46,8 @@ group's tile, and must be uniform across that group. The public operations include :doc:`load_store`, :doc:`scan`, :doc:`reduce`, :doc:`reduce_batched`, :doc:`exchange`, :doc:`shuffle`, :doc:`adjacent_difference`, :doc:`discontinuity`, :doc:`histogram`, -:doc:`radix_sort_keys`, :doc:`radix_sort_pairs` and :doc:`radix_rank_keys`. +:doc:`radix_sort_keys`, :doc:`radix_sort_pairs`, :doc:`radix_rank_keys`, +:doc:`merge_sort_keys` and :doc:`merge_sort_pairs`. Scan returns fresh per-group prefixes; Reduce returns a total only at each group's root. Batched reduction returns a fresh payload whose extent and valid slots follow its Warp width and output layout. @@ -57,11 +58,14 @@ Histogram returns fresh striped counters with an independent type and extent. Radix Sort Keys stably sorts integer keys into a fresh blocked payload. Radix Sort Pairs carries associated numeric values through that stable order and returns two fresh payloads. Radix Rank Keys returns each input's stable -selected-digit position as a fresh ``i32`` payload. Histogram, Radix Sort and -Radix Rank require generic native-registration support beyond the currently -pinned Oxide revision; their guides explain the compiler requirement. +selected-digit position as a fresh ``i32`` payload. Merge Sort Keys sorts numeric +keys or a selected prefix into a fresh payload. Merge Sort Pairs carries +associated numeric values into independent fresh payloads. Histogram, Radix Sort, +Radix Rank and Merge Sort Pairs require generic native-registration support +beyond the currently pinned Oxide revision; their guides explain the compiler requirement. Scan, Reduce, batched reduction, Exchange, Shuffle, Adjacent Difference, -Discontinuity, Histogram, Radix Sort and Radix Rank preserve their inputs. +Discontinuity, Histogram, Radix Sort, Radix Rank and Merge Sort preserve their +inputs. Additional primitive families are under development. The crate is unpublished and its interface can change. @@ -141,6 +145,8 @@ Read next radix_sort_keys radix_sort_pairs radix_rank_keys + merge_sort_keys + merge_sort_pairs provider Generate the API reference with ``cargo doc --workspace --no-deps`` from diff --git a/docs/rust/merge_sort_keys.rst b/docs/rust/merge_sort_keys.rst new file mode 100644 index 000000000000..f947b98ea178 --- /dev/null +++ b/docs/rust/merge_sort_keys.rst @@ -0,0 +1,116 @@ +.. Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. +.. SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +Merge-sorting numeric keys +=========================== + +``merge_sort_keys`` sorts a group's keys into a fresh blocked payload: + +.. code-block:: rust + + let sorted = coop::merge_sort_keys(block, &keys, Default::default()); + +The default sorts every key in ascending order. The result retains the input's +numeric type, item count and minimum alignment. Input values and validity remain +unchanged. Merge Sort supports all ten ``Numeric`` types, including ``f32`` and +``f64``. Equal keys have no stable-tie guarantee; floating keys must satisfy the +comparison's strict weak ordering requirements. No NaN total order is supplied. + +Stateless comparators +--------------------- + +Full-input Block ``i32`` Keys also accepts function items and zero-capture +``Fn(i32, i32) -> bool + Copy`` closures through ``order``: + +.. code-block:: rust + + let descending = coop::merge_sort_keys( + block, &keys, + coop::MergeSortOptions::default().order(|left, right| left > right), + ); + +True means that the left key precedes the right. The comparator must define a +strict weak ordering; ties remain unordered. Callable sorting keeps the same +fresh result, input validity and temporary-storage rules. Captures, function +pointers, partial tiles, Warp groups and Pairs are rejected at compile time for +callable orders. Closures that call methods on their keys may need explicit +``i32`` parameter types. +Built-in ascending and descending orders retain their existing type and group +support. + +A complete program +------------------ + +This helper keeps the item count and alignment generic. It returns a full +ascending result, a descending prefix and the original payload. The partial +operation takes a group-uniform signed count and a sentinel of the key type. +The second helper forwards ``MergeSortOptions`` under the existing nongeneric +``MergeSortOrder`` bound, preserving the actual comparator and selected storage: + +.. literalinclude:: ../../rust/cuda-coop/examples/merge_sort_keys/helper/src/lib.rs + :language: rust + :start-after: // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +The host instantiates one and four items per thread on two 4-by-4-by-4 Blocks. +It uses ``i64`` keys with at least 32-byte alignment, checks both sorted results, +verifies input preservation and result validity, and mutates one result while +retaining the other. A second kernel uses three 128-thread Blocks with ``i32`` +keys, including duplicates and both integer extremes. For one and four items +per thread it combines built-in ascending, an even-first function item and a +descending closure. Ascending and even-first reuse auto-sync scratch; the +descending closure uses implicit storage. All three results remain live, and +the host checks every result, validity bit and preserved input: + +.. literalinclude:: ../../rust/cuda-coop/examples/merge_sort_keys/src/main.rs + :language: rust + :start-after: // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +With the matching Oxide and ``1.99.1`` installed, run from a complete +CCCL checkout: + +.. code-block:: shell + + cd rust/cuda-coop/examples/merge_sort_keys + 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. ``auto_shared`` uses the compiler's +complete temporary-storage plan. + +Counts and result validity +-------------------------- + +``MergeSortOptions::default().valid_items(count, sentinel)`` sorts a blocked +prefix. Count and sentinel are selected together and must be uniform across the +group. The sentinel must compare after every valid key in the selected order: +for example, ``i64::MAX`` for ascending or ``i64::MIN`` for descending keys. +``order(coop::Descending)`` preserves both controls, and the builders can be +applied in either order. + +Count is ``i64`` and must be in ``0..=group_size * ITEMS_PER_THREAD``. It is +checked before narrowing to CUB's integer type. Every input slot must be +initialized, including slots outside the prefix and all slots at count zero. +Invalid counts or unreadable inputs trap before native numeric reads. + +Only the result's blocked prefix is initialized. ``get`` returns ``None`` for +suffix slots, and Store can write the valid prefix without touching the rest +of the destination. Count zero still participates and returns an empty result. + +Groups and temporary storage +---------------------------- + +Every member of the selected group participates. Complete XYZ Blocks require a +power-of-two thread count. Physical Warps and logical widths 1, 2, 4, 8, 16 and +32 are supported with complete physical warps in the enclosing Block; that +Block's total need not be a power of two. Per-thread item counts are positive. +Input and result slots use blocked order within each group. + +Implicit storage is compiler-planned for each participating group. Blocks can +also borrow existing ``TempStorage`` through ``temp_storage(&mut scratch)``. +Manual storage requires Block completion before conflicting reuse; auto-sync +storage supplies its trailing barrier. Explicit storage is not a Warp option. +Internal CUB synchronization does not complete unrelated global-memory accesses. + +See :doc:`payload`, :doc:`groups` and :doc:`load_store` for payload initialization, +participation and storage reuse. :doc:`radix_sort_keys` provides stable ordering +of selected integer-key digits. diff --git a/docs/rust/merge_sort_pairs.rst b/docs/rust/merge_sort_pairs.rst new file mode 100644 index 000000000000..35c145fd317e --- /dev/null +++ b/docs/rust/merge_sort_pairs.rst @@ -0,0 +1,84 @@ +.. Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. +.. SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +Merge-sorting key-value pairs +=============================== + +``merge_sort_pairs`` sorts keys and carries each associated value into two fresh +blocked payloads: + +.. code-block:: rust + + let (sorted_keys, sorted_values) = + coop::merge_sort_pairs(block, &keys, &values, Default::default()); + +Keys and values independently support the ten ``Numeric`` types. They have the +same per-thread item count and may have different minimum alignments. Each +result retains its corresponding input's type and alignment. Both inputs and +their validity remain unchanged. The default sorts the full group in ascending +key order with implicit temporary storage. + +Equal keys have no stable-tie guarantee. Values stay associated with their keys, +including floating-point value bit patterns. Floating keys must satisfy the +comparison's strict weak ordering requirements; no NaN total order is supplied. +Use :doc:`radix_sort_pairs` for stable selected-digit sorting of integer keys. + +A complete program +------------------ + +This helper preserves generic item counts and independent key/value alignments. +It returns a full ascending pair, a descending prefix pair and both inputs: + +.. literalinclude:: ../../rust/cuda-coop/examples/merge_sort_pairs/helper/src/lib.rs + :language: rust + :start-after: // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +The host instantiates one and four items per thread on two 4-by-4-by-4 Blocks. +It uses ``i32`` keys with at least 16-byte alignment and ``f64`` values with at +least 32-byte alignment. Its checks use key ordering and pair-bit multisets so +equal-key permutations are allowed, while association, preserved inputs, prefix +validity and fresh result mutations are observed independently: + +.. literalinclude:: ../../rust/cuda-coop/examples/merge_sort_pairs/src/main.rs + :language: rust + :start-after: // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +With the pinned Oxide revision and Rust ``1.99.1`` +installed, run from a complete CCCL checkout: + +.. code-block:: shell + + cd rust/cuda-coop/examples/merge_sort_pairs + 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. ``auto_shared`` uses the compiler's +complete temporary-storage plan. + +Counts, validity and storage +---------------------------- + +``MergeSortOptions::default().valid_items(count, sentinel)`` selects a blocked +prefix of pairs. Count is a group-uniform ``i64`` in +``0..=group_size * ITEMS_PER_THREAD``. The sentinel has the **key** type and must +compare after every valid key in the chosen order. There is no value sentinel. +The count and sentinel are selected together; ``order`` and ``temp_storage`` +retain them regardless of builder order. + +Every slot in both input payloads must be initialized, even outside the prefix +and when count is zero. Native preflight checks both validity maps before +reading any numeric input. Invalid counts or unreadable inputs trap. Both fresh +results have exactly the prefix initialized; ``get`` returns ``None`` for either +suffix, and a prefix Store leaves destination suffix slots untouched. + +Complete XYZ Blocks require a power-of-two thread count. Physical Warps and +logical widths 1, 2, 4, 8, 16 and 32 use complete physical warps in the enclosing +Block. Implicit storage is planned separately for participating groups. Blocks +can borrow existing ``TempStorage`` with ``temp_storage(&mut scratch)``; manual +reuse requires Block completion, while auto-sync storage supplies its trailing +barrier. Explicit storage is not a Warp option. Both result payloads remain +independent of scratch reuse. Internal CUB synchronization does not complete +unrelated global-memory accesses. + +See :doc:`merge_sort_keys`, :doc:`payload`, :doc:`groups` and :doc:`load_store` +for shared ordering options, initialization, participation and storage rules. diff --git a/rust/cuda-coop-codegen/src/lib.rs b/rust/cuda-coop-codegen/src/lib.rs index acb62afb4f3a..a92b7241aa8a 100644 --- a/rust/cuda-coop-codegen/src/lib.rs +++ b/rust/cuda-coop-codegen/src/lib.rs @@ -5,7 +5,7 @@ //! //! This private development crate compiles Load/Store, Scan, Reduce, batched //! Reduce, Exchange, Shuffle, Adjacent Difference, Discontinuity, Histogram -//! and Radix Sort Keys/Pairs and Rank Keys with NVRTC. +//! and Radix Sort Keys/Pairs, Rank Keys and Merge Sort Keys/Pairs 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. @@ -60,8 +60,9 @@ //! uniform bit range, preserving input and producing fresh blocked keys. //! [`RadixSortPairsRequest`] carries independently typed values with each key //! into fresh paired outputs. [`RadixRankKeysRequest`] returns stable selected-digit -//! i32 ranks in original input-slot order. [`Provider::compile_native`] batches these families -//! in one native compilation. +//! i32 ranks in original input-slot order. [`MergeSortKeysRequest`] sorts numeric +//! Block/Warp keys with full or prefix-valid fresh output. [`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 @@ -95,6 +96,7 @@ mod discontinuity; mod exchange; mod headers; mod histogram; +mod merge_sort; mod radix_sort; mod reduce; mod reduce_batched; @@ -104,6 +106,9 @@ pub use adjacent_difference::{AdjacentDifferenceRecipe, AdjacentDifferenceReques pub use discontinuity::{DiscontinuityRecipe, DiscontinuityRequest}; pub use exchange::{ExchangeMode, ExchangeRecipe, ExchangeRequest}; pub use histogram::{HistogramAlgorithm, HistogramRecipe, HistogramRequest}; +pub use merge_sort::{ + MergeSortKeysRecipe, MergeSortKeysRequest, MergeSortPairsRecipe, MergeSortPairsRequest, +}; pub use radix_sort::{ RadixRankKeysRecipe, RadixRankKeysRequest, RadixSortKeysRecipe, RadixSortKeysRequest, RadixSortPairsRecipe, RadixSortPairsRequest, @@ -514,6 +519,10 @@ pub enum NativeRequest { RadixSortPairs(RadixSortPairsRequest), /// Stable selected-digit Block ranks with preserved keys and fresh i32 output. RadixRankKeys(RadixRankKeysRequest), + /// Numeric Block merge sort preserving inputs with fresh full/partial output. + MergeSortKeys(MergeSortKeysRequest), + /// Numeric Block/Warp merge sort preserving associated key/value inputs. + MergeSortPairs(MergeSortPairsRequest), } impl NativeRequest { @@ -531,6 +540,8 @@ impl NativeRequest { Self::RadixSortKeys(request) => request.block_shape(), Self::RadixSortPairs(request) => request.block_shape(), Self::RadixRankKeys(request) => request.block_shape(), + Self::MergeSortKeys(request) => request.block_shape(), + Self::MergeSortPairs(request) => request.block_shape(), } } @@ -548,6 +559,8 @@ impl NativeRequest { Self::RadixSortKeys(request) => request.source(symbol), Self::RadixSortPairs(request) => request.source(symbol), Self::RadixRankKeys(request) => request.source(symbol), + Self::MergeSortKeys(request) => request.source(symbol), + Self::MergeSortPairs(request) => request.source(symbol), } } @@ -565,6 +578,8 @@ impl NativeRequest { Self::RadixSortKeys(request) => request.recipe.token().into(), Self::RadixSortPairs(request) => request.recipe.token(request.value_dtype), Self::RadixRankKeys(request) => request.recipe.token(), + Self::MergeSortKeys(request) => request.recipe.token().into(), + Self::MergeSortPairs(request) => request.recipe.token(request.value_dtype), } } @@ -581,7 +596,9 @@ impl NativeRequest { | Self::Histogram(_) | Self::RadixSortKeys(_) | Self::RadixSortPairs(_) - | Self::RadixRankKeys(_) => false, + | Self::RadixRankKeys(_) + | Self::MergeSortKeys(_) + | Self::MergeSortPairs(_) => false, } } @@ -601,7 +618,9 @@ impl NativeRequest { | Self::Histogram(_) | Self::RadixSortKeys(_) | Self::RadixSortPairs(_) - | Self::RadixRankKeys(_) => false, + | Self::RadixRankKeys(_) + | Self::MergeSortKeys(_) + | Self::MergeSortPairs(_) => false, } } } @@ -628,6 +647,8 @@ pub struct Layout { /// and [`RadixSortKeysRequest`]. [`RadixSortPairsRequest`] returns fresh keys and /// independently typed associated values while preserving both inputs. /// [`RadixRankKeysRequest`] returns fresh i32 ranks independently of key type. +/// [`MergeSortKeysRequest`] sorts Block/Warp numeric keys with full or prefix-valid output. +/// [`MergeSortPairsRequest`] carries independently typed associated values. /// /// 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. @@ -869,6 +890,9 @@ impl Provider { (NativeRequest::Reduce(reduce), Some(callback)) => { reduce.source_with_callback(symbol, Some(&callback.symbol)) } + (NativeRequest::MergeSortKeys(sort), Some(callback)) => { + sort.source_with_callback(symbol, Some(&callback.symbol)) + } (_, None) => request.source(symbol), _ => unreachable!("callback recipes are validated before native compilation"), }; @@ -880,10 +904,10 @@ impl Provider { 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 token = if callback.is_some() { - "reduce_callable".to_owned() - } else { - request.token() + let token = match (request, callback) { + (NativeRequest::Reduce(_), Some(_)) => "reduce_callable".to_owned(), + (NativeRequest::MergeSortKeys(_), Some(_)) => "merge_sort_keys_callable".to_owned(), + _ => request.token(), }; let symbol = format!("cuda_coop_{token}_{suffix}"); source.push_str(&source_for(&symbol)); diff --git a/rust/cuda-coop-codegen/src/merge_sort.rs b/rust/cuda-coop-codegen/src/merge_sort.rs new file mode 100644 index 000000000000..2141501df2a9 --- /dev/null +++ b/rust/cuda-coop-codegen/src/merge_sort.rs @@ -0,0 +1,637 @@ +// Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +//! Validated Block/Warp merge sorting with preserved inputs and fresh outputs. + +use crate::{Dtype, Error, Result}; + +/// Built-in ordering and full/partial output selection for Merge Sort. +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub enum MergeSortKeysRecipe { + /// Sort the complete tile in ascending order. + AscendingFull, + /// Sort the complete tile in descending order. + DescendingFull, + /// Sort a uniform valid prefix in ascending order. + AscendingPartial, + /// Sort a uniform valid prefix in descending order. + DescendingPartial, +} + +impl MergeSortKeysRecipe { + pub(crate) fn token(self) -> &'static str { + match self { + Self::AscendingFull => "merge_sort_keys_ascending_full_validity_bits_v1", + Self::DescendingFull => "merge_sort_keys_descending_full_validity_bits_v1", + Self::AscendingPartial => "merge_sort_keys_ascending_partial_validity_bits_v1", + Self::DescendingPartial => "merge_sort_keys_descending_partial_validity_bits_v1", + } + } + + pub(crate) fn parse(operation: &str, algorithm: &str) -> Result { + if !matches!(algorithm, "block" | "warp") { + return Err(Error::message( + "Merge Sort Keys requires a Block or Warp recipe", + )); + } + match operation { + "merge_sort_keys_ascending_full_validity_bits_v1" => Ok(Self::AscendingFull), + "merge_sort_keys_descending_full_validity_bits_v1" => Ok(Self::DescendingFull), + "merge_sort_keys_ascending_partial_validity_bits_v1" => Ok(Self::AscendingPartial), + "merge_sort_keys_descending_partial_validity_bits_v1" => Ok(Self::DescendingPartial), + _ => Err(Error::message("unsupported Merge Sort Keys operation")), + } + } + + pub(crate) fn partial(self) -> bool { + matches!(self, Self::AscendingPartial | Self::DescendingPartial) + } +} + +/// Validated numeric Block/Warp merge sort with fresh same-extent output. +/// +/// Full native calls use `(const T* input, T* output, u8* scratch, +/// const u8* input_bits, u8* output_bits)`. Partial calls append +/// `(i64 valid_items, T sentinel)`. These controls are uniform within the group; +/// `0 <= valid_items <= group_threads*items` is checked before narrowing or +/// numeric input reads. The sentinel must be ordered after every valid key. +/// It has exactly the key type. Floating comparisons must satisfy strict weak +/// ordering; no NaN total order or stable ordering of equal keys is promised. +/// +/// Every lane supplies `items` initialized keys and `ceil(items/8)` packed, +/// low-bit-first validity bytes. All input bits are checked even for a zero +/// count. Input keys and masks are preserved; sorting uses private copies. +/// Output keys and masks are disjoint from input and scratch. Full calls write +/// all outputs; partial calls write only the blocked valid prefix. Other output +/// bits, including padding bits, are cleared, and invalid keys are not written. +/// Block rank is `x + X*(y + Y*z)`; Warp rank is that value modulo its width. +/// Slot `i` has group index `rank*items+i`. +/// +/// Every member of the power-of-two Block or aligned power-of-two Warp group +/// participates in converged control flow, including at count zero. Warp calls +/// require complete physical warps in the enclosing Block; sibling groups may +/// have different controls. All pointers remain live with their +/// stated capacities and natural alignment. Scratch uses the complete queried +/// CUB TempStorage, including its padding element. The caller synchronizes +/// members before scratch reuse; the provider adds no trailing barrier. Warp +/// callers supply an independent scratch replica for each actual group. +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub struct MergeSortKeysRequest { + pub(crate) dtype: Dtype, + pub(crate) recipe: MergeSortKeysRecipe, + block: [u32; 3], + items: u32, + width: Option, +} + +impl MergeSortKeysRequest { + /// Validate built-in merge sorting for a complete XYZ Block. + /// + /// `dtype` is any supported numeric type, `recipe` selects ordering and + /// full/partial validity, and `items` is the positive per-lane extent. + /// + /// # Errors + /// Rejects zero/overflowing axes, a non-power-of-two thread count or more + /// than 1024 threads, zero items, and a tile whose padded extent exceeds + /// CUB's signed-int template capacity. + pub fn block( + dtype: Dtype, + recipe: MergeSortKeysRecipe, + block: [u32; 3], + items: u32, + ) -> Result { + Self::validate(dtype, recipe, block, items, None) + } + + /// Validate an aligned Warp group within complete physical warps. + /// + /// Widths 1/2/4/8/16/32 are supported. The exact enclosing XYZ Block need + /// not have a power-of-two total. The caller owns scratch replication and + /// completion; this request does not establish participation. + /// + /// # Errors + /// Rejects invalid axes, more than 1024 threads, incomplete physical warps, + /// unsupported widths, zero items, and a padded group tile exceeding CUB's + /// signed-int template capacity. + pub fn warp( + dtype: Dtype, + recipe: MergeSortKeysRecipe, + block: [u32; 3], + width: u32, + items: u32, + ) -> Result { + Self::validate(dtype, recipe, block, items, Some(width)) + } + + fn validate( + dtype: Dtype, + recipe: MergeSortKeysRecipe, + 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( + "Merge Sort Keys requires positive axes and at most 1024 block threads", + ) + })?; + let group_threads = if let Some(width) = width { + if !matches!(width, 1 | 2 | 4 | 8 | 16 | 32) || threads % 32 != 0 { + return Err(Error::message( + "Warp Merge Sort Keys requires a power-of-two width in 1..=32 and complete physical warps", + )); + } + width + } else { + if !threads.is_power_of_two() { + return Err(Error::message( + "Merge Sort Keys requires a power-of-two Block", + )); + } + threads + }; + if items == 0 + || group_threads + .checked_mul(items) + .and_then(|tile| tile.checked_add(1)) + .filter(|padded| *padded <= i32::MAX as u32) + .is_none() + { + return Err(Error::message( + "Merge Sort Keys requires positive items and a padded tile fitting CUB int capacity", + )); + } + Ok(Self { + dtype, + recipe, + block, + items, + width, + }) + } + + pub(crate) fn block_shape(self) -> [u32; 3] { + self.block + } + + 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 tile = self.width.unwrap_or(x * y * z) * items; + let (header, primitive, rank) = if let Some(width) = self.width { + ( + "warp/warp_merge_sort.cuh", + format!("cub::WarpMergeSort"), + format!("(threadIdx.x + {x} * (threadIdx.y + {y} * threadIdx.z)) % {width}"), + ) + } else { + ( + "block/block_merge_sort.cuh", + format!("cub::BlockMergeSort"), + format!("threadIdx.x + {x} * (threadIdx.y + {y} * threadIdx.z)"), + ) + }; + let bytes = items / 8 + u32::from(items % 8 != 0); + let comparison = match self.recipe { + MergeSortKeysRecipe::AscendingFull | MergeSortKeysRecipe::AscendingPartial => "<", + MergeSortKeysRecipe::DescendingFull | MergeSortKeysRecipe::DescendingPartial => ">", + }; + let (mut parameters, bounds, arguments, count) = if self.recipe.partial() { + ( + format!(", long long valid_items, {dtype} sentinel"), + format!("if (valid_items < 0 || valid_items > {tile}) __trap();"), + ", static_cast(valid_items), sentinel", + "valid_items".into(), + ) + } else { + (String::new(), String::new(), "", tile.to_string()) + }; + let declaration = callback.map_or_else(String::new, |name| { + format!("extern \"C\" __device__ int {name}(const void*, int, int);\n") + }); + let environment = if callback.is_some() { + parameters.push_str(", const void* environment"); + "environment" + } else { + "" + }; + let comparator = callback.map_or_else( + || format!( + "struct compare {{\n __device__ __forceinline__ bool operator()(value_type lhs, value_type rhs) const {{\n return lhs {comparison} rhs;\n }}\n}};" + ), + |name| format!( + "struct compare {{ const void* environment;\n __device__ __forceinline__ bool operator()(int lhs, int rhs) const {{\n return {name}(environment, lhs, rhs) != 0;\n }}\n}};" + ), + ); + format!( + r#" +#include +#include +{declaration}namespace {symbol}_detail {{ +using value_type = {dtype}; +using primitive = {primitive}; +using scratch_storage = primitive::TempStorage; +{comparator} +}} +extern "C" __device__ void {symbol}(const {dtype}* input, {dtype}* output, + unsigned char* scratch, const unsigned char* input_bits, unsigned char* output_bits{parameters}) {{ + using namespace {symbol}_detail; + {bounds} +#pragma unroll + for (int i = 0; i < {bytes}; ++i) output_bits[i] = 0; +#pragma unroll + for (int i = 0; i < {items}; ++i) + if ((input_bits[i / 8] & (1u << (i % 8))) == 0) __trap(); + value_type local_keys[{items}]; +#pragma unroll + for (int i = 0; i < {items}; ++i) local_keys[i] = input[i]; + primitive operation{{*reinterpret_cast(scratch)}}; + operation.Sort(local_keys, compare{{{environment}}}{arguments}); + const int rank = static_cast({rank}); +#pragma unroll + for (int i = 0; i < {items}; ++i) {{ + if (rank * {items} + i < {count}) {{ + output[i] = local_keys[i]; + output_bits[i / 8] |= static_cast(1u << (i % 8)); + }} + }} +}} +"# + ) + } +} + +/// Built-in ordering and full/partial selection for associated key/value pairs. +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub enum MergeSortPairsRecipe { + /// Sort all pairs in ascending key order. + AscendingFull, + /// Sort all pairs in descending key order. + DescendingFull, + /// Sort a uniform prefix of pairs in ascending key order. + AscendingPartial, + /// Sort a uniform prefix of pairs in descending key order. + DescendingPartial, +} + +impl MergeSortPairsRecipe { + fn keys(self) -> MergeSortKeysRecipe { + match self { + Self::AscendingFull => MergeSortKeysRecipe::AscendingFull, + Self::DescendingFull => MergeSortKeysRecipe::DescendingFull, + Self::AscendingPartial => MergeSortKeysRecipe::AscendingPartial, + Self::DescendingPartial => MergeSortKeysRecipe::DescendingPartial, + } + } + pub(crate) fn partial(self) -> bool { + self.keys().partial() + } + pub(crate) fn token(self, value: Dtype) -> String { + let form = match self { + Self::AscendingFull => "ascending_full", + Self::DescendingFull => "descending_full", + Self::AscendingPartial => "ascending_partial", + Self::DescendingPartial => "descending_partial", + }; + let value = match value { + Dtype::I8 => "i8", + Dtype::U8 => "u8", + Dtype::I16 => "i16", + Dtype::U16 => "u16", + Dtype::I32 => "i32", + Dtype::U32 => "u32", + Dtype::I64 => "i64", + Dtype::U64 => "u64", + Dtype::F32 => "f32", + Dtype::F64 => "f64", + }; + format!("merge_sort_pairs_{form}_{value}_validity_bits_v1") + } + pub(crate) fn parse(operation: &str, algorithm: &str) -> Result<(Self, Dtype)> { + if !matches!(algorithm, "block" | "warp") { + return Err(Error::message( + "Merge Sort Pairs requires a Block or Warp recipe", + )); + } + let unsupported = || Error::message("unsupported Merge Sort Pairs operation"); + let operation = operation + .strip_prefix("merge_sort_pairs_") + .and_then(|op| op.strip_suffix("_validity_bits_v1")) + .ok_or_else(unsupported)?; + let (recipe, value) = [ + ("ascending_full_", Self::AscendingFull), + ("descending_full_", Self::DescendingFull), + ("ascending_partial_", Self::AscendingPartial), + ("descending_partial_", Self::DescendingPartial), + ] + .into_iter() + .find_map(|(prefix, recipe)| operation.strip_prefix(prefix).map(|v| (recipe, v))) + .ok_or_else(unsupported)?; + let value = match value { + "i8" => Dtype::I8, + "u8" => Dtype::U8, + "i16" => Dtype::I16, + "u16" => Dtype::U16, + "i32" => Dtype::I32, + "u32" => Dtype::U32, + "i64" => Dtype::I64, + "u64" => Dtype::U64, + "f32" => Dtype::F32, + "f64" => Dtype::F64, + _ => return Err(unsupported()), + }; + Ok((recipe, value)) + } +} + +/// Validated Block/Warp merge sort with independently typed associated values. +/// +/// Keys and values independently use any [`Dtype`], with the same positive +/// per-lane extent. Full native calls use `(const K* keys_in, K* keys_out, +/// u8* scratch, const V* values_in, V* values_out, const u8* keys_in_bits, +/// u8* keys_out_bits, const u8* values_in_bits, u8* values_out_bits)`. +/// Partial calls append `(i64 valid_items, K sentinel)`. The count and sentinel +/// are uniform within the group; `0 <= valid_items <= group_threads*items` is +/// checked before narrowing or numeric reads. The key-typed sentinel must +/// follow all valid keys in the chosen order. Values are carried unchanged, +/// not compared. Equal keys have no stable-tie guarantee; floating keys must +/// satisfy strict weak ordering, with no NaN total-order guarantee. +/// +/// Each lane supplies fully initialized keys AND values, even for a zero +/// count, with separate packed low-bit-first `ceil(items/8)` validity maps. +/// Every input bit is checked before reading either array. Both inputs and +/// maps are preserved; sorting uses private copies. Outputs, output maps and +/// scratch are mutually disjoint and disjoint from both input arrays/maps. +/// Read-only inputs may alias. Pointers remain live and naturally aligned. +/// Full calls initialize both complete outputs; partial calls write only the +/// same blocked prefix in both outputs. All other bits, including padding, +/// are cleared; invalid output slots are not written. +/// +/// Group geometry, rank, participation and completion follow +/// [`MergeSortKeysRequest`]. Scratch is the complete queried *pair* CUB +/// TempStorage including padding and value exchange, not the keys-only +/// layout. Each Warp group has an independent replica. No trailing barrier +/// is inserted; the caller completes accesses before scratch reuse. +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub struct MergeSortPairsRequest { + pub(crate) recipe: MergeSortPairsRecipe, + pub(crate) key_dtype: Dtype, + pub(crate) value_dtype: Dtype, + keys: MergeSortKeysRequest, +} + +impl MergeSortPairsRequest { + /// Validate independently typed pairs for a complete XYZ Block. + /// + /// `recipe` selects key ordering and full/partial output. `items` is the + /// shared per-lane extent of both inputs and outputs. + /// + /// # Errors + /// Rejects the same geometry, zero extent and padded signed-int capacity + /// violations as [`MergeSortKeysRequest::block`]. + pub fn block( + key_dtype: Dtype, + value_dtype: Dtype, + recipe: MergeSortPairsRecipe, + block: [u32; 3], + items: u32, + ) -> Result { + let keys = MergeSortKeysRequest::block(key_dtype, recipe.keys(), block, items)?; + Ok(Self { + recipe, + key_dtype, + value_dtype, + keys, + }) + } + /// Validate independently typed pairs for a physical/logical Warp group. + /// + /// Supports widths 1/2/4/8/16/32 within complete physical warps. `items` + /// is the shared per-lane extent; scratch is queried for one group. + /// + /// # Errors + /// Rejects the same geometry, width and padded capacity violations as + /// [`MergeSortKeysRequest::warp`]. + pub fn warp( + key_dtype: Dtype, + value_dtype: Dtype, + recipe: MergeSortPairsRecipe, + block: [u32; 3], + width: u32, + items: u32, + ) -> Result { + let keys = MergeSortKeysRequest::warp(key_dtype, recipe.keys(), block, width, items)?; + Ok(Self { + recipe, + key_dtype, + value_dtype, + keys, + }) + } + pub(crate) fn block_shape(self) -> [u32; 3] { + self.keys.block + } + pub(crate) fn source(self, symbol: &str) -> String { + let (key, _) = self.key_dtype.native(); + let (value, _) = self.value_dtype.native(); + let [x, y, z] = self.keys.block; + let items = self.keys.items; + let tile = self.keys.width.unwrap_or(x * y * z) * items; + let (header, primitive, rank) = if let Some(width) = self.keys.width { + ( + "warp/warp_merge_sort.cuh", + format!("cub::WarpMergeSort"), + format!("(threadIdx.x + {x} * (threadIdx.y + {y} * threadIdx.z)) % {width}"), + ) + } else { + ( + "block/block_merge_sort.cuh", + format!("cub::BlockMergeSort"), + format!("threadIdx.x + {x} * (threadIdx.y + {y} * threadIdx.z)"), + ) + }; + let bytes = items / 8 + u32::from(items % 8 != 0); + let comparison = match self.recipe { + MergeSortPairsRecipe::AscendingFull | MergeSortPairsRecipe::AscendingPartial => "<", + MergeSortPairsRecipe::DescendingFull | MergeSortPairsRecipe::DescendingPartial => ">", + }; + let (parameters, bounds, arguments, count) = if self.recipe.partial() { + ( + format!(", long long valid_items, {key} sentinel"), + format!("if (valid_items < 0 || valid_items > {tile}) __trap();"), + ", static_cast(valid_items), sentinel", + "valid_items".into(), + ) + } else { + (String::new(), String::new(), "", tile.to_string()) + }; + format!( + r#" +#include +namespace {symbol}_detail {{ +using key_type = {key}; +using value_type = {value}; +using primitive = {primitive}; +using scratch_storage = primitive::TempStorage; +struct compare {{ + __device__ __forceinline__ bool operator()(key_type lhs, key_type rhs) const {{ + return lhs {comparison} rhs; + }} +}}; +}} +extern "C" __device__ void {symbol}(const {key}* keys_in, {key}* keys_out, + unsigned char* scratch, const {value}* values_in, {value}* values_out, + const unsigned char* keys_in_bits, unsigned char* keys_out_bits, + const unsigned char* values_in_bits, unsigned char* values_out_bits{parameters}) {{ + using namespace {symbol}_detail; + {bounds} +#pragma unroll + for (int i = 0; i < {bytes}; ++i) {{ keys_out_bits[i] = 0; values_out_bits[i] = 0; }} +#pragma unroll + for (int i = 0; i < {items}; ++i) {{ + const unsigned int bit = 1u << (i % 8); + if ((keys_in_bits[i / 8] & bit) == 0 || (values_in_bits[i / 8] & bit) == 0) __trap(); + }} + key_type local_keys[{items}]; + value_type local_values[{items}]; +#pragma unroll + for (int i = 0; i < {items}; ++i) {{ local_keys[i] = keys_in[i]; local_values[i] = values_in[i]; }} + primitive operation{{*reinterpret_cast(scratch)}}; + operation.Sort(local_keys, local_values, compare{{}}{arguments}); + const int rank = static_cast({rank}); +#pragma unroll + for (int i = 0; i < {items}; ++i) {{ + if (rank * {items} + i < {count}) {{ + keys_out[i] = local_keys[i]; values_out[i] = local_values[i]; + keys_out_bits[i / 8] |= static_cast(1u << (i % 8)); + values_out_bits[i / 8] |= static_cast(1u << (i % 8)); + }} + }} +}} +"# + ) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn complete_block_and_padded_int_capacity() { + for dtype in [ + Dtype::I8, + Dtype::U8, + Dtype::I16, + Dtype::U16, + Dtype::I32, + Dtype::U32, + Dtype::I64, + Dtype::U64, + Dtype::F32, + Dtype::F64, + ] { + for (block, items) in [ + ([1, 1, 1], i32::MAX as u32 - 1), + ([4, 4, 4], 4), + ([2, 2, 2], 9), + ] { + assert!( + MergeSortKeysRequest::block( + dtype, + MergeSortKeysRecipe::AscendingFull, + block, + items + ) + .is_ok() + ); + } + } + for (block, items) in [ + ([0, 1, 1], 1), + ([1, 0, 1], 1), + ([1, 1, 0], 1), + ([3, 1, 1], 1), + ([3, 5, 2], 1), + ([2048, 1, 1], 1), + ([u32::MAX, 2, 2], 1), + ([64, 1, 1], 0), + ([1, 1, 1], i32::MAX as u32), + ([1024, 1, 1], 1 << 21), + ([64, 1, 1], u32::MAX), + ] { + assert!( + MergeSortKeysRequest::block( + Dtype::I64, + MergeSortKeysRecipe::DescendingPartial, + block, + items + ) + .is_err() + ); + } + } + + #[test] + fn warp_groups_use_group_capacity_with_complete_physical_warps() { + for width in [1, 2, 4, 8, 16, 32] { + for block in [[64, 1, 1], [4, 4, 6]] { + assert!( + MergeSortKeysRequest::warp( + Dtype::I64, + MergeSortKeysRecipe::DescendingPartial, + block, + width, + 9 + ) + .is_ok() + ); + } + } + // Capacity is per logical group, not the enclosing Block. + assert!( + MergeSortKeysRequest::warp( + Dtype::U8, + MergeSortKeysRecipe::AscendingFull, + [1024, 1, 1], + 1, + i32::MAX as u32 - 1 + ) + .is_ok() + ); + for (block, width, items) in [ + ([0, 1, 1], 8, 1), + ([u32::MAX, 2, 2], 8, 1), + ([2048, 1, 1], 8, 1), + ([33, 1, 1], 8, 1), + ([4, 4, 1], 8, 1), + ([64, 1, 1], 0, 1), + ([64, 1, 1], 3, 1), + ([64, 1, 1], 17, 1), + ([64, 1, 1], 64, 1), + ([64, 1, 1], 8, 0), + ([64, 1, 1], 1, i32::MAX as u32), + ([64, 1, 1], 32, (i32::MAX as u32 / 32) + 1), + ] { + assert!( + MergeSortKeysRequest::warp( + Dtype::I64, + MergeSortKeysRecipe::AscendingFull, + block, + width, + items + ) + .is_err() + ); + } + } +} diff --git a/rust/cuda-coop-codegen/src/oxide.rs b/rust/cuda-coop-codegen/src/oxide.rs index 54dadb026875..cdae849cc8b0 100644 --- a/rust/cuda-coop-codegen/src/oxide.rs +++ b/rust/cuda-coop-codegen/src/oxide.rs @@ -6,7 +6,8 @@ //! 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-request-v2` serves compiler-admitted callable Reduce and +//! Merge Sort Keys comparator 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 @@ -27,16 +28,20 @@ //! carries an independent numeric value type into a fresh paired output. //! [`RadixRankKeysRecipe`] selects stable i32 ranks with a static digit width //! and signed64 begin/end-constraint controls, preserving input-slot order. +//! [`MergeSortKeysRecipe`] selects numeric Block/Warp sorting with full output or a +//! signed64 count and key-typed sentinel for prefix-valid output. //! 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 +//! and root-result ABI. Callable Merge Sort Keys uses full-input Block i32 and a +//! stateless comparator. Their thunks and borrowed environments come from Oxide's //! checked callback marker; the provider does not select or execute Rust code. use crate::{ AdjacentDifferenceRecipe, AdjacentDifferenceRequest, Algorithm, Artifact, DiscontinuityRecipe, DiscontinuityRequest, Dtype, Error, ExchangeRecipe, ExchangeRequest, HistogramAlgorithm, - HistogramRecipe, HistogramRequest, Layout, NativeRequest, Operation, Provider, + HistogramRecipe, HistogramRequest, Layout, MergeSortKeysRecipe, MergeSortKeysRequest, + MergeSortPairsRecipe, MergeSortPairsRequest, NativeRequest, Operation, Provider, RadixRankKeysRecipe, RadixRankKeysRequest, RadixSortKeysRecipe, RadixSortKeysRequest, RadixSortPairsRecipe, RadixSortPairsRequest, ReduceAlgorithm, ReduceBatchedRecipe, ReduceBatchedRequest, ReduceForm, ReduceOperator, ReduceRecipe, ReduceRequest, Request, Result, @@ -153,13 +158,9 @@ struct CallbackRequest { 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() - { + if self.dtype != Scalar::Signed(32) || self.warp_width.is_some() { return Err(Error::message( - "callable Reduce requires the i32 Block payload raking recipe", + "callable recipes require i32 Block payloads", )); } if self.callback.symbol.is_empty() @@ -178,18 +179,32 @@ impl CallbackRequest { || 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", + "callable recipes require 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) + match (self.operation.as_str(), self.algorithm.as_str()) { + ("reduce_callable_payload_validity_bits_v1", "raking") => ReduceRequest::block( + Dtype::I32, + ReduceOperator::Sum, + ReduceAlgorithm::Raking, + ReduceForm::Payload, + self.block, + self.items_per_thread, + ) + .map(NativeRequest::Reduce), + ("merge_sort_keys_callable_payload_validity_bits_v1", "block") => { + MergeSortKeysRequest::block( + Dtype::I32, + MergeSortKeysRecipe::AscendingFull, + self.block, + self.items_per_thread, + ) + .map(NativeRequest::MergeSortKeys) + } + _ => Err(Error::message( + "unsupported callable operation or algorithm", + )), + } } } @@ -207,6 +222,51 @@ struct SpecializationRequest { impl SpecializationRequest { fn native(&self) -> Result { + if self.operation.starts_with("merge_sort_pairs_") { + let dtype = self.dtype.dtype()?; + let (recipe, value) = MergeSortPairsRecipe::parse(&self.operation, &self.algorithm)?; + return match (self.algorithm.as_str(), self.warp_width) { + ("block", None) => MergeSortPairsRequest::block( + dtype, + value, + recipe, + self.block, + self.items_per_thread, + ), + ("warp", Some(width)) => MergeSortPairsRequest::warp( + dtype, + value, + recipe, + self.block, + width, + self.items_per_thread, + ), + _ => Err(Error::message( + "Merge Sort Pairs scope and width must agree", + )), + } + .map(NativeRequest::MergeSortPairs); + } + if self.operation.starts_with("merge_sort_keys_") { + let recipe = MergeSortKeysRecipe::parse(&self.operation, &self.algorithm)?; + let dtype = self.dtype.dtype()?; + return match (self.algorithm.as_str(), self.warp_width) { + ("block", None) => { + MergeSortKeysRequest::block(dtype, recipe, self.block, self.items_per_thread) + } + ("warp", Some(width)) => MergeSortKeysRequest::warp( + dtype, + recipe, + self.block, + width, + self.items_per_thread, + ), + _ => Err(Error::message( + "Merge Sort Keys scope and warp width disagree", + )), + } + .map(NativeRequest::MergeSortKeys); + } if self.operation.starts_with("radix_rank_keys_") { let recipe = RadixRankKeysRecipe::parse(&self.operation, &self.algorithm)?; if self.warp_width.is_some() { @@ -453,6 +513,37 @@ struct BatchResponse { fn arguments(request: NativeRequest) -> Vec { let request = match request { + NativeRequest::MergeSortPairs(request) => { + let key = Scalar::from(request.key_dtype); + let value = Scalar::from(request.value_dtype); + let mut arguments = vec![ + Argument::ConstPointer(key), + Argument::MutPointer(key), + Argument::ConstPointer(value), + Argument::MutPointer(value), + Argument::ConstPointer(Scalar::Unsigned(8)), + Argument::MutPointer(Scalar::Unsigned(8)), + Argument::ConstPointer(Scalar::Unsigned(8)), + Argument::MutPointer(Scalar::Unsigned(8)), + ]; + if request.recipe.partial() { + arguments.extend([Argument::Value(Scalar::Signed(64)), Argument::Value(key)]); + } + return arguments; + } + NativeRequest::MergeSortKeys(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.partial() { + arguments.extend([Argument::Value(Scalar::Signed(64)), Argument::Value(dtype)]); + } + return arguments; + } NativeRequest::RadixRankKeys(request) => { return vec![ Argument::ConstPointer(Scalar::from(request.dtype)), @@ -754,6 +845,27 @@ pub fn reduce_callable_registration(executable: &Path) -> Result { .map_err(|error| Error::message(format!("cannot encode provider registration: {error}"))) } +/// Register full-input i32 Block Merge Sort Keys with a stateless Rust comparator. +/// +/// The compiler supplies a genuine i32-return thunk under protocol version two; +/// the native comparator converts its result to bool. Output retains all keys. +/// +/// # Errors +/// Rejects a nonabsolute executable or a path that cannot be encoded as JSON. +pub fn merge_sort_keys_callable_registration(executable: &Path) -> 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: "merge_sort_keys_callable_payload_validity_bits_v1", + algorithm: "block", + }) + .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 @@ -849,6 +961,118 @@ pub fn shuffle_registration(executable: &Path, recipe: ShuffleRecipe) -> Result< .map_err(|error| Error::message(format!("cannot encode provider registration: {error}"))) } +/// Encode one of four Block Merge Sort Keys order/prefix registrations. +/// +/// Key dtype, exact XYZ and per-lane extent arrive with specialization. Four +/// typed pointers describe preserved input and fresh output with independent +/// packed validity. Full recipes have no controls; partial recipes append a +/// signed64 count and sentinel of exactly the key type. Scratch is inserted at +/// slot 2, giving five or seven native arguments. See [`MergeSortKeysRequest`] +/// for input validity, sentinel ordering, scratch and participation obligations. +/// +/// # Errors +/// Rejects a nonabsolute executable or a path that cannot be encoded as JSON. +pub fn merge_sort_keys_registration( + executable: &Path, + recipe: MergeSortKeysRecipe, +) -> 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: "block", + }) + .map_err(|error| Error::message(format!("cannot encode provider registration: {error}"))) +} + +/// Encode one of four physical/logical Warp Merge Sort Keys registrations. +/// +/// Uses the same typed five/seven-argument native ABI as +/// [`merge_sort_keys_registration`]. The genuine checked Warp marker supplies +/// the logical width and exact enclosing XYZ Block at specialization. Only +/// widths 1/2/4/8/16/32 and complete physical warps are supported. Controls agree +/// within each group; the compiler supplies its independent scratch replica. +/// See [`MergeSortKeysRequest`] for input, validity and sentinel obligations. +/// +/// # Errors +/// Rejects a nonabsolute executable or a path that cannot be encoded as JSON. +pub fn warp_merge_sort_keys_registration( + executable: &Path, + recipe: MergeSortKeysRecipe, +) -> 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: "warp", + }) + .map_err(|error| Error::message(format!("cannot encode provider registration: {error}"))) +} + +/// Encode one of forty independently typed Block Merge Sort Pairs registrations. +/// +/// Key dtype, exact XYZ and shared extent arrive with specialization. Value +/// dtype is selected here; eight payload pointers preserve both input types +/// and their separate validity maps. Scratch at slot 2 gives nine native +/// arguments, or eleven when a partial recipe appends `(i64 count, K sentinel)`. +/// See [`MergeSortPairsRequest`] for participation, validity and scratch rules. +/// +/// # Errors +/// Rejects a nonabsolute executable or a path that cannot be encoded as JSON. +pub fn merge_sort_pairs_registration( + executable: &Path, + recipe: MergeSortPairsRecipe, + value_dtype: Dtype, +) -> 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(value_dtype), + algorithm: "block", + }) + .map_err(|error| Error::message(format!("cannot encode provider registration: {error}"))) +} + +/// Encode one of forty independently typed Warp Merge Sort Pairs registrations. +/// +/// Key dtype, exact XYZ and shared extent arrive with specialization. Value +/// dtype is selected here; eight payload pointers preserve both input types +/// and their separate validity maps. Scratch at slot 2 gives nine native +/// arguments, or eleven when a partial recipe appends `(i64 count, K sentinel)`. +/// See [`MergeSortPairsRequest`] for participation, validity and scratch rules. +/// +/// # Errors +/// Rejects a nonabsolute executable or a path that cannot be encoded as JSON. +pub fn warp_merge_sort_pairs_registration( + executable: &Path, + recipe: MergeSortPairsRecipe, + value_dtype: Dtype, +) -> 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(value_dtype), + algorithm: "warp", + }) + .map_err(|error| Error::message(format!("cannot encode provider registration: {error}"))) +} + /// Encode one of sixteen Block Radix Rank Keys order/width registrations. /// /// `recipe` selects ascending/descending order and 1..=8 radix bits. Key type, @@ -1341,12 +1565,295 @@ mod tests { Ok(()) } + fn comparator_request(symbol: &str, items: u32) -> serde_json::Value { + let mut request = callable_request(symbol, items); + request["operation"] = + serde_json::json!("merge_sort_keys_callable_payload_validity_bits_v1"); + request["algorithm"] = serde_json::json!("block"); + request + } + + #[test] + fn callable_merge_sort_is_closed_before_toolkit() { + let mut base = comparator_request("rust_compare", 4); + base["warp_width"] = serde_json::Value::Null; + let parsed: CallbackRequest = serde_json::from_value(base.clone()).unwrap(); + assert!(matches!( + parsed.native().unwrap(), + NativeRequest::MergeSortKeys(_) + )); + assert!(serde_json::from_value::(base.clone()).is_err()); + for (path, value) in [ + ( + "/operation", + serde_json::json!("merge_sort_keys_ascending_partial_validity_bits_v1"), + ), + ("/algorithm", serde_json::json!("warp")), + ("/dtype/kind", serde_json::json!("unsigned")), + ("/dtype/bits", serde_json::json!(64)), + ("/warp_width", serde_json::json!(8)), + ("/block/0", serde_json::json!(33)), + ("/items_per_thread", serde_json::json!(0)), + ] { + let mut request = base.clone(); + *request.pointer_mut(path).unwrap() = value; + let parsed: CallbackRequest = serde_json::from_value(request).unwrap(); + assert!(parsed.native().is_err(), "{path}"); + } + } + + #[test] + #[ignore = "requires CUDA_HOME and a Linux CUDA toolkit with NVRTC LTO support"] + fn native_callable_merge_sort_identity_and_layout() + -> std::result::Result<(), Box> { + let requests = [ + comparator_request("rust_even_first", 1), + comparator_request("rust_even_first", 4), + comparator_request("rust_descending", 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 builtin = provider.compile_native(&native[..2], &target)?; + let identities = artifact + .specializations + .iter() + .map(|s| s.identity) + .collect::>(); + assert_eq!(identities.len(), 3); + for (index, entry) in artifact.specializations.iter().enumerate() { + let original = &builtin.specializations[index.min(1)]; + assert_ne!(entry.identity, original.identity); + assert_eq!(entry.cub_layout, original.cub_layout); + assert_eq!(entry.scratch, original.scratch); + assert!(!entry.scratch_free && !entry.mutates_payload && entry.scratch.size > 0); + } + 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])?; + 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::{ ReduceAlgorithm, ReduceBatchedLayout, ReduceOperator, ScanAlgorithm, ScanOperation, ScanOperator, }; + #[test] + fn merge_sort_pairs_registration_preserves_independent_payload_and_sentinel() { + for value in [ + Dtype::I8, + Dtype::U8, + Dtype::I16, + Dtype::U16, + Dtype::I32, + Dtype::U32, + Dtype::I64, + Dtype::U64, + Dtype::F32, + Dtype::F64, + ] { + for recipe in [ + MergeSortPairsRecipe::AscendingFull, + MergeSortPairsRecipe::DescendingFull, + MergeSortPairsRecipe::AscendingPartial, + MergeSortPairsRecipe::DescendingPartial, + ] { + for warp in [false, true] { + let encoded = if warp { + warp_merge_sort_pairs_registration(Path::new("/provider"), recipe, value) + } else { + merge_sort_pairs_registration(Path::new("/provider"), recipe, value) + } + .unwrap(); + let registration: serde_json::Value = serde_json::from_str(&encoded).unwrap(); + let mut request = SpecializationRequest { + operation: registration["operation"].as_str().unwrap().into(), + algorithm: registration["algorithm"].as_str().unwrap().into(), + dtype: Scalar::Float(64), + block: [4, 4, 4], + items_per_thread: 9, + warp_width: warp.then_some(8), + }; + let native = request.native().unwrap(); + let expected = if warp { + MergeSortPairsRequest::warp(Dtype::F64, value, recipe, [4, 4, 4], 8, 9) + } else { + MergeSortPairsRequest::block(Dtype::F64, value, recipe, [4, 4, 4], 9) + } + .unwrap(); + assert_eq!(native, NativeRequest::MergeSortPairs(expected)); + let mut abi = vec![ + Argument::ConstPointer(Scalar::Float(64)), + Argument::MutPointer(Scalar::Float(64)), + Argument::ConstPointer(Scalar::from(value)), + Argument::MutPointer(Scalar::from(value)), + Argument::ConstPointer(Scalar::Unsigned(8)), + Argument::MutPointer(Scalar::Unsigned(8)), + Argument::ConstPointer(Scalar::Unsigned(8)), + Argument::MutPointer(Scalar::Unsigned(8)), + ]; + if recipe.partial() { + abi.extend([ + Argument::Value(Scalar::Signed(64)), + Argument::Value(Scalar::Float(64)), + ]); + } + assert_eq!(arguments(native), abi); + assert!(!native.mutates_payload() && !native.scratch_free()); + request.warp_width = if warp { None } else { Some(8) }; + assert!(request.native().is_err()); + } + } + } + for operation in [ + "merge_sort_pairs_ascending_full_f16_validity_bits_v1", + "merge_sort_pairs_ascending_full_f64_validity_bits_v2", + "merge_sort_pairs_ascending_partial_validity_bits_v1", + ] { + assert!(MergeSortPairsRecipe::parse(operation, "block").is_err()); + } + assert!( + MergeSortPairsRecipe::parse( + "merge_sort_pairs_ascending_full_f64_validity_bits_v1", + "custom" + ) + .is_err() + ); + assert!( + merge_sort_pairs_registration( + Path::new("relative"), + MergeSortPairsRecipe::AscendingFull, + Dtype::F64 + ) + .is_err() + ); + assert!( + warp_merge_sort_pairs_registration( + Path::new("relative"), + MergeSortPairsRecipe::AscendingFull, + Dtype::F64 + ) + .is_err() + ); + } + + #[test] + fn merge_sort_keys_registration_preserves_count_and_sentinel_types() { + for (recipe, dtype) in [ + (MergeSortKeysRecipe::AscendingFull, Scalar::Signed(8)), + (MergeSortKeysRecipe::DescendingFull, Scalar::Float(32)), + (MergeSortKeysRecipe::AscendingPartial, Scalar::Unsigned(64)), + (MergeSortKeysRecipe::DescendingPartial, Scalar::Float(64)), + ] { + let registration: serde_json::Value = serde_json::from_str( + &merge_sort_keys_registration(Path::new("/provider"), recipe).unwrap(), + ) + .unwrap(); + let mut request = SpecializationRequest { + operation: registration["operation"].as_str().unwrap().into(), + algorithm: registration["algorithm"].as_str().unwrap().into(), + dtype, + block: [4, 4, 4], + items_per_thread: 9, + warp_width: None, + }; + let native = request.native().unwrap(); + assert_eq!( + native, + NativeRequest::MergeSortKeys( + MergeSortKeysRequest::block(dtype.dtype().unwrap(), recipe, [4, 4, 4], 9) + .unwrap() + ) + ); + let mut expected = vec![ + Argument::ConstPointer(dtype), + Argument::MutPointer(dtype), + Argument::ConstPointer(Scalar::Unsigned(8)), + Argument::MutPointer(Scalar::Unsigned(8)), + ]; + if recipe.partial() { + expected.extend([Argument::Value(Scalar::Signed(64)), Argument::Value(dtype)]); + } + assert_eq!(arguments(native), expected); + assert!(!native.mutates_payload() && !native.scratch_free()); + request.warp_width = Some(32); + assert!(request.native().is_err()); + request.warp_width = None; + request.block = [3, 5, 2]; + assert!(request.native().is_err()); + + let warp: serde_json::Value = serde_json::from_str( + &warp_merge_sort_keys_registration(Path::new("/provider"), recipe).unwrap(), + ) + .unwrap(); + request.operation = warp["operation"].as_str().unwrap().into(); + request.algorithm = warp["algorithm"].as_str().unwrap().into(); + request.block = [4, 4, 6]; + assert!(request.native().is_err()); // No authenticated width. + request.warp_width = Some(8); + let native = request.native().unwrap(); + assert_eq!( + native, + NativeRequest::MergeSortKeys( + MergeSortKeysRequest::warp(dtype.dtype().unwrap(), recipe, [4, 4, 6], 8, 9) + .unwrap() + ) + ); + assert_eq!(arguments(native), expected); + assert!(!native.mutates_payload() && !native.scratch_free()); + request.warp_width = Some(17); + assert!(request.native().is_err()); + } + assert!( + merge_sort_keys_registration(Path::new("relative"), MergeSortKeysRecipe::AscendingFull) + .is_err() + ); + assert!( + warp_merge_sort_keys_registration( + Path::new("relative"), + MergeSortKeysRecipe::AscendingFull + ) + .is_err() + ); + for (operation, algorithm) in [ + ("merge_sort_keys_ascending_full_validity_bits_v1", "other"), + ("merge_sort_keys_ascending_full_validity_bits_v2", "block"), + ("merge_sort_keys_custom_full_validity_bits_v1", "block"), + ("merge_sort_pairs_ascending_full_validity_bits_v1", "block"), + ] { + assert!(MergeSortKeysRecipe::parse(operation, algorithm).is_err()); + } + } + #[test] fn radix_sort_pairs_registration_preserves_independent_value_abi() { for value in [ diff --git a/rust/cuda-coop-codegen/tests/native.rs b/rust/cuda-coop-codegen/tests/native.rs index e10f80c5bbe1..68fefa9fcd62 100644 --- a/rust/cuda-coop-codegen/tests/native.rs +++ b/rust/cuda-coop-codegen/tests/native.rs @@ -6,7 +6,8 @@ use cuda_coop_codegen::{ AdjacentDifferenceRecipe, AdjacentDifferenceRequest, Algorithm, DiscontinuityRecipe, DiscontinuityRequest, Dtype, ExchangeMode, ExchangeRecipe, ExchangeRequest, HistogramAlgorithm, - HistogramRecipe, HistogramRequest, Layout, NativeRequest, Operation, Provider, + HistogramRecipe, HistogramRequest, Layout, MergeSortKeysRecipe, MergeSortKeysRequest, + MergeSortPairsRecipe, MergeSortPairsRequest, NativeRequest, Operation, Provider, RadixRankKeysRecipe, RadixRankKeysRequest, RadixSortKeysRecipe, RadixSortKeysRequest, RadixSortPairsRecipe, RadixSortPairsRequest, ReduceAlgorithm, ReduceBatchedLayout, ReduceBatchedRecipe, ReduceBatchedRequest, ReduceForm, ReduceOperator, ReduceRequest, Request, @@ -14,6 +15,278 @@ use cuda_coop_codegen::{ }; use std::{env, path::Path, time::Instant}; +#[test] +#[ignore = "requires CUDA_HOME and a Linux CUDA toolkit with NVRTC LTO support"] +fn native_mixed_merge_sort_pairs_recipes() -> Result<(), Box> { + use MergeSortPairsRecipe::{ + AscendingFull, AscendingPartial, DescendingFull, DescendingPartial, + }; + 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)?; + // First N1/N4 example, then independent numeric types and Warp widths. + // None means complete Block; Some(width) means one Warp group's layout. + let rows = [ + (Dtype::I32, Dtype::F64, [4, 4, 4], None, 1, AscendingFull), + ( + Dtype::I32, + Dtype::F64, + [4, 4, 4], + None, + 1, + DescendingPartial, + ), + (Dtype::I32, Dtype::F64, [4, 4, 4], None, 4, AscendingFull), + ( + Dtype::I32, + Dtype::F64, + [4, 4, 4], + None, + 4, + DescendingPartial, + ), + ( + Dtype::I8, + Dtype::U64, + [64, 1, 1], + Some(1), + 9, + AscendingPartial, + ), + ( + Dtype::U8, + Dtype::I64, + [64, 1, 1], + Some(2), + 4, + DescendingFull, + ), + ( + Dtype::I16, + Dtype::F32, + [64, 1, 1], + Some(4), + 4, + AscendingFull, + ), + ( + Dtype::U16, + Dtype::I32, + [4, 4, 6], + Some(8), + 9, + DescendingPartial, + ), + ( + Dtype::U32, + Dtype::I8, + [64, 1, 1], + Some(16), + 4, + AscendingPartial, + ), + ( + Dtype::I64, + Dtype::U8, + [64, 1, 1], + Some(32), + 4, + DescendingFull, + ), + (Dtype::U64, Dtype::I16, [1, 1, 1], None, 1, AscendingFull), + ( + Dtype::F32, + Dtype::U16, + [4, 4, 4], + None, + 9, + DescendingPartial, + ), + (Dtype::F64, Dtype::U32, [4, 4, 4], None, 4, AscendingPartial), + ]; + let mut requests = Vec::new(); + for (key, value, block, width, items, recipe) in rows { + let request = if let Some(width) = width { + MergeSortPairsRequest::warp(key, value, recipe, block, width, items)? + } else { + MergeSortPairsRequest::block(key, value, recipe, block, items)? + }; + requests.push(NativeRequest::MergeSortPairs(request)); + } + requests.extend([ + NativeRequest::MergeSortKeys(MergeSortKeysRequest::block( + Dtype::I32, + MergeSortKeysRecipe::AscendingFull, + [4, 4, 4], + 4, + )?), + NativeRequest::MergeSortKeys(MergeSortKeysRequest::warp( + Dtype::F64, + MergeSortKeysRecipe::DescendingPartial, + [64, 1, 1], + 8, + 4, + )?), + ]); + let artifact = provider.compile_native(&requests, &target)?; + assert_eq!(artifact.specializations.len(), 15); + assert!(!artifact.ltoir.is_empty()); + let mut identities = std::collections::HashSet::new(); + for (request, entry) in requests.iter().zip(&artifact.specializations) { + assert!(identities.insert(entry.identity)); + assert!(!entry.mutates_payload && !entry.scratch_free); + assert_eq!(entry.scratch, entry.cub_layout); + assert!(entry.scratch.size > 0 && entry.scratch.alignment.is_power_of_two()); + eprintln!("{request:?}: {entry:?}"); + } + // Full/partial use the complete pair union, including its padded element. + for (full, partial) in [(0, 1), (2, 3)] { + assert_eq!( + artifact.specializations[full].scratch, + artifact.specializations[partial].scratch + ); + assert_ne!( + artifact.specializations[full].identity, + artifact.specializations[partial].identity + ); + } + assert_eq!( + artifact.specializations[0].scratch, + Layout { + size: 65 * 8, + alignment: 8 + } + ); + assert_eq!( + artifact.specializations[2].scratch, + Layout { + size: 257 * 8, + alignment: 8 + } + ); + // The i32-only control cannot stand in for the mixed i32/f64 pair layout. + assert!(artifact.specializations[2].scratch.size > artifact.specializations[13].scratch.size); + Ok(()) +} + +#[test] +#[ignore = "requires CUDA_HOME and a Linux CUDA toolkit with NVRTC LTO support"] +fn native_mixed_merge_sort_keys_recipes() -> Result<(), Box> { + use MergeSortKeysRecipe::{AscendingFull, AscendingPartial, DescendingFull, DescendingPartial}; + 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)?; + // The public N1/N4 example, then distributed numeric types and bitmap tails. + let rows = [ + (Dtype::I64, [4, 4, 4], 1, AscendingFull), + (Dtype::I64, [4, 4, 4], 1, DescendingPartial), + (Dtype::I64, [4, 4, 4], 4, AscendingFull), + (Dtype::I64, [4, 4, 4], 4, DescendingPartial), + (Dtype::I8, [1, 1, 1], 9, AscendingPartial), + (Dtype::U8, [2, 2, 2], 9, DescendingFull), + (Dtype::I16, [64, 1, 1], 4, AscendingFull), + (Dtype::U16, [64, 1, 1], 4, DescendingPartial), + (Dtype::I32, [8, 4, 1], 9, AscendingPartial), + (Dtype::U32, [32, 1, 1], 1, DescendingFull), + (Dtype::U64, [64, 1, 1], 4, AscendingPartial), + (Dtype::F32, [4, 4, 4], 4, AscendingFull), + (Dtype::F64, [4, 4, 4], 4, DescendingPartial), + ]; + let mut requests = Vec::new(); + for (dtype, block, items, recipe) in rows { + requests.push(NativeRequest::MergeSortKeys(MergeSortKeysRequest::block( + dtype, recipe, block, items, + )?)); + } + requests.extend([ + NativeRequest::RadixSortKeys(RadixSortKeysRequest::block( + Dtype::I64, + RadixSortKeysRecipe::Ascending, + [4, 4, 4], + 4, + )?), + NativeRequest::Histogram(HistogramRequest::block( + Dtype::U8, + HistogramRecipe::new(HistogramAlgorithm::Sort, Dtype::U64, 65, 2)?, + [64, 1, 1], + 4, + )?), + ]); + // Append the six actual Warp widths; preserve the original fifteen entries. + let warp_rows = [ + (Dtype::U8, [64, 1, 1], 1, 1, AscendingFull), + (Dtype::I16, [64, 1, 1], 2, 4, DescendingPartial), + (Dtype::U32, [64, 1, 1], 4, 4, AscendingFull), + (Dtype::I32, [4, 4, 6], 8, 9, DescendingPartial), + (Dtype::U64, [64, 1, 1], 16, 4, DescendingPartial), + (Dtype::I64, [64, 1, 1], 32, 4, AscendingPartial), + ]; + for (dtype, block, width, items, recipe) in warp_rows { + requests.push(NativeRequest::MergeSortKeys(MergeSortKeysRequest::warp( + dtype, recipe, block, width, items, + )?)); + } + let artifact = provider.compile_native(&requests, &target)?; + assert_eq!(artifact.specializations.len(), 21); + 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)); + assert!(!entry.mutates_payload && !entry.scratch_free); + if matches!(request, NativeRequest::Histogram(_)) { + assert!(entry.scratch.size >= entry.cub_layout.size); + assert!(entry.scratch.alignment >= entry.cub_layout.alignment); + } else { + assert_eq!(entry.scratch, entry.cub_layout); + } + assert!(entry.scratch.size > 0 && entry.scratch.alignment.is_power_of_two()); + eprintln!("{request:?}: {entry:?}"); + } + // Full/partial ordering uses the same complete CUB storage, including padding. + for (full, partial) in [(0, 1), (2, 3)] { + assert_eq!( + artifact.specializations[full].scratch, + artifact.specializations[partial].scratch + ); + assert_ne!( + artifact.specializations[full].identity, + artifact.specializations[partial].identity + ); + } + for (index, (_, block, items, _)) in rows.iter().enumerate() { + let entry = &artifact.specializations[index]; + let scalar_bytes = match rows[index].0 { + Dtype::I8 | Dtype::U8 => 1, + Dtype::I16 | Dtype::U16 => 2, + Dtype::I32 | Dtype::U32 | Dtype::F32 => 4, + Dtype::I64 | Dtype::U64 | Dtype::F64 => 8, + }; + let padded_tile = (block.iter().product::() * items + 1) as usize; + assert!(entry.scratch.size >= padded_tile * scalar_bytes); + } + for (index, (dtype, _, width, items, _)) in warp_rows.iter().enumerate() { + let scalar_bytes = match dtype { + Dtype::U8 => 1, + Dtype::I16 => 2, + Dtype::U32 | Dtype::I32 => 4, + Dtype::U64 | Dtype::I64 => 8, + _ => unreachable!(), + }; + let padded_tile = (width * items + 1) as usize; + assert!(artifact.specializations[15 + index].scratch.size >= padded_tile * scalar_bytes); + } + eprintln!( + "{} LTO bytes, {} specializations", + artifact.ltoir.len(), + artifact.specializations.len() + ); + Ok(()) +} + #[test] #[ignore = "requires CUDA_HOME and a Linux CUDA toolkit with NVRTC LTO support"] fn native_mixed_radix_rank_keys_recipes() -> Result<(), Box> { diff --git a/rust/cuda-coop/build.rs b/rust/cuda-coop/build.rs index ed72ea69ee37..21e2fc33e68c 100644 --- a/rust/cuda-coop/build.rs +++ b/rust/cuda-coop/build.rs @@ -4,9 +4,10 @@ //! Register the owning host provider for native cooperative operations. use cuda_coop_codegen::{ AdjacentDifferenceRecipe, Algorithm, DiscontinuityRecipe, Dtype, ExchangeMode, ExchangeRecipe, - HistogramAlgorithm, Operation, RadixRankKeysRecipe, RadixSortKeysRecipe, RadixSortPairsRecipe, - ReduceAlgorithm, ReduceBatchedLayout, ReduceBatchedRecipe, ReduceForm, ReduceOperator, - ReduceRecipe, ScanAlgorithm, ScanOperation, ScanOperator, ScanRecipe, ShuffleRecipe, oxide, + HistogramAlgorithm, MergeSortKeysRecipe, MergeSortPairsRecipe, Operation, RadixRankKeysRecipe, + RadixSortKeysRecipe, RadixSortPairsRecipe, ReduceAlgorithm, ReduceBatchedLayout, + ReduceBatchedRecipe, ReduceForm, ReduceOperator, ReduceRecipe, ScanAlgorithm, ScanOperation, + ScanOperator, ScanRecipe, ShuffleRecipe, oxide, }; fn main() { @@ -25,6 +26,51 @@ fn main() { let callable = oxide::reduce_callable_registration(&executable) .expect("cuda.coop callable Reduce registration"); println!("cargo::rustc-env=COOP_REDUCE_RAKING_CALLABLE_PAYLOAD={callable}"); + let comparator = oxide::merge_sort_keys_callable_registration(&executable) + .expect("cuda.coop callable Merge Sort Keys registration"); + println!("cargo::rustc-env=COOP_MERGE_SORT_KEYS_CALLABLE_FULL={comparator}"); + for (recipe, name) in [ + (MergeSortKeysRecipe::AscendingFull, "ASCENDING_FULL"), + (MergeSortKeysRecipe::DescendingFull, "DESCENDING_FULL"), + (MergeSortKeysRecipe::AscendingPartial, "ASCENDING_PARTIAL"), + (MergeSortKeysRecipe::DescendingPartial, "DESCENDING_PARTIAL"), + ] { + let value = oxide::merge_sort_keys_registration(&executable, recipe) + .expect("cuda.coop merge sort keys registration"); + println!("cargo::rustc-env=COOP_MERGE_SORT_KEYS_{name}={value}"); + let value = oxide::warp_merge_sort_keys_registration(&executable, recipe) + .expect("cuda.coop Warp merge sort keys registration"); + println!("cargo::rustc-env=COOP_MERGE_SORT_KEYS_WARP_{name}={value}"); + } + for (recipe, form) in [ + (MergeSortPairsRecipe::AscendingFull, "ASCENDING_FULL"), + (MergeSortPairsRecipe::DescendingFull, "DESCENDING_FULL"), + (MergeSortPairsRecipe::AscendingPartial, "ASCENDING_PARTIAL"), + ( + MergeSortPairsRecipe::DescendingPartial, + "DESCENDING_PARTIAL", + ), + ] { + for (dtype, name) in [ + (Dtype::I8, "I8"), + (Dtype::U8, "U8"), + (Dtype::I16, "I16"), + (Dtype::U16, "U16"), + (Dtype::I32, "I32"), + (Dtype::U32, "U32"), + (Dtype::I64, "I64"), + (Dtype::U64, "U64"), + (Dtype::F32, "F32"), + (Dtype::F64, "F64"), + ] { + let value = oxide::merge_sort_pairs_registration(&executable, recipe, dtype) + .expect("cuda.coop merge sort pairs registration"); + println!("cargo::rustc-env=COOP_MERGE_SORT_PAIRS_{form}_{name}={value}"); + let value = oxide::warp_merge_sort_pairs_registration(&executable, recipe, dtype) + .expect("cuda.coop Warp merge sort pairs registration"); + println!("cargo::rustc-env=COOP_MERGE_SORT_PAIRS_WARP_{form}_{name}={value}"); + } + } for bits in 1..=8 { for (recipe, order) in [ (RadixRankKeysRecipe::Ascending(bits), "ASCENDING"), diff --git a/rust/cuda-coop/examples/merge_sort_keys/.gitignore b/rust/cuda-coop/examples/merge_sort_keys/.gitignore new file mode 100644 index 000000000000..74bffa8af4a4 --- /dev/null +++ b/rust/cuda-coop/examples/merge_sort_keys/.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/merge_sort_keys/Cargo.lock b/rust/cuda-coop/examples/merge_sort_keys/Cargo.lock new file mode 100644 index 000000000000..cf43197b40c7 --- /dev/null +++ b/rust/cuda-coop/examples/merge_sort_keys/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-merge-sort-keys" +version = "0.0.0" +dependencies = [ + "coop-merge-sort-keys-helper", + "cuda-async", + "cuda-coop", + "cuda-core", + "cuda-device", + "cuda-host", +] + +[[package]] +name = "coop-merge-sort-keys-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/merge_sort_keys/Cargo.toml b/rust/cuda-coop/examples/merge_sort_keys/Cargo.toml new file mode 100644 index 000000000000..776539c57078 --- /dev/null +++ b/rust/cuda-coop/examples/merge_sort_keys/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-merge-sort-keys" +version = "0.0.0" +edition = "2024" +publish = false + +[workspace] +members = ["helper"] + +[dependencies] +cuda-coop = { path = "../.." } +coop-merge-sort-keys-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/merge_sort_keys/helper/Cargo.toml b/rust/cuda-coop/examples/merge_sort_keys/helper/Cargo.toml new file mode 100644 index 000000000000..c374ccb60da7 --- /dev/null +++ b/rust/cuda-coop/examples/merge_sort_keys/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-merge-sort-keys-helper" +version = "0.0.0" +edition = "2024" +publish = false + +[dependencies] +cuda-coop = { path = "../../.." } diff --git a/rust/cuda-coop/examples/merge_sort_keys/helper/src/lib.rs b/rust/cuda-coop/examples/merge_sort_keys/helper/src/lib.rs new file mode 100644 index 000000000000..f45f7f49ba2c --- /dev/null +++ b/rust/cuda-coop/examples/merge_sort_keys/helper/src/lib.rs @@ -0,0 +1,61 @@ +// Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +//! Load inferred numeric keys and return fresh full and partial sort results. +#![no_std] +use cuda_coop as coop; + +/// Full ascending, partial descending, then the preserved input payload. +pub type SortedTile = ( + coop::ThreadData, + coop::ThreadData, + coop::ThreadData, +); + +/// Sort two ways while borrowing the same fully initialized key tile. +/// +/// Count and sentinel are uniform; the sentinel follows valid descending keys. +pub fn sort_tile( + block: coop::ThreadGroup, + keys: &[T], + offset: usize, + valid_items: i64, + sentinel: T, +) -> SortedTile +where + coop::Alignment: coop::SupportedAlignment, +{ + let mut input = coop::ThreadData::<_, ITEMS_PER_THREAD, ALIGN>::new(); + coop::load( + block, + keys, + &mut input, + coop::LoadOptions::default().offset(offset), + ); + let ascending = coop::merge_sort_keys(block, &input, Default::default()); + let descending = coop::merge_sort_keys( + block, + &input, + coop::MergeSortOptions::default() + .valid_items(valid_items, sentinel) + .order(coop::Descending), + ); + (ascending, descending, input) +} + +/// Sort full i32 Keys with either a built-in order or a stateless comparator. +/// +/// The comparator defines a strict weak ordering. Its concrete type and the +/// selected storage are retained through this unchanged generic order bound. +pub fn sort_values( + block: coop::ThreadGroup, + input: &coop::ThreadData, + options: coop::MergeSortOptions, +) -> coop::ThreadData { + coop::merge_sort_keys(block, input, options) +} + +/// Order even keys before odd keys, then order each parity numerically. +pub fn even_first(left: i32, right: i32) -> bool { + (left & 1) < (right & 1) || ((left & 1) == (right & 1) && left < right) +} diff --git a/rust/cuda-coop/examples/merge_sort_keys/src/main.rs b/rust/cuda-coop/examples/merge_sort_keys/src/main.rs new file mode 100644 index 000000000000..9f7cca524faa --- /dev/null +++ b/rust/cuda-coop/examples/merge_sort_keys/src/main.rs @@ -0,0 +1,290 @@ +// 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 = 3, block = (4, 4, 4), dynamic_shared = inferred)] + pub fn sort_tile( + keys: &[i64], + mut ascending: coop::DisjointSlice>, + mut descending: coop::DisjointSlice>, + mut preserved: coop::DisjointSlice>, + mut validity: coop::DisjointSlice>, + mut mutated: coop::DisjointSlice>, + valid_items: i64, + ) { + let block = coop::this_block(); + let offset = + block.rank_at::() as usize * block.count() as usize * ITEMS_PER_THREAD; + let (mut ascending_keys, descending_keys, input) = + coop_merge_sort_keys_helper::sort_tile::<_, ITEMS_PER_THREAD, 32>( + block, + keys, + offset, + valid_items, + i64::MIN, + ); + let mut flags = coop::ThreadData::<_, ITEMS_PER_THREAD>::new(); + for item in 0..ITEMS_PER_THREAD { + flags.write(item, i32::from(descending_keys.get(item).is_some())); + } + coop::store( + block, + &mut ascending, + &ascending_keys, + coop::StoreOptions::default().offset(offset), + ); + coop::store( + block, + &mut descending, + &descending_keys, + coop::StoreOptions::default() + .offset(offset) + .valid_items(valid_items as usize), + ); + // Observe a mutation of the fresh result separately from the original. + ascending_keys.write(0, -9999); + coop::store( + block, + &mut mutated, + &ascending_keys, + coop::StoreOptions::default().offset(offset), + ); + coop::store( + block, + &mut preserved, + &input, + coop::StoreOptions::default().offset(offset), + ); + coop::store( + block, + &mut validity, + &flags, + coop::StoreOptions::default().offset(offset), + ); + } + #[kernel] + #[launch_contract(domain = 3, block = (128, 1, 1), dynamic_shared = inferred)] + pub fn compare_tile( + keys: &[i32], + mut ascending: coop::DisjointSlice>, + mut even_first: coop::DisjointSlice>, + mut descending: coop::DisjointSlice>, + mut preserved: coop::DisjointSlice>, + mut validity: coop::DisjointSlice>, + mut mutated: coop::DisjointSlice>, + ) { + let block = coop::this_block(); + let offset = + block.rank_at::() as usize * block.count() as usize * ITEMS_PER_THREAD; + let mut input = coop::ThreadData::<_, ITEMS_PER_THREAD>::new(); + coop::load( + block, + keys, + &mut input, + coop::LoadOptions::default().offset(offset), + ); + let mut scratch = coop::TempStorage::new().auto_sync(); + let mut ascending_keys = coop_merge_sort_keys_helper::sort_values( + block, + &input, + coop::MergeSortOptions::default().temp_storage(&mut scratch), + ); + let even_keys = coop_merge_sort_keys_helper::sort_values( + block, + &input, + coop::MergeSortOptions::default() + .order(coop_merge_sort_keys_helper::even_first) + .temp_storage(&mut scratch), + ); + let descending_keys = coop_merge_sort_keys_helper::sort_values( + block, + &input, + coop::MergeSortOptions::default().order(|left, right| left > right), + ); + // All three fresh results survive explicit and implicit scratch use. + let mut flags = coop::ThreadData::<_, ITEMS_PER_THREAD>::new(); + for item in 0..ITEMS_PER_THREAD { + flags.write( + item, + i32::from(ascending_keys.get(item).is_some()) + | (i32::from(even_keys.get(item).is_some()) << 1) + | (i32::from(descending_keys.get(item).is_some()) << 2), + ); + } + coop::store( + block, + &mut ascending, + &ascending_keys, + coop::StoreOptions::default().offset(offset), + ); + ascending_keys.write(0, -9999); + coop::store( + block, + &mut mutated, + &ascending_keys, + coop::StoreOptions::default().offset(offset), + ); + coop::store( + block, + &mut even_first, + &even_keys, + coop::StoreOptions::default().offset(offset), + ); + coop::store( + block, + &mut descending, + &descending_keys, + coop::StoreOptions::default().offset(offset), + ); + coop::store( + block, + &mut preserved, + &input, + coop::StoreOptions::default().offset(offset), + ); + coop::store( + block, + &mut validity, + &flags, + coop::StoreOptions::default().offset(offset), + ); + } +} + +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, LaunchConfig3D}; + use std::cmp::Reverse; + + 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 = 64 * $items; + let count = tile - 1; + let keys: Vec = (0..2 * tile).map(|i| { + let local = i % tile; + let block = i / tile; + // Repeated signed keys exercise comparisons. The excluded key + // would lead the descending result if the count were ignored. + if local == count { 1000 + block as i64 } + else { ((local * 17 + block * 11) % 23) as i64 - 11 } + }).collect(); + let mut expected_ascending = Vec::with_capacity(keys.len()); + let mut expected_descending = Vec::with_capacity(keys.len()); + let mut expected_validity = Vec::with_capacity(keys.len()); + for input in keys.chunks_exact(tile) { + let mut ascending = input.to_vec(); + ascending.sort_unstable(); + expected_ascending.extend(ascending); + let mut descending = input[..count].to_vec(); + descending.sort_unstable_by_key(|value| Reverse(*value)); + expected_descending.extend(descending); + expected_descending.push(-777); + expected_validity.extend(std::iter::repeat_n(1_i32, count)); + expected_validity.push(0); + } + let mut expected_mutated = expected_ascending.clone(); + for lane in expected_mutated.chunks_exact_mut($items) { lane[0] = -9999; } + let source = DeviceBuffer::from_host(&stream, &keys)?; + let ascending_target = DeviceBuffer::from_host(&stream, &vec![-777_i64; keys.len()])?; + let descending_target = DeviceBuffer::from_host(&stream, &vec![-777_i64; keys.len()])?; + let preserved_target = DeviceBuffer::from_host(&stream, &vec![17_i64; keys.len()])?; + let validity_target = DeviceBuffer::from_host(&stream, &vec![-1_i32; keys.len()])?; + let mutated_target = DeviceBuffer::from_host(&stream, &vec![17_i64; keys.len()])?; + let plan = module.prepare_sort_tile::<$items>(LaunchConfig3D::auto_shared((2, 1, 1), (4, 4, 4)))?; + let (source, ascending_target, descending_target, preserved_target, validity_target, mutated_target) = + module.sort_tile_async_owned::<$items, _, _, _, _, _, _>( + &plan, source, ascending_target, descending_target, preserved_target, validity_target, mutated_target, count as i64, + ).sync()?; + let mut ascending = vec![0_i64; keys.len()]; + let mut descending = vec![0_i64; keys.len()]; + let mut preserved = vec![0_i64; keys.len()]; + let mut source_after = vec![0_i64; keys.len()]; + let mut validity = vec![0_i32; keys.len()]; + let mut mutated = vec![0_i64; keys.len()]; + ascending_target.copy_to_host(&stream, &mut ascending)?; + descending_target.copy_to_host(&stream, &mut descending)?; + preserved_target.copy_to_host(&stream, &mut preserved)?; + source.copy_to_host(&stream, &mut source_after)?; + validity_target.copy_to_host(&stream, &mut validity)?; + mutated_target.copy_to_host(&stream, &mut mutated)?; + stream.synchronize()?; + assert_eq!(ascending, expected_ascending); + assert_eq!(descending, expected_descending); + assert_eq!(preserved, keys); + assert_eq!(source_after, keys); + assert_eq!(validity, expected_validity); + assert_eq!(mutated, expected_mutated); + println!("items/thread={}: {} sorted values, {} preserved keys, {} source keys, {} validity flags and {} mutated values checked", $items, keys.len() + 2 * count, keys.len(), keys.len(), keys.len(), keys.len()); + }}; + } + run!(1); + run!(4); + macro_rules! run_callable { + ($items:literal) => {{ + let tile = 128 * $items; + let keys: Vec = (0..3 * tile).map(|i| match i % tile { + 0 => i32::MIN, + 1 => i32::MAX, + _ => ((i * 17 + i / tile * 11) % 23) as i32 - 11, + }).collect(); + let mut expected_ascending = Vec::new(); + let mut expected_even = Vec::new(); + let mut expected_descending = Vec::new(); + for input in keys.chunks_exact(tile) { + let mut sorted = input.to_vec(); + sorted.sort_unstable(); + expected_ascending.extend_from_slice(&sorted); + sorted.sort_unstable_by_key(|value| (value & 1, *value)); + expected_even.extend_from_slice(&sorted); + sorted.sort_unstable_by_key(|value| Reverse(*value)); + expected_descending.extend(sorted); + } + let mut expected_mutated = expected_ascending.clone(); + for lane in expected_mutated.chunks_exact_mut($items) { lane[0] = -9999; } + let source = DeviceBuffer::from_host(&stream, &keys)?; + let ascending_target = DeviceBuffer::from_host(&stream, &vec![-777_i32; keys.len()])?; + let even_target = DeviceBuffer::from_host(&stream, &vec![-777_i32; keys.len()])?; + let descending_target = DeviceBuffer::from_host(&stream, &vec![-777_i32; keys.len()])?; + let preserved_target = DeviceBuffer::from_host(&stream, &vec![-777_i32; keys.len()])?; + let validity_target = DeviceBuffer::from_host(&stream, &vec![-1_i32; keys.len()])?; + let mutated_target = DeviceBuffer::from_host(&stream, &vec![-777_i32; keys.len()])?; + let plan = module.prepare_compare_tile::<$items>(LaunchConfig3D::auto_shared((3, 1, 1), (128, 1, 1)))?; + let (source, ascending_target, even_target, descending_target, preserved_target, validity_target, mutated_target) = + module.compare_tile_async_owned::<$items, _, _, _, _, _, _, _>( + &plan, source, ascending_target, even_target, descending_target, preserved_target, validity_target, mutated_target, + ).sync()?; + for (target, expected) in [ + (&ascending_target, &expected_ascending), + (&even_target, &expected_even), + (&descending_target, &expected_descending), + (&preserved_target, &keys), + (&source, &keys), + (&validity_target, &vec![7_i32; keys.len()]), + (&mutated_target, &expected_mutated), + ] { + let mut actual = vec![0_i32; keys.len()]; + target.copy_to_host(&stream, &mut actual)?; + stream.synchronize()?; + assert_eq!(&actual, expected); + } + println!("callable items/thread={}: {} sorted values, {} preserved keys, {} source keys, {} validity flags and {} mutated values checked", $items, 3 * keys.len(), keys.len(), keys.len(), keys.len(), keys.len()); + }}; + } + run_callable!(1); + run_callable!(4); + Ok(()) +} diff --git a/rust/cuda-coop/examples/merge_sort_pairs/.gitignore b/rust/cuda-coop/examples/merge_sort_pairs/.gitignore new file mode 100644 index 000000000000..74bffa8af4a4 --- /dev/null +++ b/rust/cuda-coop/examples/merge_sort_pairs/.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/merge_sort_pairs/Cargo.lock b/rust/cuda-coop/examples/merge_sort_pairs/Cargo.lock new file mode 100644 index 000000000000..705671d8221e --- /dev/null +++ b/rust/cuda-coop/examples/merge_sort_pairs/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-merge-sort-pairs" +version = "0.0.0" +dependencies = [ + "coop-merge-sort-pairs-helper", + "cuda-async", + "cuda-coop", + "cuda-core", + "cuda-device", + "cuda-host", +] + +[[package]] +name = "coop-merge-sort-pairs-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/merge_sort_pairs/Cargo.toml b/rust/cuda-coop/examples/merge_sort_pairs/Cargo.toml new file mode 100644 index 000000000000..da1bb780d7d2 --- /dev/null +++ b/rust/cuda-coop/examples/merge_sort_pairs/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-merge-sort-pairs" +version = "0.0.0" +edition = "2024" +publish = false + +[workspace] +members = ["helper"] + +[dependencies] +cuda-coop = { path = "../.." } +coop-merge-sort-pairs-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/merge_sort_pairs/helper/Cargo.toml b/rust/cuda-coop/examples/merge_sort_pairs/helper/Cargo.toml new file mode 100644 index 000000000000..2efe543be1c2 --- /dev/null +++ b/rust/cuda-coop/examples/merge_sort_pairs/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-merge-sort-pairs-helper" +version = "0.0.0" +edition = "2024" +publish = false + +[dependencies] +cuda-coop = { path = "../../.." } diff --git a/rust/cuda-coop/examples/merge_sort_pairs/helper/src/lib.rs b/rust/cuda-coop/examples/merge_sort_pairs/helper/src/lib.rs new file mode 100644 index 000000000000..efc3968e0d48 --- /dev/null +++ b/rust/cuda-coop/examples/merge_sort_pairs/helper/src/lib.rs @@ -0,0 +1,72 @@ +// Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +//! Infer both payload types from Load, then merge-sort their associated pairs. +#![no_std] +use cuda_coop as coop; + +/// Ascending pairs, descending pairs, then the two preserved input payloads. +pub type SortedTiles = ( + coop::ThreadData, + coop::ThreadData, + coop::ThreadData, + coop::ThreadData, + coop::ThreadData, + coop::ThreadData, +); + +/// Return full ascending pairs, partial descending pairs and both intact inputs. +/// +/// Count and key-typed sentinel agree across the Block. Keys and values have +/// independent minimum alignments; every input slot is initialized. +pub fn sort_tile< + K: coop::Numeric, + V: coop::MergeSortValue, + const ITEMS_PER_THREAD: usize, + const KEY_ALIGN: usize, + const VALUE_ALIGN: usize, +>( + block: coop::ThreadGroup, + keys: &[K], + values: &[V], + offset: usize, + valid_items: i64, + sentinel: K, +) -> SortedTiles +where + coop::Alignment: coop::SupportedAlignment, + coop::Alignment: coop::SupportedAlignment, +{ + let mut input_keys = coop::ThreadData::<_, ITEMS_PER_THREAD, KEY_ALIGN>::new(); + let mut input_values = coop::ThreadData::<_, ITEMS_PER_THREAD, VALUE_ALIGN>::new(); + coop::load( + block, + keys, + &mut input_keys, + coop::LoadOptions::default().offset(offset), + ); + coop::load( + block, + values, + &mut input_values, + coop::LoadOptions::default().offset(offset), + ); + let (ascending_keys, ascending_values) = + coop::merge_sort_pairs(block, &input_keys, &input_values, Default::default()); + let (descending_keys, descending_values) = coop::merge_sort_pairs( + block, + &input_keys, + &input_values, + coop::MergeSortOptions::default() + .valid_items(valid_items, sentinel) + .order(coop::Descending), + ); + ( + ascending_keys, + ascending_values, + descending_keys, + descending_values, + input_keys, + input_values, + ) +} diff --git a/rust/cuda-coop/examples/merge_sort_pairs/src/main.rs b/rust/cuda-coop/examples/merge_sort_pairs/src/main.rs new file mode 100644 index 000000000000..238e55c4e7f7 --- /dev/null +++ b/rust/cuda-coop/examples/merge_sort_pairs/src/main.rs @@ -0,0 +1,248 @@ +// Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +// The kernel and generated launch wrapper expose separate typed result buffers. +#![allow(clippy::too_many_arguments)] + +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 = 3, block = (4, 4, 4), dynamic_shared = inferred)] + pub fn sort_tile( + keys: &[i32], + values: &[f64], + mut ascending_keys: coop::DisjointSlice>, + mut ascending_values: coop::DisjointSlice>, + mut descending_keys: coop::DisjointSlice>, + mut descending_values: coop::DisjointSlice>, + mut preserved_keys: coop::DisjointSlice>, + mut preserved_values: coop::DisjointSlice>, + mut mutated_keys: coop::DisjointSlice>, + mut mutated_values: coop::DisjointSlice>, + mut validity: coop::DisjointSlice>, + valid_items: i64, + ) { + let block = coop::this_block(); + let offset = + block.rank_at::() as usize * block.count() as usize * ITEMS_PER_THREAD; + let (mut ascending_k, mut ascending_v, descending_k, descending_v, input_k, input_v) = + coop_merge_sort_pairs_helper::sort_tile::<_, _, ITEMS_PER_THREAD, 16, 32>( + block, + keys, + values, + offset, + valid_items, + i32::MIN, + ); + let mut flags = coop::ThreadData::<_, ITEMS_PER_THREAD>::new(); + for item in 0..ITEMS_PER_THREAD { + flags.write( + item, + i32::from(input_k.get(item).is_some()) + | (i32::from(input_v.get(item).is_some()) << 1) + | (i32::from(descending_k.get(item).is_some()) << 2) + | (i32::from(descending_v.get(item).is_some()) << 3), + ); + } + coop::store( + block, + &mut validity, + &flags, + coop::StoreOptions::default().offset(offset), + ); + coop::store( + block, + &mut ascending_keys, + &ascending_k, + coop::StoreOptions::default().offset(offset), + ); + coop::store( + block, + &mut ascending_values, + &ascending_v, + coop::StoreOptions::default().offset(offset), + ); + // Modify fresh outputs after observing the complete original sort. + // The next Stores observe both independent results and both inputs. + ascending_k.write(0, 123_456_789); + ascending_v.write(0, -4096.5); + coop::store( + block, + &mut descending_keys, + &descending_k, + coop::StoreOptions::default() + .offset(offset) + .valid_items(valid_items as usize), + ); + coop::store( + block, + &mut descending_values, + &descending_v, + coop::StoreOptions::default() + .offset(offset) + .valid_items(valid_items as usize), + ); + coop::store( + block, + &mut preserved_keys, + &input_k, + coop::StoreOptions::default().offset(offset), + ); + coop::store( + block, + &mut preserved_values, + &input_v, + coop::StoreOptions::default().offset(offset), + ); + coop::store( + block, + &mut mutated_keys, + &ascending_k, + coop::StoreOptions::default().offset(offset), + ); + coop::store( + block, + &mut mutated_values, + &ascending_v, + coop::StoreOptions::default().offset(offset), + ); + } +} + +// CUB Merge Sort does not promise stable ties. Verify numeric key order and +// the exact associated (key, value bits) multiset, not a tie permutation. +fn check_pairs( + input_k: &[i32], + input_v: &[f64], + output_k: &[i32], + output_v: &[f64], + descending: bool, +) { + assert!(output_k.windows(2).all(|w| if descending { + w[0] >= w[1] + } else { + w[0] <= w[1] + })); + let mut expected: Vec<_> = input_k + .iter() + .copied() + .zip(input_v.iter().map(|v| v.to_bits())) + .collect(); + let mut actual: Vec<_> = output_k + .iter() + .copied() + .zip(output_v.iter().map(|v| v.to_bits())) + .collect(); + expected.sort_unstable(); + actual.sort_unstable(); + assert_eq!(actual, expected); +} + +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, 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)? }; + macro_rules! run { + ($items:literal) => {{ + let tile = 64 * $items; + let count = tile - 1; + let keys: Vec = (0..2 * tile).map(|i| { + let local = i % tile; + let block = i / tile; + // Repeated signed keys allow arbitrary tie order. The excluded + // suffix would lead descending output if the count were ignored. + if local == count { 1000 + block as i32 } + else { ((local * 17 + block * 11) % 23) as i32 - 11 } + }).collect(); + let values: Vec = (0..keys.len()).map(|i| { + // Values are carried, never compared: preserve NaN payloads, + // signed zeros, and distinct finite association tags bitwise. + match i { + 0 => 0.0, + 1 => -0.0, + _ if i % 7 == 0 => f64::from_bits(0x7ff8_0000_0000_0000 | i as u64), + _ => i as f64 + 0.25, + } + }).collect(); + let mut expected_validity = Vec::with_capacity(keys.len()); + for _ in 0..2 { + expected_validity.extend(std::iter::repeat_n(15_i32, count)); + expected_validity.push(3); + } + let source_keys = DeviceBuffer::from_host(&stream, &keys)?; + let source_values = DeviceBuffer::from_host(&stream, &values)?; + let ascending_keys_target = DeviceBuffer::from_host(&stream, &vec![17 as i32; keys.len()])?; + let ascending_values_target = DeviceBuffer::from_host(&stream, &vec![17 as f64; keys.len()])?; + let descending_keys_target = DeviceBuffer::from_host(&stream, &vec![17 as i32; keys.len()])?; + let descending_values_target = DeviceBuffer::from_host(&stream, &vec![17 as f64; keys.len()])?; + let preserved_keys_target = DeviceBuffer::from_host(&stream, &vec![17 as i32; keys.len()])?; + let preserved_values_target = DeviceBuffer::from_host(&stream, &vec![17 as f64; keys.len()])?; + let mutated_keys_target = DeviceBuffer::from_host(&stream, &vec![17 as i32; keys.len()])?; + let mutated_values_target = DeviceBuffer::from_host(&stream, &vec![17 as f64; keys.len()])?; + let validity_target = DeviceBuffer::from_host(&stream, &vec![-1_i32; keys.len()])?; + let plan = module.prepare_sort_tile::<$items>(LaunchConfig3D::auto_shared((2, 1, 1), (4, 4, 4)))?; + let (source_keys, source_values, ascending_keys_target, ascending_values_target, descending_keys_target, descending_values_target, preserved_keys_target, preserved_values_target, mutated_keys_target, mutated_values_target, validity_target) = module.sort_tile_async_owned::<$items, _, _, _, _, _, _, _, _, _, _, _>(&plan, source_keys, source_values, ascending_keys_target, ascending_values_target, descending_keys_target, descending_values_target, preserved_keys_target, preserved_values_target, mutated_keys_target, mutated_values_target, validity_target, count as i64).sync()?; + let mut ascending_keys = vec![0 as i32; keys.len()]; + ascending_keys_target.copy_to_host(&stream, &mut ascending_keys)?; + let mut ascending_values = vec![0 as f64; keys.len()]; + ascending_values_target.copy_to_host(&stream, &mut ascending_values)?; + let mut descending_keys = vec![0 as i32; keys.len()]; + descending_keys_target.copy_to_host(&stream, &mut descending_keys)?; + let mut descending_values = vec![0 as f64; keys.len()]; + descending_values_target.copy_to_host(&stream, &mut descending_values)?; + let mut preserved_keys = vec![0 as i32; keys.len()]; + preserved_keys_target.copy_to_host(&stream, &mut preserved_keys)?; + let mut preserved_values = vec![0 as f64; keys.len()]; + preserved_values_target.copy_to_host(&stream, &mut preserved_values)?; + let mut mutated_keys = vec![0 as i32; keys.len()]; + mutated_keys_target.copy_to_host(&stream, &mut mutated_keys)?; + let mut mutated_values = vec![0 as f64; keys.len()]; + mutated_values_target.copy_to_host(&stream, &mut mutated_values)?; + let mut source_keys_after = vec![0_i32; keys.len()]; + let mut source_values_after = vec![0_f64; values.len()]; + source_keys.copy_to_host(&stream, &mut source_keys_after)?; + source_values.copy_to_host(&stream, &mut source_values_after)?; + let mut validity = vec![0_i32; keys.len()]; + validity_target.copy_to_host(&stream, &mut validity)?; + stream.synchronize()?; + for block in 0..2 { + let base = block * tile; + check_pairs(&keys[base..base+tile], &values[base..base+tile], + &ascending_keys[base..base+tile], &ascending_values[base..base+tile], false); + check_pairs(&keys[base..base+count], &values[base..base+count], + &descending_keys[base..base+count], &descending_values[base..base+count], true); + assert_eq!(descending_keys[base+count], 17); + assert_eq!(descending_values[base+count].to_bits(), 17_f64.to_bits()); + } + let mut expected_mutated_keys = ascending_keys.clone(); + let mut expected_mutated_values = ascending_values.clone(); + for index in (0..keys.len()).step_by($items) { + expected_mutated_keys[index] = 123_456_789; + expected_mutated_values[index] = -4096.5; + } + assert_eq!(validity, expected_validity); + assert_eq!(preserved_keys, keys); + assert_eq!(preserved_values.iter().map(|x| x.to_bits()).collect::>(), values.iter().map(|x| x.to_bits()).collect::>()); + assert_eq!(mutated_keys, expected_mutated_keys); + assert_eq!(mutated_values.iter().map(|x| x.to_bits()).collect::>(), expected_mutated_values.iter().map(|x| x.to_bits()).collect::>()); + assert_eq!(source_keys_after, keys); + assert_eq!(source_values_after.iter().map(|x| x.to_bits()).collect::>(), values.iter().map(|x| x.to_bits()).collect::>()); + println!("items/thread={}: {} sorted key/value observations, {} preserved inputs, {} source values, {} validity flags and {} mutated outputs checked", $items, 2 * keys.len() + 4 * count, 2 * keys.len(), 2 * keys.len(), keys.len(), 2 * keys.len()); + }}; + } + run!(1); + run!(4); + Ok(()) +} diff --git a/rust/cuda-coop/src/lib.rs b/rust/cuda-coop/src/lib.rs index a015aa5c07fd..1d2763bc3315 100644 --- a/rust/cuda-coop/src/lib.rs +++ b/rust/cuda-coop/src/lib.rs @@ -75,6 +75,11 @@ //! ordering, returning fresh key/value payloads with independent alignments. //! [`radix_rank_keys`] returns fresh naturally aligned i32 ranks for selected //! key digits. [`RadixRankOptions`] selects static width/order and begin/end. +//! [`merge_sort_keys`] sorts numeric keys on complete power-of-two XYZ Blocks +//! or physical/logical Warps, with existing explicit storage available to Blocks. +//! [`merge_sort_pairs`] carries independently typed associated values into two fresh outputs. +//! [`MergeSortOptions`] selects built-in order and an optional count/sentinel; +//! fresh results preserve the input shape and expose only the sorted prefix. //! Host payload tests need no GPU; //! executing hardware queries requires an Oxide-compiled CUDA kernel. @@ -89,6 +94,7 @@ mod exchange; mod group; mod histogram; mod load_store; +mod merge_sort; mod radix_sort; mod reduce; mod reduce_batched; @@ -116,6 +122,10 @@ pub use load_store::{ Algorithm, Direct, LoadOptions, LoadSource, LoadStoreGroup, StoreOptions, StoreValue, Striped, Transpose, Vectorize, WarpTranspose, WarpTransposeTimesliced, load, store, }; +pub use merge_sort::{ + MergeSortGroup, MergeSortOptions, MergeSortOrder, MergeSortPrefix, MergeSortValue, + merge_sort_keys, merge_sort_pairs, +}; pub use radix_sort::{ Ascending, Descending, RadixBit, RadixKey, RadixRankBits, RadixRankOptions, RadixRankWidth, RadixSortBit, RadixSortOptions, RadixSortOrder, RadixValue, radix_rank_keys, radix_sort_keys, diff --git a/rust/cuda-coop/src/merge_sort.rs b/rust/cuda-coop/src/merge_sort.rs new file mode 100644 index 000000000000..eee01ae602ab --- /dev/null +++ b/rust/cuda-coop/src/merge_sort.rs @@ -0,0 +1,571 @@ +// 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, Ascending, Block, Descending, GroupScope, Numeric, SupportedAlignment, TempStorage, + TempStorageLike, ThreadData, ThreadGroup, ThreadsWithinWarp, Warp, +}; + +use cuda_device::library; + +mod boundary; + +/// A sealed Merge Sort order or stateless i32 key comparator. +/// +/// [`Ascending`] and [`Descending`] compare keys in their original numeric type. +/// Full-input Block i32 Keys also accepts function items and zero-capture +/// `Fn(i32, i32) -> bool + Copy` closures. True means the left key precedes the +/// right; the comparator must define a strict weak ordering. Merge Sort does +/// not guarantee the relative order of equal keys. +pub trait MergeSortOrder: sealed::Order {} + +/// Sealed associated-value types for [`merge_sort_pairs`]. +/// +/// All ten [`Numeric`] types are supported independently of the key type. +/// Values follow their keys unchanged, including floating-point bit patterns; +/// they do not participate in comparisons. +pub trait MergeSortValue: Numeric + sealed::Value {} + +/// A sealed full or partial tile selection for [`MergeSortOptions`]. +/// +/// Omitted `()` selects the full tile. The [`MergeSortOptions::valid_items`] +/// builder selects `(i64, T)`, keeping a uniform count and key-typed sentinel +/// together. This type bound alone does not prove their uniformity. +pub trait MergeSortPrefix: sealed::Prefix {} + +/// A sealed Merge Sort group and storage combination. +/// +/// Complete power-of-two Blocks accept implicit or explicit storage. Physical +/// [`Warp`] and logical widths 1/2/4/8/16/32 accept implicit storage only and +/// require complete physical warps in the enclosing Block. +/// +/// ```compile_fail,E0308 +/// use cuda_coop as coop; +/// fn invalid(keys: &coop::ThreadData, scratch: &mut coop::TempStorage) { +/// coop::merge_sort_keys(coop::this_warp(), keys, +/// coop::MergeSortOptions::default().temp_storage(scratch)); +/// } +/// ``` +/// ```compile_fail,E0277 +/// use cuda_coop as coop; +/// fn invalid(keys: &coop::ThreadData) { +/// let group = coop::this_warp().group_by_non_exhaustive::<17>(); +/// coop::merge_sort_keys(group, keys, Default::default()); +/// } +/// ``` +pub trait MergeSortGroup: GroupScope + sealed::Group {} +impl MergeSortGroup for Block {} +impl MergeSortGroup<()> 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 N: 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:literal $(, $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 N: usize>( + registration: R, _storage: S, uniform: U, payload: P, + ) { + // SAFETY: MergeSortGroup permits only implicit Warp storage; + // the genuine marker checks participation and the exact width. + unsafe { library::checked_warp_payload_call_v1::(registration, uniform, payload) } + } + } + }; +} +warp_group!(Warp, 32); +macro_rules! logical_groups { + ($($width:literal),+) => {$ ( + impl MergeSortGroup<()> for ThreadsWithinWarp<$width, E> {} + warp_group!(ThreadsWithinWarp<$width, E>, $width, E); + )+}; +} +logical_groups!(1, 2, 4, 8, 16, 32); + +/// Order and optional count/sentinel for [`merge_sort_keys`] or [`merge_sort_pairs`]. +/// +/// `T` is inferred from the input. Only canonical [`Default`] is implemented: +/// ascending order, every input and implicit scratch. Builders consume options +/// and preserve the other selections, including the actual order value and storage. +/// +/// ``` +/// use cuda_coop as coop; +/// let _: coop::MergeSortOptions = Default::default(); +/// let _ = coop::MergeSortOptions::::default() +/// .order(coop::Descending).valid_items(7, i64::MIN); +/// let _ = coop::MergeSortOptions::::default() +/// .valid_items(7, i64::MIN).order(coop::Descending); +/// ``` +pub struct MergeSortOptions { + order: O, + prefix: Prefix, + storage: S, + selection: PhantomData, +} + +impl Default for MergeSortOptions { + fn default() -> Self { + Self { + order: Ascending, + prefix: (), + storage: (), + selection: PhantomData, + } + } +} + +impl MergeSortOptions { + /// Select a built-in order or stateless i32 Keys comparator, retaining other options. + /// + /// Closures that call methods on their keys may need explicit `i32` parameter types. + pub fn order(self, order: O2) -> MergeSortOptions { + MergeSortOptions { + order, + prefix: self.prefix, + storage: self.storage, + selection: PhantomData, + } + } + + /// Sort this blocked prefix using a sentinel of exactly the key type. + /// + /// Count and sentinel must be uniform across the selected group. Count + /// must be in `0..=group_size * ITEMS_PER_THREAD`; it is checked as signed + /// `i64` before conversion to CUB's int. Widen smaller integer counts + /// losslessly. The sentinel must be ordered after every valid key in the + /// selected order. A literal sentinel is inferred from the input key type. + /// + /// Every input slot must still be initialized, including at count zero. + /// The result has initialized slots only within this prefix. Calling this + /// builder again replaces both count and sentinel. + /// + /// ```compile_fail,E0308 + /// use cuda_coop as coop; + /// let _ = coop::MergeSortOptions::::default().valid_items(7, 0_i32); + /// ``` + pub fn valid_items(self, count: i64, sentinel: T) -> MergeSortOptions { + MergeSortOptions { + order: self.order, + prefix: (count, sentinel), + storage: self.storage, + selection: PhantomData, + } + } + /// Borrow existing Block scratch, retaining order and paired controls. + /// + /// Manual scratch requires Block completion before conflicting reuse. + /// Auto-sync storage inserts its existing trailing barrier. Warp rejects + /// this option; consumed options permit later ordinary mutable reborrows. + pub fn temp_storage( + self, + storage: &mut TempStorage, + ) -> MergeSortOptions> { + MergeSortOptions { + order: self.order, + prefix: self.prefix, + storage, + selection: PhantomData, + } + } +} + +/// Merge-sort a group's keys into a fresh blocked payload. +/// +/// Supports all ten [`Numeric`] types and a positive `ITEMS_PER_THREAD`. The +/// result preserves the input type, extent and minimum alignment. Every input +/// slot must be initialized, even outside a selected prefix or at count zero; +/// input values and validity bits remain unchanged through shared borrows. +/// +/// All members of the selected group participate. Complete XYZ Blocks require a +/// power-of-two total; Warp supports physical and logical widths described by +/// [`MergeSortGroup`]. Input/output index is `rank * ITEMS_PER_THREAD + item`, +/// using x-fastest Block rank or rank within each Warp group. The provider checks +/// tile capacity and CUB's +/// padded scratch extent before converting dimensions to int. +/// +/// Default options sort the full tile ascending and initialize every result. +/// A partial count and typed sentinel are selected together with +/// [`MergeSortOptions::valid_items`]. Only the result's blocked prefix is then +/// initialized; [`ThreadData::get`] returns `None` for the suffix. Count zero +/// still participates, validates the full input and returns an empty result. +/// Invalid runtime counts or unreadable inputs trap before native numeric reads. +/// +/// Equal keys have no stable-tie guarantee. Floating keys must obey the +/// comparison's strict weak ordering requirements; no NaN total order is +/// supplied. The sentinel must follow valid keys in the chosen order. +/// +/// Implicit scratch is compiler-planned, with synchronization before dirty reuse. +/// Explicit Block storage retains its manual/auto-sync/exclusive policy. CUB +/// synchronization does not complete unrelated global accesses; use +/// `cuda_device::thread::sync_threads()` when needed. Full-input Block i32 Keys +/// also accepts the stateless comparators described by [`MergeSortOrder`]. +/// Execution requires an +/// Oxide-compiled CUDA kernel. Generic T/N/ALIGN callers need no +/// `generic_const_exprs` feature. +/// +/// ```no_run +/// use cuda_coop as coop; +/// fn sorted( +/// block: coop::ThreadGroup, +/// keys: &coop::ThreadData, +/// ) -> coop::ThreadData +/// where coop::Alignment: coop::SupportedAlignment { +/// coop::merge_sort_keys(block, keys, Default::default()) +/// } +/// ``` +#[inline(always)] +pub fn merge_sort_keys( + _group: ThreadGroup, + keys: &ThreadData, + options: MergeSortOptions, +) -> ThreadData +where + T: Numeric, + O: MergeSortOrder, + P: MergeSortPrefix, + S: TempStorageLike, + G: MergeSortGroup, + Alignment: SupportedAlignment, +{ + O::keys::(keys, options) +} + +/// Merge-sort associated keys and values into two fresh blocked payloads. +/// +/// Keys support all ten [`Numeric`] types, and values independently support +/// all ten [`MergeSortValue`] types. Both inputs share `ITEMS_PER_THREAD`; +/// their minimum alignments `KEY_ALIGN` and `VALUE_ALIGN` are independent. +/// Each output keeps its corresponding input's type, extent and alignment. +/// Both borrowed inputs and their validity maps remain unchanged. +/// +/// The group, built-in order, count/sentinel, storage and completion rules +/// are the same as [`merge_sort_keys`]. Default options sort the complete tile +/// ascending with implicit scratch. Partial options take a uniform signed64 +/// count and a sentinel of the key type. ALL keys and values must be initialized, +/// even beyond the prefix or at count zero. Invalid counts or either missing +/// input bit trap before numeric reads. Both outputs have the same valid +/// prefix; suffix [`ThreadData::get`] calls return `None`. +/// Callable comparator orders are rejected at compile time for Pairs. +/// +/// Each value follows its associated key without numeric conversion. Equal +/// keys have no stable-tie guarantee. Floating keys must satisfy strict weak +/// ordering; no NaN total order is supplied. Values are not compared and may +/// contain any numeric bit pattern. Scratch is queried for the pair primitive, +/// including its padded value storage, and cannot be inferred from key size. +/// Execution requires an Oxide-compiled CUDA kernel; no `generic_const_exprs` +/// feature is needed for generic K/V/N/alignment callers. +/// +/// ```no_run +/// use cuda_coop as coop; +/// fn sorted( +/// block: coop::ThreadGroup, +/// keys: &coop::ThreadData, values: &coop::ThreadData, +/// ) -> (coop::ThreadData, coop::ThreadData) +/// where coop::Alignment: coop::SupportedAlignment, +/// coop::Alignment: coop::SupportedAlignment { +/// coop::merge_sort_pairs(block, keys, values, Default::default()) +/// } +/// ``` +#[inline(always)] +pub fn merge_sort_pairs< + K, + V, + G, + O, + P, + S, + const ITEMS_PER_THREAD: usize, + const KEY_ALIGN: usize, + const VALUE_ALIGN: usize, +>( + _group: ThreadGroup, + keys: &ThreadData, + values: &ThreadData, + options: MergeSortOptions, +) -> ( + ThreadData, + ThreadData, +) +where + K: Numeric, + V: MergeSortValue, + O: MergeSortOrder, + P: MergeSortPrefix, + S: TempStorageLike, + G: MergeSortGroup, + Alignment: SupportedAlignment, + Alignment: SupportedAlignment, +{ + boundary::merge_sort_pairs::( + keys, values, options, + ) +} + +mod sealed { + use super::*; + pub trait Value { + const ASCENDING_FULL_REGISTRATION: &'static str; + const DESCENDING_FULL_REGISTRATION: &'static str; + const ASCENDING_PARTIAL_REGISTRATION: &'static str; + const DESCENDING_PARTIAL_REGISTRATION: &'static str; + const WARP_ASCENDING_FULL_REGISTRATION: &'static str; + const WARP_DESCENDING_FULL_REGISTRATION: &'static str; + const WARP_ASCENDING_PARTIAL_REGISTRATION: &'static str; + const WARP_DESCENDING_PARTIAL_REGISTRATION: &'static str; + } + pub trait Group { + const WARP: bool; + unsafe fn call< + T: Copy, + R: Fn() -> &'static str + Copy, + S, + U: Copy, + P: Copy, + const N: usize, + >( + registration: R, + storage: S, + uniform: U, + payload: P, + ); + } + pub trait Order { + const DESCENDING: bool; + fn keys( + input: &ThreadData, + options: MergeSortOptions, + ) -> ThreadData + where + Self: Sized, + T: Numeric, + P: MergeSortPrefix, + S: TempStorageLike, + G: MergeSortGroup, + Alignment: SupportedAlignment; + } + pub trait Prefix: Copy { + const PARTIAL: bool; + } +} + +impl sealed::Order for Ascending { + const DESCENDING: bool = false; + fn keys( + input: &ThreadData, + options: MergeSortOptions, + ) -> ThreadData + where + T: Numeric, + P: MergeSortPrefix, + S: TempStorageLike, + G: MergeSortGroup, + Alignment: SupportedAlignment, + { + boundary::merge_sort_keys::(input, options) + } +} +impl MergeSortOrder for Ascending {} +impl sealed::Order for Descending { + const DESCENDING: bool = true; + fn keys( + input: &ThreadData, + options: MergeSortOptions, + ) -> ThreadData + where + T: Numeric, + P: MergeSortPrefix, + S: TempStorageLike, + G: MergeSortGroup, + Alignment: SupportedAlignment, + { + boundary::merge_sort_keys::(input, options) + } +} +impl MergeSortOrder for Descending {} + +struct CallableBounds(PhantomData<(T, G, P, F)>); +impl, F> CallableBounds { + const VALID: () = { + assert!(!G::WARP, "callable Merge Sort requires Block"); + assert!(!P::PARTIAL, "callable Merge Sort requires full input"); + assert!( + core::mem::size_of::() == 0, + "callable Merge Sort requires a stateless operator" + ); + }; +} +impl bool + Copy> sealed::Order for F { + const DESCENDING: bool = panic!("callable Merge Sort Pairs is unsupported"); + fn keys( + input: &ThreadData, + options: MergeSortOptions, + ) -> ThreadData + where + T: Numeric, + P: MergeSortPrefix, + S: TempStorageLike, + G: MergeSortGroup, + Alignment: SupportedAlignment, + { + let () = CallableBounds::::VALID; + boundary::callable(input, options.order, options.storage) + } +} +impl bool + Copy> MergeSortOrder for F {} +impl sealed::Prefix for () { + const PARTIAL: bool = false; +} +impl MergeSortPrefix for () {} +impl sealed::Prefix for (i64, T) { + const PARTIAL: bool = true; +} +impl MergeSortPrefix for (i64, T) {} + +macro_rules! pair_values { + ($(($ty:ty, $ascending_full:literal, $descending_full:literal, $ascending_partial:literal, $descending_partial:literal, $warp_ascending_full:literal, $warp_descending_full:literal, $warp_ascending_partial:literal, $warp_descending_partial:literal)), + $(,)?) => {$ ( + impl sealed::Value for $ty { + const ASCENDING_FULL_REGISTRATION: &'static str = library::registration!($ascending_full); + const DESCENDING_FULL_REGISTRATION: &'static str = library::registration!($descending_full); + const ASCENDING_PARTIAL_REGISTRATION: &'static str = library::registration!($ascending_partial); + const DESCENDING_PARTIAL_REGISTRATION: &'static str = library::registration!($descending_partial); + const WARP_ASCENDING_FULL_REGISTRATION: &'static str = library::registration!($warp_ascending_full); + const WARP_DESCENDING_FULL_REGISTRATION: &'static str = library::registration!($warp_descending_full); + const WARP_ASCENDING_PARTIAL_REGISTRATION: &'static str = library::registration!($warp_ascending_partial); + const WARP_DESCENDING_PARTIAL_REGISTRATION: &'static str = library::registration!($warp_descending_partial); + } + impl MergeSortValue for $ty {} + )+}; +} +pair_values!( + ( + i8, + "COOP_MERGE_SORT_PAIRS_ASCENDING_FULL_I8", + "COOP_MERGE_SORT_PAIRS_DESCENDING_FULL_I8", + "COOP_MERGE_SORT_PAIRS_ASCENDING_PARTIAL_I8", + "COOP_MERGE_SORT_PAIRS_DESCENDING_PARTIAL_I8", + "COOP_MERGE_SORT_PAIRS_WARP_ASCENDING_FULL_I8", + "COOP_MERGE_SORT_PAIRS_WARP_DESCENDING_FULL_I8", + "COOP_MERGE_SORT_PAIRS_WARP_ASCENDING_PARTIAL_I8", + "COOP_MERGE_SORT_PAIRS_WARP_DESCENDING_PARTIAL_I8" + ), + ( + u8, + "COOP_MERGE_SORT_PAIRS_ASCENDING_FULL_U8", + "COOP_MERGE_SORT_PAIRS_DESCENDING_FULL_U8", + "COOP_MERGE_SORT_PAIRS_ASCENDING_PARTIAL_U8", + "COOP_MERGE_SORT_PAIRS_DESCENDING_PARTIAL_U8", + "COOP_MERGE_SORT_PAIRS_WARP_ASCENDING_FULL_U8", + "COOP_MERGE_SORT_PAIRS_WARP_DESCENDING_FULL_U8", + "COOP_MERGE_SORT_PAIRS_WARP_ASCENDING_PARTIAL_U8", + "COOP_MERGE_SORT_PAIRS_WARP_DESCENDING_PARTIAL_U8" + ), + ( + i16, + "COOP_MERGE_SORT_PAIRS_ASCENDING_FULL_I16", + "COOP_MERGE_SORT_PAIRS_DESCENDING_FULL_I16", + "COOP_MERGE_SORT_PAIRS_ASCENDING_PARTIAL_I16", + "COOP_MERGE_SORT_PAIRS_DESCENDING_PARTIAL_I16", + "COOP_MERGE_SORT_PAIRS_WARP_ASCENDING_FULL_I16", + "COOP_MERGE_SORT_PAIRS_WARP_DESCENDING_FULL_I16", + "COOP_MERGE_SORT_PAIRS_WARP_ASCENDING_PARTIAL_I16", + "COOP_MERGE_SORT_PAIRS_WARP_DESCENDING_PARTIAL_I16" + ), + ( + u16, + "COOP_MERGE_SORT_PAIRS_ASCENDING_FULL_U16", + "COOP_MERGE_SORT_PAIRS_DESCENDING_FULL_U16", + "COOP_MERGE_SORT_PAIRS_ASCENDING_PARTIAL_U16", + "COOP_MERGE_SORT_PAIRS_DESCENDING_PARTIAL_U16", + "COOP_MERGE_SORT_PAIRS_WARP_ASCENDING_FULL_U16", + "COOP_MERGE_SORT_PAIRS_WARP_DESCENDING_FULL_U16", + "COOP_MERGE_SORT_PAIRS_WARP_ASCENDING_PARTIAL_U16", + "COOP_MERGE_SORT_PAIRS_WARP_DESCENDING_PARTIAL_U16" + ), + ( + i32, + "COOP_MERGE_SORT_PAIRS_ASCENDING_FULL_I32", + "COOP_MERGE_SORT_PAIRS_DESCENDING_FULL_I32", + "COOP_MERGE_SORT_PAIRS_ASCENDING_PARTIAL_I32", + "COOP_MERGE_SORT_PAIRS_DESCENDING_PARTIAL_I32", + "COOP_MERGE_SORT_PAIRS_WARP_ASCENDING_FULL_I32", + "COOP_MERGE_SORT_PAIRS_WARP_DESCENDING_FULL_I32", + "COOP_MERGE_SORT_PAIRS_WARP_ASCENDING_PARTIAL_I32", + "COOP_MERGE_SORT_PAIRS_WARP_DESCENDING_PARTIAL_I32" + ), + ( + u32, + "COOP_MERGE_SORT_PAIRS_ASCENDING_FULL_U32", + "COOP_MERGE_SORT_PAIRS_DESCENDING_FULL_U32", + "COOP_MERGE_SORT_PAIRS_ASCENDING_PARTIAL_U32", + "COOP_MERGE_SORT_PAIRS_DESCENDING_PARTIAL_U32", + "COOP_MERGE_SORT_PAIRS_WARP_ASCENDING_FULL_U32", + "COOP_MERGE_SORT_PAIRS_WARP_DESCENDING_FULL_U32", + "COOP_MERGE_SORT_PAIRS_WARP_ASCENDING_PARTIAL_U32", + "COOP_MERGE_SORT_PAIRS_WARP_DESCENDING_PARTIAL_U32" + ), + ( + i64, + "COOP_MERGE_SORT_PAIRS_ASCENDING_FULL_I64", + "COOP_MERGE_SORT_PAIRS_DESCENDING_FULL_I64", + "COOP_MERGE_SORT_PAIRS_ASCENDING_PARTIAL_I64", + "COOP_MERGE_SORT_PAIRS_DESCENDING_PARTIAL_I64", + "COOP_MERGE_SORT_PAIRS_WARP_ASCENDING_FULL_I64", + "COOP_MERGE_SORT_PAIRS_WARP_DESCENDING_FULL_I64", + "COOP_MERGE_SORT_PAIRS_WARP_ASCENDING_PARTIAL_I64", + "COOP_MERGE_SORT_PAIRS_WARP_DESCENDING_PARTIAL_I64" + ), + ( + u64, + "COOP_MERGE_SORT_PAIRS_ASCENDING_FULL_U64", + "COOP_MERGE_SORT_PAIRS_DESCENDING_FULL_U64", + "COOP_MERGE_SORT_PAIRS_ASCENDING_PARTIAL_U64", + "COOP_MERGE_SORT_PAIRS_DESCENDING_PARTIAL_U64", + "COOP_MERGE_SORT_PAIRS_WARP_ASCENDING_FULL_U64", + "COOP_MERGE_SORT_PAIRS_WARP_DESCENDING_FULL_U64", + "COOP_MERGE_SORT_PAIRS_WARP_ASCENDING_PARTIAL_U64", + "COOP_MERGE_SORT_PAIRS_WARP_DESCENDING_PARTIAL_U64" + ), + ( + f32, + "COOP_MERGE_SORT_PAIRS_ASCENDING_FULL_F32", + "COOP_MERGE_SORT_PAIRS_DESCENDING_FULL_F32", + "COOP_MERGE_SORT_PAIRS_ASCENDING_PARTIAL_F32", + "COOP_MERGE_SORT_PAIRS_DESCENDING_PARTIAL_F32", + "COOP_MERGE_SORT_PAIRS_WARP_ASCENDING_FULL_F32", + "COOP_MERGE_SORT_PAIRS_WARP_DESCENDING_FULL_F32", + "COOP_MERGE_SORT_PAIRS_WARP_ASCENDING_PARTIAL_F32", + "COOP_MERGE_SORT_PAIRS_WARP_DESCENDING_PARTIAL_F32" + ), + ( + f64, + "COOP_MERGE_SORT_PAIRS_ASCENDING_FULL_F64", + "COOP_MERGE_SORT_PAIRS_DESCENDING_FULL_F64", + "COOP_MERGE_SORT_PAIRS_ASCENDING_PARTIAL_F64", + "COOP_MERGE_SORT_PAIRS_DESCENDING_PARTIAL_F64", + "COOP_MERGE_SORT_PAIRS_WARP_ASCENDING_FULL_F64", + "COOP_MERGE_SORT_PAIRS_WARP_DESCENDING_FULL_F64", + "COOP_MERGE_SORT_PAIRS_WARP_ASCENDING_PARTIAL_F64", + "COOP_MERGE_SORT_PAIRS_WARP_DESCENDING_PARTIAL_F64" + ), +); diff --git a/rust/cuda-coop/src/merge_sort/boundary.rs b/rust/cuda-coop/src/merge_sort/boundary.rs new file mode 100644 index 000000000000..41a7cf0dadc1 --- /dev/null +++ b/rust/cuda-coop/src/merge_sort/boundary.rs @@ -0,0 +1,217 @@ +// Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +use super::{MergeSortGroup, MergeSortOptions, MergeSortOrder, MergeSortPrefix, MergeSortValue}; +use crate::{Alignment, Numeric, SupportedAlignment, TempStorageLike, ThreadData}; +use cuda_device::library; + +fn ascending_full() -> &'static str { + library::registration!("COOP_MERGE_SORT_KEYS_ASCENDING_FULL") +} +fn descending_full() -> &'static str { + library::registration!("COOP_MERGE_SORT_KEYS_DESCENDING_FULL") +} +fn ascending_partial() -> &'static str { + library::registration!("COOP_MERGE_SORT_KEYS_ASCENDING_PARTIAL") +} +fn descending_partial() -> &'static str { + library::registration!("COOP_MERGE_SORT_KEYS_DESCENDING_PARTIAL") +} + +fn warp_ascending_full() -> &'static str { + library::registration!("COOP_MERGE_SORT_KEYS_WARP_ASCENDING_FULL") +} +fn warp_descending_full() -> &'static str { + library::registration!("COOP_MERGE_SORT_KEYS_WARP_DESCENDING_FULL") +} +fn warp_ascending_partial() -> &'static str { + library::registration!("COOP_MERGE_SORT_KEYS_WARP_ASCENDING_PARTIAL") +} +fn warp_descending_partial() -> &'static str { + library::registration!("COOP_MERGE_SORT_KEYS_WARP_DESCENDING_PARTIAL") +} + +#[inline(always)] +pub(super) fn merge_sort_keys( + input: &ThreadData, + options: MergeSortOptions, +) -> ThreadData +where + T: Numeric, + O: MergeSortOrder, + P: MergeSortPrefix, + S: TempStorageLike, + G: MergeSortGroup, + Alignment: SupportedAlignment, +{ + let mut output = ThreadData::new(); + let (input_pointer, input_bits) = input.block_io_read_parts(); + let (output_pointer, output_bits) = output.block_io_parts(); + let payload = (input_pointer, output_pointer, input_bits, output_bits); + let _order = options.order; + let storage = options.storage.into_raw(); + macro_rules! call { + ($registration:ident) => { + G::call::($registration, storage, options.prefix, payload) + }; + } + // SAFETY: borrowed input and fresh same-shape output are distinct aligned + // lane-local buffers with independent N-bit maps. The closed recipe checks + // wide count bounds and ALL input bits, even at count zero, before numeric + // reads. It sorts a private copy, marks only the written output prefix and + // clears all other output bits. No pointer escapes. Sealed options select + // the exact full/partial scalar ABI. The genuine marker proves complete + // group participation and uniform controls, and plans scratch/reuse; the + // provider validates geometry, integer capacity and full padded layout. + unsafe { + match (G::WARP, O::DESCENDING, P::PARTIAL) { + (false, false, false) => call!(ascending_full), + (false, true, false) => call!(descending_full), + (false, false, true) => call!(ascending_partial), + (false, true, true) => call!(descending_partial), + (true, false, false) => call!(warp_ascending_full), + (true, true, false) => call!(warp_descending_full), + (true, false, true) => call!(warp_ascending_partial), + (true, true, true) => call!(warp_descending_partial), + } + } + output +} + +fn pairs_ascending_full() -> &'static str { + V::ASCENDING_FULL_REGISTRATION +} +fn pairs_descending_full() -> &'static str { + V::DESCENDING_FULL_REGISTRATION +} +fn pairs_ascending_partial() -> &'static str { + V::ASCENDING_PARTIAL_REGISTRATION +} +fn pairs_descending_partial() -> &'static str { + V::DESCENDING_PARTIAL_REGISTRATION +} +fn pairs_warp_ascending_full() -> &'static str { + V::WARP_ASCENDING_FULL_REGISTRATION +} +fn pairs_warp_descending_full() -> &'static str { + V::WARP_DESCENDING_FULL_REGISTRATION +} +fn pairs_warp_ascending_partial() -> &'static str { + V::WARP_ASCENDING_PARTIAL_REGISTRATION +} +fn pairs_warp_descending_partial() -> &'static str { + V::WARP_DESCENDING_PARTIAL_REGISTRATION +} + +#[inline(always)] +pub(super) fn merge_sort_pairs( + keys: &ThreadData, + values: &ThreadData, + options: MergeSortOptions, +) -> (ThreadData, ThreadData) +where + K: Numeric, + V: MergeSortValue, + O: MergeSortOrder, + P: MergeSortPrefix, + S: TempStorageLike, + G: MergeSortGroup, + Alignment: SupportedAlignment, + Alignment: SupportedAlignment, +{ + let mut output_keys = ThreadData::new(); + let mut output_values = ThreadData::new(); + let (keys_pointer, keys_bits) = keys.block_io_read_parts(); + let (values_pointer, values_bits) = values.block_io_read_parts(); + let (output_keys_pointer, output_keys_bits) = output_keys.block_io_parts(); + let (output_values_pointer, output_values_bits) = output_values.block_io_parts(); + let payload = ( + keys_pointer, + output_keys_pointer, + values_pointer, + output_values_pointer, + keys_bits, + output_keys_bits, + values_bits, + output_values_bits, + ); + let _order = options.order; + let storage = options.storage.into_raw(); + macro_rules! call { + ($registration:ident) => { + G::call::($registration::, storage, options.prefix, payload) + }; + } + // SAFETY: shared inputs and two fresh same-shape outputs are live, aligned + // lane-local buffers with independent N-bit maps. The closed recipe checks + // wide count bounds and EVERY bit in both inputs before reading either + // array, including at count zero. It sorts private copies, writes paired + // valid-prefix outputs, and clears all other output bits. No pointer + // escapes. Sealed value constants and options preserve exact K/V/count/ + // sentinel ABI. The genuine marker proves participation, uniform controls + // and storage reuse; the provider queries complete pair scratch. + unsafe { + match (G::WARP, O::DESCENDING, P::PARTIAL) { + (false, false, false) => call!(pairs_ascending_full), + (false, true, false) => call!(pairs_descending_full), + (false, false, true) => call!(pairs_ascending_partial), + (false, true, true) => call!(pairs_descending_partial), + (true, false, false) => call!(pairs_warp_ascending_full), + (true, true, false) => call!(pairs_warp_descending_full), + (true, false, true) => call!(pairs_warp_ascending_partial), + (true, true, true) => call!(pairs_warp_descending_partial), + } + } + (output_keys, output_values) +} + +fn callable_full() -> &'static str { + library::registration!("COOP_MERGE_SORT_KEYS_CALLABLE_FULL") +} + +unsafe extern "C" fn compare_thunk bool>( + environment: *const F, + left: i32, + right: i32, +) -> i32 { + // SAFETY: the marker borrows the actual operator for this synchronous call. + (unsafe { (&*environment)(left, right) }) as i32 +} + +#[inline(always)] +pub(super) fn callable< + T: Numeric, + F: Fn(i32, i32) -> bool + Copy, + S: TempStorageLike, + const N: usize, + const ALIGN: usize, +>( + input: &ThreadData, + operator: F, + storage: S, +) -> ThreadData +where + Alignment: SupportedAlignment, +{ + let mut output = ThreadData::::new(); + let (input_pointer, input_bits) = input.block_io_read_parts(); + let (output_pointer, output_bits) = output.block_io_parts(); + let payload = (input_pointer, output_pointer, input_bits, output_bits); + let storage = storage.into_raw(); + // SAFETY: input and fresh same-shape output are distinct live payloads. The + // full recipe checks all input bits before copying, preserves input, and + // writes every output bit. The compiler admits the actual i32 callable, + // complete Block participation and scratch reuse. Its shared environment + // remains live for this synchronous call; the C thunk converts bool to i32. + unsafe { + library::checked_block_payload_callback_call_v1::( + callable_full, + storage, + (), + payload, + compare_thunk::, + &operator, + ) + } + output +} diff --git a/rust/cuda-coop/src/radix_sort.rs b/rust/cuda-coop/src/radix_sort.rs index 0eb95c97a112..316633b80e96 100644 --- a/rust/cuda-coop/src/radix_sort.rs +++ b/rust/cuda-coop/src/radix_sort.rs @@ -10,11 +10,11 @@ use crate::{ mod boundary; -/// Sort selected digits in ascending order (the default Radix Sort order). +/// Select ascending ordering. #[derive(Clone, Copy, Debug, Default)] pub struct Ascending; -/// Sort selected digits in descending order, preserving the order of ties. +/// Select descending ordering. #[derive(Clone, Copy, Debug, Default)] pub struct Descending; diff --git a/rust/cuda-coop/tests/check_load_store.py b/rust/cuda-coop/tests/check_load_store.py index 3f897769a8f6..f54a08a50103 100644 --- a/rust/cuda-coop/tests/check_load_store.py +++ b/rust/cuda-coop/tests/check_load_store.py @@ -1873,6 +1873,220 @@ ), ) +# Merge Sort uses complete groups, typed partial controls and existing scratch epochs. +merge_sort = "let result=coop::merge_sort_keys(block,&values,Default::default());" +merge_sort_manual = merge_sort.replace( + "Default::default()", "coop::MergeSortOptions::default().temp_storage(&mut scratch)" +) +cases.update( + merge_sort_generic=( + "fn sorted,O:coop::MergeSortOrder," + "P:coop::MergeSortPrefix,S:coop::TempStorageLike,const N:usize,const A:usize>" + "(g:coop::ThreadGroup,v:&coop::ThreadData,o:coop::MergeSortOptions)" + "->coop::ThreadData where coop::Alignment:coop::SupportedAlignment" + "{coop::merge_sort_keys(g,v,o)}" + + load + + "let _=sorted(block,&values,Default::default());" + "let group=coop::this_warp().group_by::<8>();" + "let _=sorted(group,&values,coop::MergeSortOptions::default()" + ".order(coop::Descending).valid_items(8,i32::MIN));" + "let _=sorted(group,&values,coop::MergeSortOptions::default()" + ".valid_items(8,i32::MIN).order(coop::Descending));", + None, + ), + merge_sort_storage_completed=( + load + + new_storage + + merge_sort_manual + + sync + + merge_sort_manual + + sync + + store, + None, + ), + merge_sort_auto_epoch_completed=( + reload.replace(".valid_items(valid)", ".valid_items(valid).oob_default(0)") + + sync + + new_storage.replace("new()", "new().auto_sync()") + + merge_sort_manual + + merge_sort_manual + + sync + + store, + None, + ), + merge_sort_nonpower_partial_block=(load + merge_sort, "power-of-two Block"), + merge_sort_warp_partial_geometry=( + load + merge_sort.replace("(block,", "(coop::this_warp(),"), + "complete physical warps", + ), + merge_sort_varying_count=( + load + + merge_sort.replace( + "Default::default()", + "coop::MergeSortOptions::default().valid_items(thread::threadIdx_x() as i64,i32::MAX)", + ), + "not proved block-uniform", + ), + merge_sort_varying_sentinel=( + load + + merge_sort.replace( + "Default::default()", + "coop::MergeSortOptions::default().valid_items(1,thread::threadIdx_x() as i32)", + ), + "not proved block-uniform", + ), + merge_sort_result_control=( + load + merge_sort + "if result[0]==0 {" + merge_sort + "}", + "lane-varying", + ), + merge_sort_storage_missing_barrier=( + load + new_storage + merge_sort_manual + merge_sort_manual, + "requires a full-block barrier", + ), + merge_sort_read_epoch=( + reload.replace(".valid_items(valid)", ".valid_items(valid).oob_default(0)") + + merge_sort + + store, + "requires a barrier covering both accesses", + ), +) + +# Callable Keys retains the existing order bound and admits only the first closed recipe. +merge_compare = "coop::merge_sort_keys(block,&values,coop::MergeSortOptions::default().order(|left,right| left" + "(block:coop::ThreadGroup,values:&coop::ThreadData," + "options:coop::MergeSortOptions)->coop::ThreadData" + "{coop::merge_sort_keys(block,values,options)}" + "fn before(left:i32,right:i32)->bool {leftright));" + + store.replace("&values", "&result"), + None, + ), + merge_sort_callable_warp=( + load + merge_compare.replace("(block,", "(coop::this_warp(),"), + "callable Merge Sort requires Block", + ), + merge_sort_callable_partial=( + load + merge_compare.replace(".order(", ".valid_items(1,i32::MAX).order("), + "callable Merge Sort requires full input", + ), + merge_sort_callable_pairs=( + load + + merge_compare.replace( + "merge_sort_keys(block,&values,", "merge_sort_pairs(block,&values,&values," + ), + "callable Merge Sort Pairs is unsupported", + ), + merge_sort_callable_capture=( + load + + "let shift=1_i32;" + + merge_compare.replace( + "|left,right| leftbool=|left,right| left::new();for i in 0..4 {keys.write(i,0); }" + + merge_compare.replace("&values", "&keys"), + "callback version one supports i32 only", + ), + merge_sort_callable_cuda_body=( + load + + merge_compare.replace("left," + "O:coop::MergeSortOrder,P:coop::MergeSortPrefix,S:coop::TempStorageLike," + "const N:usize,const KA:usize,const VA:usize>(g:coop::ThreadGroup," + "keys:&coop::ThreadData,values:&coop::ThreadData," + "o:coop::MergeSortOptions)" + "->(coop::ThreadData,coop::ThreadData) " + "where coop::Alignment:coop::SupportedAlignment,coop::Alignment:coop::SupportedAlignment" + "{coop::merge_sort_pairs(g,keys,values,o)}" + + load + + pair_values + + "let _=sorted(block,&values,&associated,Default::default());" + "let (result,pair_result)=sorted(block,&values,&associated,coop::MergeSortOptions::default()" + ".valid_items(valid as i64,i32::MIN).order(coop::Descending));" + "let _:coop::ThreadData=pair_result;" + + store.replace("&values", "&result"), + None, + ), + merge_pairs_storage_completed=( + load + + pair_values + + new_storage + + merge_pairs_manual.replace("result,pair_result", "first,first_values") + + sync + + merge_pairs_manual + + sync + + "values.write(0,first[0]+result[0]+first_values[0] as i32+pair_result[0] as i32);" + + "let mut scratch=coop::TempStorage::new().auto_sync();" + + merge_pairs_manual + + merge_pairs_manual + + sync + + store, + None, + ), + merge_pairs_unequal_extent=( + load + + pair_values.replace("f64,4,32", "f64,3,32").replace("0..4", "0..3") + + merge_pairs, + "mismatched types", + ), + merge_pairs_value_sentinel=( + load + + pair_values + + merge_pairs.replace( + "Default::default()", + "coop::MergeSortOptions::default().valid_items(1,0_f64)", + ), + "mismatched types", + ), + merge_pairs_value_result_type=( + load + + pair_values + + merge_pairs + + "let _:coop::ThreadData=pair_result;", + "mismatched types", + ), + merge_pairs_value_control=( + load + + pair_values + + merge_pairs + + "if pair_result[0]==0.0 {" + + merge_pairs + + "}", + "lane-varying", + ), + merge_pairs_storage_missing_barrier=( + load + pair_values + new_storage + merge_pairs_manual + merge_pairs_manual, + "requires a full-block barrier", + ), +) + if args.case: unknown = set(args.case).difference(cases) if unknown: @@ -1933,6 +2147,8 @@ case_prefix = case_prefix.replace( "domain=1,block=(128,1,1)", "domain=3,block=(4,4,8)" ) + if name == "merge_sort_warp_partial_geometry": + case_prefix = case_prefix.replace("block=(128,1,1)", "block=(48,1,1)") if name == "warp_partial_geometry": case_prefix = case_prefix.replace("block=(64,1,1)", "block=(48,1,1)") elif name.endswith("_partial_block"): diff --git a/rust/cuda-coop/tests/merge_sort_keys/.gitignore b/rust/cuda-coop/tests/merge_sort_keys/.gitignore new file mode 100644 index 000000000000..74bffa8af4a4 --- /dev/null +++ b/rust/cuda-coop/tests/merge_sort_keys/.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/merge_sort_keys/Cargo.lock b/rust/cuda-coop/tests/merge_sort_keys/Cargo.lock new file mode 100644 index 000000000000..ebce343eedd9 --- /dev/null +++ b/rust/cuda-coop/tests/merge_sort_keys/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-merge-sort-keys-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/merge_sort_keys/Cargo.toml b/rust/cuda-coop/tests/merge_sort_keys/Cargo.toml new file mode 100644 index 000000000000..a0a3a959efd8 --- /dev/null +++ b/rust/cuda-coop/tests/merge_sort_keys/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-merge-sort-keys-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/merge_sort_keys/src/main.rs b/rust/cuda-coop/tests/merge_sort_keys/src/main.rs new file mode 100644 index 000000000000..adc5e9af83c3 --- /dev/null +++ b/rust/cuda-coop/tests/merge_sort_keys/src/main.rs @@ -0,0 +1,645 @@ +// Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +// Each root observes fresh results, prefix validity and preserved input. +#![allow(clippy::too_many_arguments)] +use cuda_coop as coop; +use cuda_device::{kernel, launch_contract, thread}; +use cuda_host::cuda_module; + +#[cuda_module] +mod kernels { + use super::*; + + macro_rules! exercise { + ($n:expr, $alignment:literal, $threads:literal, $first_group:expr, + $second_group:expr, $first_options:expr, $second_options:expr, $manual:literal, + $input:ident, $first:ident, $second:ident, $keep:ident, $flags:ident, $changed:ident) => {{ + let block = coop::this_block(); + let offset = block.rank_at::() as usize * $threads * $n; + let mut keys = coop::ThreadData::<_, $n, $alignment>::new(); + coop::load( + block, + $input, + &mut keys, + coop::LoadOptions::default().offset(offset), + ); + let mut first = coop::merge_sort_keys($first_group, &keys, $first_options); + if $manual { + thread::sync_threads(); + } + coop::store( + block, + &mut $first, + &first, + coop::StoreOptions::default().offset(offset), + ); + first.write(0, 111 as _); + let second = coop::merge_sort_keys($second_group, &keys, $second_options); + if $manual { + thread::sync_threads(); + } + let mut visible = coop::ThreadData::<_, $n>::new(); + let mut flags = coop::ThreadData::<_, $n>::new(); + for item in 0..$n { + // Observe validity safely, never reading an invalid numeric suffix. + visible.write(item, second.get(item).unwrap_or(99 as _)); + flags.write( + item, + 2 * i32::from(keys.get(item).is_some()) + i32::from(second.get(item).is_some()), + ); + } + coop::store( + block, + &mut $second, + &visible, + coop::StoreOptions::default().offset(offset), + ); + coop::store( + block, + &mut $keep, + &keys, + coop::StoreOptions::default().offset(offset), + ); + coop::store( + block, + &mut $flags, + &flags, + coop::StoreOptions::default().offset(offset), + ); + coop::store( + block, + &mut $changed, + &first, + coop::StoreOptions::default().offset(offset), + ); + }}; + } + + #[kernel] + #[launch_contract(domain = 3, block = (4, 4, 4), dynamic_shared = inferred)] + pub fn f32_manual( + input: &[f32], + mut first_output: coop::DisjointSlice>, + mut second_output: coop::DisjointSlice>, + mut preserved: coop::DisjointSlice>, + mut validity: coop::DisjointSlice>, + mut mutated: coop::DisjointSlice>, + count: i64, + invalid: u32, + ) { + let block = coop::this_block(); + let mut scratch = coop::TempStorage::new(); + if invalid != 0 { + let rank = (thread::threadIdx_x() + + 4 * (thread::threadIdx_y() + 4 * thread::threadIdx_z())) + as usize; + let mut keys = coop::ThreadData::<_, ITEMS_PER_THREAD, 32>::new(); + for item in 0..ITEMS_PER_THREAD { + if !((invalid == 1 || invalid == 4) && rank == 63 && item == 8) { + keys.write(item, input[rank * ITEMS_PER_THREAD + item]); + } + } + let _ = coop::merge_sort_keys( + block, + &keys, + coop::MergeSortOptions::default() + .order(coop::Descending) + .valid_items(count, f32::MIN) + .temp_storage(&mut scratch), + ); + thread::sync_threads(); + // Independent report only: losing preflight must not be hidden by + // an invalid-result read or a device assertion. + let mut report = coop::ThreadData::<_, ITEMS_PER_THREAD>::new(); + for item in 0..ITEMS_PER_THREAD { + report.write(item, 99_f32); + } + coop::store(block, &mut mutated, &report, Default::default()); + return; + } + exercise!( + ITEMS_PER_THREAD, + 32, + 64, + block, + block, + coop::MergeSortOptions::default().temp_storage(&mut scratch), + coop::MergeSortOptions::default() + .valid_items(count, f32::MIN) + .order(coop::Descending) + .temp_storage(&mut scratch), + true, + input, + first_output, + second_output, + preserved, + validity, + mutated + ); + } + + #[kernel] + #[launch_contract(domain = 1, block = (64, 1, 1), dynamic_shared = inferred)] + pub fn f64_auto( + input: &[f64], + mut first_output: coop::DisjointSlice>, + mut second_output: coop::DisjointSlice>, + mut preserved: coop::DisjointSlice>, + mut validity: coop::DisjointSlice>, + mut mutated: coop::DisjointSlice>, + count: i64, + invalid: u32, + ) { + let block = coop::this_block(); + let mut scratch = coop::TempStorage::new().auto_sync(); + let _ = invalid; + exercise!( + ITEMS_PER_THREAD, + 16, + 64, + block, + block, + coop::MergeSortOptions::default().temp_storage(&mut scratch), + coop::MergeSortOptions::default() + .valid_items(count, f64::MIN) + .order(coop::Descending) + .temp_storage(&mut scratch), + false, + input, + first_output, + second_output, + preserved, + validity, + mutated + ); + } + + #[kernel] + #[launch_contract(domain = 1, block = (64, 1, 1), dynamic_shared = inferred)] + pub fn warp1( + input: &[u8], + mut first_output: coop::DisjointSlice>, + mut second_output: coop::DisjointSlice>, + mut preserved: coop::DisjointSlice>, + mut validity: coop::DisjointSlice>, + mut mutated: coop::DisjointSlice>, + count: i64, + invalid: u32, + ) { + let warp = coop::this_warp(); + let _ = invalid; + exercise!( + ITEMS_PER_THREAD, + 1, + 64, + warp.group_by::<1>(), + warp.group_by::<1>(), + coop::MergeSortOptions::default(), + coop::MergeSortOptions::default() + .valid_items(count, u8::MIN) + .order(coop::Descending), + false, + input, + first_output, + second_output, + preserved, + validity, + mutated + ); + } + + #[kernel] + #[launch_contract(domain = 1, block = (64, 1, 1), dynamic_shared = inferred)] + pub fn warp2( + input: &[i16], + mut first_output: coop::DisjointSlice>, + mut second_output: coop::DisjointSlice>, + mut preserved: coop::DisjointSlice>, + mut validity: coop::DisjointSlice>, + mut mutated: coop::DisjointSlice>, + count: i64, + invalid: u32, + ) { + let warp = coop::this_warp(); + let _ = invalid; + exercise!( + ITEMS_PER_THREAD, + 2, + 64, + warp.group_by::<2>(), + warp.group_by::<2>(), + coop::MergeSortOptions::default(), + coop::MergeSortOptions::default() + .valid_items(count, i16::MIN) + .order(coop::Descending), + false, + input, + first_output, + second_output, + preserved, + validity, + mutated + ); + } + + #[kernel] + #[launch_contract(domain = 1, block = (64, 1, 1), dynamic_shared = inferred)] + pub fn warp4( + input: &[u32], + mut first_output: coop::DisjointSlice>, + mut second_output: coop::DisjointSlice>, + mut preserved: coop::DisjointSlice>, + mut validity: coop::DisjointSlice>, + mut mutated: coop::DisjointSlice>, + count: i64, + invalid: u32, + ) { + let warp = coop::this_warp(); + let _ = invalid; + exercise!( + ITEMS_PER_THREAD, + 16, + 64, + warp.group_by::<4>(), + warp.group_by::<4>(), + coop::MergeSortOptions::default(), + coop::MergeSortOptions::default() + .valid_items(count, u32::MIN) + .order(coop::Descending), + false, + input, + first_output, + second_output, + preserved, + validity, + mutated + ); + } + + #[kernel] + #[launch_contract(domain = 3, block = (4, 4, 6), dynamic_shared = inferred)] + pub fn warp8( + input: &[i32], + mut first_output: coop::DisjointSlice>, + mut second_output: coop::DisjointSlice>, + mut preserved: coop::DisjointSlice>, + mut validity: coop::DisjointSlice>, + mut mutated: coop::DisjointSlice>, + count: i64, + invalid: u32, + ) { + let block = coop::this_block(); + let warp = coop::this_warp(); + if invalid != 0 { + let rank = (thread::threadIdx_x() + + 4 * (thread::threadIdx_y() + 4 * thread::threadIdx_z())) + as usize; + let mut keys = coop::ThreadData::<_, ITEMS_PER_THREAD, 32>::new(); + for item in 0..ITEMS_PER_THREAD { + if !((invalid == 1 || invalid == 4) && rank == 95 && item == 8) { + keys.write(item, input[rank * ITEMS_PER_THREAD + item]); + } + } + if invalid == 4 { + let _ = coop::merge_sort_keys(warp.group_by::<8>(), &keys, Default::default()); + } else { + let _ = coop::merge_sort_keys( + warp.group_by::<8>(), + &keys, + coop::MergeSortOptions::default() + .order(coop::Descending) + .valid_items(count, i32::MIN), + ); + } + // Independent report only: losing preflight must not be hidden by + // an invalid-result read or a device assertion. + let mut report = coop::ThreadData::<_, ITEMS_PER_THREAD>::new(); + for item in 0..ITEMS_PER_THREAD { + report.write(item, 99_i32); + } + coop::store(block, &mut mutated, &report, Default::default()); + return; + } + exercise!( + ITEMS_PER_THREAD, + 32, + 96, + warp.group_by::<8>(), + warp.group_by::<8>(), + coop::MergeSortOptions::default(), + coop::MergeSortOptions::default() + .valid_items(count, i32::MIN) + .order(coop::Descending), + false, + input, + first_output, + second_output, + preserved, + validity, + mutated + ); + } + + #[kernel] + #[launch_contract(domain = 1, block = (64, 1, 1), dynamic_shared = inferred)] + pub fn warp16( + input: &[u64], + mut first_output: coop::DisjointSlice>, + mut second_output: coop::DisjointSlice>, + mut preserved: coop::DisjointSlice>, + mut validity: coop::DisjointSlice>, + mut mutated: coop::DisjointSlice>, + count: i64, + invalid: u32, + ) { + let warp = coop::this_warp(); + let _ = invalid; + exercise!( + ITEMS_PER_THREAD, + 16, + 64, + warp.group_by::<16>(), + warp.group_by::<16>(), + coop::MergeSortOptions::default(), + coop::MergeSortOptions::default() + .valid_items(count, u64::MIN) + .order(coop::Descending), + false, + input, + first_output, + second_output, + preserved, + validity, + mutated + ); + } + + #[kernel] + #[launch_contract(domain = 1, block = (64, 1, 1), dynamic_shared = inferred)] + pub fn warp32( + input: &[i64], + mut first_output: coop::DisjointSlice>, + mut second_output: coop::DisjointSlice>, + mut preserved: coop::DisjointSlice>, + mut validity: coop::DisjointSlice>, + mut mutated: coop::DisjointSlice>, + count: i64, + invalid: u32, + ) { + let warp = coop::this_warp(); + let _ = invalid; + exercise!( + ITEMS_PER_THREAD, + 32, + 64, + warp, + warp.group_by::<32>(), + coop::MergeSortOptions::default(), + coop::MergeSortOptions::default() + .valid_items(count, i64::MIN) + .order(coop::Descending), + false, + input, + first_output, + second_output, + preserved, + validity, + mutated + ); + } +} + +fn check_sorted( + input: &[T], + output: &[T], + group: usize, + count: usize, + descending: bool, + bits: impl Fn(T) -> B, +) { + for (keys, sorted) in input.chunks_exact(group).zip(output.chunks_exact(group)) { + for pair in sorted[..count].windows(2) { + assert!(if descending { + pair[0] >= pair[1] + } else { + pair[0] <= pair[1] + }); + } + let mut expected: Vec<_> = keys[..count].iter().copied().map(&bits).collect(); + let mut observed: Vec<_> = sorted[..count].iter().copied().map(&bits).collect(); + expected.sort_unstable(); + observed.sort_unstable(); + // Equal floating keys can permute; their exact bit multiset must survive. + assert_eq!(observed, expected); + } +} + +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}; + let invalid = match std::env::args().nth(1).as_deref() { + None => 0, + Some("missing-zero-input") => 1, + Some("negative-count") => 2, + Some("wide-count") => 3, + Some("missing-warp-input") => 4, + Some(other) => panic!("unknown Merge Sort mode {other}"), + }; + init_device_contexts(0, 1)?; + let context = with_cuda_context(0, |context| context.clone())?; + let stream = context.default_stream(); + // SAFETY: typed declarations and embedded kernels compile together. + let module = unsafe { kernels::load_async(0)? }; + let (mut cases, mut sorted, mut kept, mut source_checks, mut flags, mut changed) = + (0, 0, 0, 0, 0, 0); + macro_rules! run { + ($t:ty, $n:literal, $threads:literal, $width:literal, $config:expr, + $prepare:ident, $launch:ident, $count:expr, $selected:expr, $data:expr, $bits:expr) => {{ + if invalid == 0 || $selected { + let blocks: usize = if invalid == 0 { 2 } else { 1 }; + let keys: Vec<$t> = (0..blocks * $threads * $n).map($data).collect(); + let bits = $bits; + let count: i64 = $count; + let input = DeviceBuffer::from_host(&stream, &keys)?; + let first = DeviceBuffer::from_host(&stream, &vec![77 as $t; keys.len()])?; + let second = DeviceBuffer::from_host(&stream, &vec![77 as $t; keys.len()])?; + let keep = DeviceBuffer::from_host(&stream, &vec![77 as $t; keys.len()])?; + let validity = DeviceBuffer::from_host(&stream, &vec![-1_i32; keys.len()])?; + let mutated = DeviceBuffer::from_host(&stream, &vec![77 as $t; keys.len()])?; + stream.synchronize()?; + let plan = module.$prepare::<$n>(($config)(blocks))?; + let (input, first, second, keep, validity, mutated) = module + .$launch::<$n, _, _, _, _, _, _>( + &plan, input, first, second, keep, validity, mutated, count, invalid, + ) + .sync()?; + stream.synchronize()?; + if invalid != 0 { + panic!("invalid Merge Sort mode {invalid} unexpectedly completed"); + } + let first = first.to_host_vec(&stream)?; + let second = second.to_host_vec(&stream)?; + let keep = keep.to_host_vec(&stream)?; + let source = input.to_host_vec(&stream)?; + let validity = validity.to_host_vec(&stream)?; + let mutated = mutated.to_host_vec(&stream)?; + check_sorted(&keys, &first, $width * $n, $width * $n, false, bits); + check_sorted(&keys, &second, $width * $n, count as usize, true, bits); + for (index, value) in second.iter().enumerate() { + #[allow(clippy::modulo_one)] // Width-one singleton groups are intentional. + let valid = index % ($width * $n) < count as usize; + assert_eq!(validity[index], 2 + i32::from(valid)); + if !valid { + assert_eq!(bits(*value), bits(99 as $t)); + } + assert_eq!(bits(keep[index]), bits(keys[index])); + assert_eq!(bits(source[index]), bits(keys[index])); + let expected = if index.is_multiple_of($n) { + 111 as $t + } else { + first[index] + }; + assert_eq!(bits(mutated[index]), bits(expected)); + } + cases += 1; + sorted += keys.len() + (keys.len() / ($width * $n)) * count as usize; + kept += keys.len(); + source_checks += keys.len(); + flags += keys.len(); + changed += keys.len(); + } + }}; + } + run!( + f32, + 9, + 64, + 64, + |blocks| LaunchConfig3D::auto_shared((blocks as u32, 1, 1), (4, 4, 4)), + prepare_f32_manual, + f32_manual_async_owned, + if invalid == 1 { + 0 + } else if invalid == 2 { + i64::MIN + } else { + 575 + }, + invalid == 1 || invalid == 2, + |i| if i % 11 == 0 { + -0.0_f32 + } else if i % 13 == 0 { + 0.0 + } else { + ((i * 17 + i / 576 * 7) % 31) as f32 - 15.0 + }, + f32::to_bits + ); + run!( + f64, + 4, + 64, + 64, + |blocks| LaunchConfig1D::auto_shared(blocks as u32, 64), + prepare_f64_auto, + f64_auto_async_owned, + 0, + false, + |i| if i % 11 == 0 { + -0.0_f64 + } else if i % 13 == 0 { + 0.0 + } else { + ((i * 17 + i / 256 * 7) % 31) as f64 - 15.0 + }, + f64::to_bits + ); + run!( + u8, + 1, + 64, + 1, + |blocks| LaunchConfig1D::auto_shared(blocks as u32, 64), + prepare_warp1, + warp1_async_owned, + 1, + false, + |i| ((i * 17 + i * 7) % 31 + 1) as u8, + |v: u8| v + ); + run!( + i16, + 4, + 64, + 2, + |blocks| LaunchConfig1D::auto_shared(blocks as u32, 64), + prepare_warp2, + warp2_async_owned, + 7, + false, + |i| ((i * 17 + i / 8 * 7) % 31) as i16 - 15, + |v: i16| v + ); + run!( + u32, + 4, + 64, + 4, + |blocks| LaunchConfig1D::auto_shared(blocks as u32, 64), + prepare_warp4, + warp4_async_owned, + 16, + false, + |i| ((i * 17 + i / 16 * 7) % 31 + 1) as u32, + |v: u32| v + ); + run!( + i32, + 9, + 96, + 8, + |blocks| LaunchConfig3D::auto_shared((blocks as u32, 1, 1), (4, 4, 6)), + prepare_warp8, + warp8_async_owned, + if invalid == 3 { (1_i64 << 32) + 1 } else { 71 }, + invalid == 3 || invalid == 4, + |i| ((i * 17 + i / 72 * 7) % 31) as i32 - 15, + |v: i32| v + ); + run!( + u64, + 4, + 64, + 16, + |blocks| LaunchConfig1D::auto_shared(blocks as u32, 64), + prepare_warp16, + warp16_async_owned, + 0, + false, + |i| ((i * 17 + i / 64 * 7) % 31 + 1) as u64, + |v: u64| v + ); + run!( + i64, + 4, + 64, + 32, + |blocks| LaunchConfig1D::auto_shared(blocks as u32, 64), + prepare_warp32, + warp32_async_owned, + 127, + false, + |i| ((i * 17 + i / 128 * 7) % 31) as i64 - 15, + |v: i64| v + ); + assert_eq!( + (cases, sorted, kept, source_checks, flags, changed), + (8, 10018, 5568, 5568, 5568, 5568) + ); + println!( + "merge_sort_keys: {cases} cases, {sorted} sorted keys, {kept} preserved keys, {source_checks} source keys, {flags} validity flags and {changed} retained mutated keys passed" + ); + Ok(()) +} diff --git a/rust/cuda-coop/tests/merge_sort_pairs/.gitignore b/rust/cuda-coop/tests/merge_sort_pairs/.gitignore new file mode 100644 index 000000000000..74bffa8af4a4 --- /dev/null +++ b/rust/cuda-coop/tests/merge_sort_pairs/.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/merge_sort_pairs/Cargo.lock b/rust/cuda-coop/tests/merge_sort_pairs/Cargo.lock new file mode 100644 index 000000000000..d0cf06d6e117 --- /dev/null +++ b/rust/cuda-coop/tests/merge_sort_pairs/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-merge-sort-pairs-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/merge_sort_pairs/Cargo.toml b/rust/cuda-coop/tests/merge_sort_pairs/Cargo.toml new file mode 100644 index 000000000000..c1c813e2a291 --- /dev/null +++ b/rust/cuda-coop/tests/merge_sort_pairs/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-merge-sort-pairs-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/merge_sort_pairs/src/main.rs b/rust/cuda-coop/tests/merge_sort_pairs/src/main.rs new file mode 100644 index 000000000000..0ce6e015e2c2 --- /dev/null +++ b/rust/cuda-coop/tests/merge_sort_pairs/src/main.rs @@ -0,0 +1,560 @@ +// Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +// Independent outputs expose association, validity and both surviving pairs. +#![allow(clippy::too_many_arguments)] +use cuda_coop as coop; +use cuda_device::{kernel, launch_contract, thread}; +use cuda_host::cuda_module; + +#[cuda_module] +mod kernels { + use super::*; + + macro_rules! exercise { + ($n:literal, $ka:literal, $va:literal, $b:literal, $group:expr, + $first_options:expr, $second_options:expr, $manual:literal, + $keys:ident, $values:ident, $first_keys:ident, $first_values:ident, + $second_keys:ident, $second_values:ident, $keep_keys:ident, + $keep_values:ident, $changed_keys:ident, $changed_values:ident, + $validity:ident) => {{ + let block = coop::this_block(); + let offset = block.rank_at::() as usize * $b * $n; + let mut keys = coop::ThreadData::<_, $n, $ka>::new(); + let mut values = coop::ThreadData::<_, $n, $va>::new(); + coop::load( + block, + $keys, + &mut keys, + coop::LoadOptions::default().offset(offset), + ); + coop::load( + block, + $values, + &mut values, + coop::LoadOptions::default().offset(offset), + ); + let (mut first_k, mut first_v) = + coop::merge_sort_pairs($group, &keys, &values, $first_options); + if $manual { + thread::sync_threads(); + } + coop::store( + block, + &mut $first_keys, + &first_k, + coop::StoreOptions::default().offset(offset), + ); + coop::store( + block, + &mut $first_values, + &first_v, + coop::StoreOptions::default().offset(offset), + ); + // Retain both fresh owners, including their mutations, over reuse. + first_k.write(0, 111 as _); + first_v.write(0, -7 as _); + let (second_k, second_v) = + coop::merge_sort_pairs($group, &keys, &values, $second_options); + if $manual { + thread::sync_threads(); + } + let mut visible_k = coop::ThreadData::<_, $n>::new(); + let mut visible_v = coop::ThreadData::<_, $n>::new(); + let mut flags = coop::ThreadData::<_, $n>::new(); + for item in 0..$n { + visible_k.write(item, second_k.get(item).unwrap_or(99 as _)); + visible_v.write(item, second_v.get(item).unwrap_or(99 as _)); + flags.write( + item, + i32::from(keys.get(item).is_some()) + + 2 * i32::from(values.get(item).is_some()) + + 4 * i32::from(second_k.get(item).is_some()) + + 8 * i32::from(second_v.get(item).is_some()), + ); + } + coop::store( + block, + &mut $second_keys, + &visible_k, + coop::StoreOptions::default().offset(offset), + ); + coop::store( + block, + &mut $second_values, + &visible_v, + coop::StoreOptions::default().offset(offset), + ); + coop::store( + block, + &mut $keep_keys, + &keys, + coop::StoreOptions::default().offset(offset), + ); + coop::store( + block, + &mut $keep_values, + &values, + coop::StoreOptions::default().offset(offset), + ); + coop::store( + block, + &mut $changed_keys, + &first_k, + coop::StoreOptions::default().offset(offset), + ); + coop::store( + block, + &mut $changed_values, + &first_v, + coop::StoreOptions::default().offset(offset), + ); + coop::store( + block, + &mut $validity, + &flags, + coop::StoreOptions::default().offset(offset), + ); + }}; + } + + #[kernel] + #[launch_contract(domain = 1, block = (64, 1, 1), dynamic_shared = inferred)] + pub fn f32_i8_warp( + keys: &[f32], + values: &[i8], + mut first_keys: coop::DisjointSlice>, + mut first_values: coop::DisjointSlice>, + mut second_keys: coop::DisjointSlice>, + mut second_values: coop::DisjointSlice>, + mut keep_keys: coop::DisjointSlice>, + mut keep_values: coop::DisjointSlice>, + mut changed_keys: coop::DisjointSlice>, + mut changed_values: coop::DisjointSlice>, + mut validity: coop::DisjointSlice>, + count: i64, + invalid: u32, + ) { + let block = coop::this_block(); + + if invalid != 0 { + let rank = thread::threadIdx_x() as usize; + let mut input_keys = coop::ThreadData::<_, 4, 32>::new(); + let mut input_values = coop::ThreadData::<_, 4, 1>::new(); + for item in 0..4 { + input_keys.write(item, keys[rank * 4 + item]); + if !(rank % 8 == 7 && item == 3) { + input_values.write(item, values[rank * 4 + item]); + } + } + let _ = coop::merge_sort_pairs( + coop::this_warp().group_by::<8>(), + &input_keys, + &input_values, + coop::MergeSortOptions::default(), + ); + thread::sync_threads(); + // Never read the invalid result. A lost preflight returns an + // independent report; the host then reports unexpected success. + let mut report = coop::ThreadData::<_, 4>::new(); + for item in 0..4 { + report.write(item, 99_f32); + } + coop::store(block, &mut changed_keys, &report, Default::default()); + return; + } + exercise!( + 4, + 32, + 1, + 64, + coop::this_warp().group_by::<8>(), + coop::MergeSortOptions::default(), + coop::MergeSortOptions::default() + .valid_items(count, f32::MIN) + .order(coop::Descending), + false, + keys, + values, + first_keys, + first_values, + second_keys, + second_values, + keep_keys, + keep_values, + changed_keys, + changed_values, + validity + ); + } + + #[kernel] + #[launch_contract(domain = 1, block = (64, 1, 1), dynamic_shared = inferred)] + pub fn i64_f64_manual( + keys: &[i64], + values: &[f64], + mut first_keys: coop::DisjointSlice>, + mut first_values: coop::DisjointSlice>, + mut second_keys: coop::DisjointSlice>, + mut second_values: coop::DisjointSlice>, + mut keep_keys: coop::DisjointSlice>, + mut keep_values: coop::DisjointSlice>, + mut changed_keys: coop::DisjointSlice>, + mut changed_values: coop::DisjointSlice>, + mut validity: coop::DisjointSlice>, + count: i64, + invalid: u32, + ) { + let block = coop::this_block(); + let mut scratch = coop::TempStorage::new(); + if invalid != 0 { + let rank = thread::threadIdx_x() as usize; + let mut input_keys = coop::ThreadData::<_, 9, 1>::new(); + let mut input_values = coop::ThreadData::<_, 9, 32>::new(); + for item in 0..9 { + input_keys.write(item, keys[rank * 9 + item]); + if !(rank == 63 && item == 8) { + input_values.write(item, values[rank * 9 + item]); + } + } + let _ = coop::merge_sort_pairs( + block, + &input_keys, + &input_values, + coop::MergeSortOptions::default() + .order(coop::Descending) + .valid_items(count, i64::MIN) + .temp_storage(&mut scratch), + ); + thread::sync_threads(); + // Never read the invalid result. A lost preflight returns an + // independent report; the host then reports unexpected success. + let mut report = coop::ThreadData::<_, 9>::new(); + for item in 0..9 { + report.write(item, 99_i64); + } + coop::store(block, &mut changed_keys, &report, Default::default()); + return; + } + exercise!( + 9, + 1, + 32, + 64, + block, + coop::MergeSortOptions::default().temp_storage(&mut scratch), + coop::MergeSortOptions::default() + .order(coop::Descending) + .valid_items(count, i64::MIN) + .temp_storage(&mut scratch), + true, + keys, + values, + first_keys, + first_values, + second_keys, + second_values, + keep_keys, + keep_values, + changed_keys, + changed_values, + validity + ); + } + + #[kernel] + #[launch_contract(domain = 1, block = (32, 1, 1), dynamic_shared = inferred)] + pub fn u32_i16_auto( + keys: &[u32], + values: &[i16], + mut first_keys: coop::DisjointSlice>, + mut first_values: coop::DisjointSlice>, + mut second_keys: coop::DisjointSlice>, + mut second_values: coop::DisjointSlice>, + mut keep_keys: coop::DisjointSlice>, + mut keep_values: coop::DisjointSlice>, + mut changed_keys: coop::DisjointSlice>, + mut changed_values: coop::DisjointSlice>, + mut validity: coop::DisjointSlice>, + count: i64, + invalid: u32, + ) { + let block = coop::this_block(); + let mut scratch = coop::TempStorage::new().auto_sync(); + let _ = invalid; + exercise!( + 1, + 32, + 16, + 32, + block, + coop::MergeSortOptions::default().temp_storage(&mut scratch), + coop::MergeSortOptions::default() + .valid_items(count, u32::MIN) + .order(coop::Descending) + .temp_storage(&mut scratch), + false, + keys, + values, + first_keys, + first_values, + second_keys, + second_values, + keep_keys, + keep_values, + changed_keys, + changed_values, + validity + ); + } +} + +// Keys are ordered, but equal keys may permute. Bitwise pair multisets check +// association and preserve value NaN payloads/signed zeros without a tie rule. +fn check_pairs( + keys: &[K], + values: &[V], + actual_keys: &[K], + actual_values: &[V], + descending: bool, + key_bits: impl Fn(K) -> u64, + value_bits: impl Fn(V) -> u64, +) { + assert!(actual_keys.windows(2).all(|w| if descending { + w[0] >= w[1] + } else { + w[0] <= w[1] + })); + let pairs = |k: &[K], v: &[V]| { + let mut result: Vec<_> = k + .iter() + .copied() + .zip(v.iter().copied()) + .map(|(k, v)| (key_bits(k), value_bits(v))) + .collect(); + result.sort_unstable(); + result + }; + assert_eq!(pairs(actual_keys, actual_values), pairs(keys, values)); +} + +fn f64_values(count: usize) -> Vec { + (0..count) + .map(|i| match i % 8 { + 0 => f64::from_bits(0x7ff8_0000_0000_0000 | (i as u64 + 1)), + 1 => f64::from_bits(0xfff8_0000_0000_0000 | (i as u64 + 1)), + 2 => 0.0, + 3 => -0.0, + 4 => f64::INFINITY, + 5 => f64::NEG_INFINITY, + _ => i as f64 + 0.25, + }) + .collect() +} + +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}; + let invalid = match std::env::args().nth(1).as_deref() { + None => 0, + Some("missing-value-zero") => 1, + Some("missing-value-warp") => 2, + Some(other) => panic!("unknown Merge Sort Pairs mode {other}"), + }; + init_device_contexts(0, 1)?; + let context = with_cuda_context(0, |context| context.clone())?; + let stream = context.default_stream(); + // SAFETY: typed declarations and embedded kernels are compiled together. + let module = unsafe { kernels::load_async(0)? }; + let (mut cases, mut sorted, mut preserved, mut sources, mut flags_checked, mut mutated) = + (0, 0, 0, 0, 0, 0); + macro_rules! run { + ($kt:ty, $vt:ty, $n:literal, $b:literal, $group:literal, $blocks:expr, + $prepare:ident, $launch:ident, $keys:expr, $values:expr, $count:expr, + $key_bits:expr, $value_bits:expr) => {{ + let keys: Vec<$kt> = $keys; + let values: Vec<$vt> = $values; + let count: usize = $count; + let group_items = $group * $n; + let key_bits = $key_bits; + let value_bits = $value_bits; + let kb = |v: &[$kt]| v.iter().copied().map(key_bits).collect::>(); + let vb = |v: &[$vt]| v.iter().copied().map(value_bits).collect::>(); + let source_k = DeviceBuffer::from_host(&stream, &keys)?; + let source_v = DeviceBuffer::from_host(&stream, &values)?; + let first_k = DeviceBuffer::from_host(&stream, &vec![17 as $kt; keys.len()])?; + let first_v = DeviceBuffer::from_host(&stream, &vec![17 as $vt; keys.len()])?; + let second_k = DeviceBuffer::from_host(&stream, &vec![17 as $kt; keys.len()])?; + let second_v = DeviceBuffer::from_host(&stream, &vec![17 as $vt; keys.len()])?; + let keep_k = DeviceBuffer::from_host(&stream, &vec![17 as $kt; keys.len()])?; + let keep_v = DeviceBuffer::from_host(&stream, &vec![17 as $vt; keys.len()])?; + let changed_k = DeviceBuffer::from_host(&stream, &vec![17 as $kt; keys.len()])?; + let changed_v = DeviceBuffer::from_host(&stream, &vec![17 as $vt; keys.len()])?; + let flags = DeviceBuffer::from_host(&stream, &vec![-1_i32; keys.len()])?; + stream.synchronize()?; + let plan = module.$prepare(LaunchConfig1D::auto_shared($blocks, $b))?; + let ( + source_k, + source_v, + first_k, + first_v, + second_k, + second_v, + keep_k, + keep_v, + changed_k, + changed_v, + flags, + ) = module + .$launch( + &plan, + source_k, + source_v, + first_k, + first_v, + second_k, + second_v, + keep_k, + keep_v, + changed_k, + changed_v, + flags, + count as i64, + invalid, + ) + .sync()?; + stream.synchronize()?; + if invalid != 0 { + panic!("invalid Merge Sort Pairs mode {invalid} unexpectedly completed"); + } + let first_k = first_k.to_host_vec(&stream)?; + let first_v = first_v.to_host_vec(&stream)?; + let second_k = second_k.to_host_vec(&stream)?; + let second_v = second_v.to_host_vec(&stream)?; + let flags = flags.to_host_vec(&stream)?; + for base in (0..keys.len()).step_by(group_items) { + let end = base + group_items; + let prefix = base + count; + check_pairs( + &keys[base..end], + &values[base..end], + &first_k[base..end], + &first_v[base..end], + false, + key_bits, + value_bits, + ); + check_pairs( + &keys[base..prefix], + &values[base..prefix], + &second_k[base..prefix], + &second_v[base..prefix], + true, + key_bits, + value_bits, + ); + assert!( + second_k[prefix..end] + .iter() + .all(|&v| key_bits(v) == key_bits(99 as $kt)) + ); + assert!( + second_v[prefix..end] + .iter() + .all(|&v| value_bits(v) == value_bits(99 as $vt)) + ); + assert!(flags[base..prefix].iter().all(|&v| v == 15)); + assert!(flags[prefix..end].iter().all(|&v| v == 3)); + } + assert_eq!(kb(&keep_k.to_host_vec(&stream)?), kb(&keys)); + assert_eq!(vb(&keep_v.to_host_vec(&stream)?), vb(&values)); + assert_eq!(kb(&source_k.to_host_vec(&stream)?), kb(&keys)); + assert_eq!(vb(&source_v.to_host_vec(&stream)?), vb(&values)); + let mut expected_k = first_k; + let mut expected_v = first_v; + for i in (0..keys.len()).step_by($n) { + expected_k[i] = 111 as $kt; + expected_v[i] = -7 as $vt; + } + assert_eq!(kb(&changed_k.to_host_vec(&stream)?), kb(&expected_k)); + assert_eq!(vb(&changed_v.to_host_vec(&stream)?), vb(&expected_v)); + cases += 1; + sorted += 2 * keys.len() + 2 * keys.len() / group_items * count; + preserved += 2 * keys.len(); + sources += 2 * keys.len(); + flags_checked += keys.len(); + mutated += 2 * keys.len(); + }}; + } + let blocks = if invalid == 0 { 2 } else { 1 }; + if invalid == 0 || invalid == 2 { + let len = blocks as usize * 64 * 4; + run!( + f32, + i8, + 4, + 64, + 8, + blocks, + prepare_f32_i8_warp, + f32_i8_warp_async_owned, + (0..len) + .map(|i| match i % 13 { + 0 => -0.0, + 1 => 0.0, + _ => ((i * 7) % 11) as f32 - 5.0, + }) + .collect(), + (0..len).map(|i| ((i * 17) % 101) as i8 - 50).collect(), + 31, + |v: f32| v.to_bits() as u64, + |v: i8| v as u8 as u64 + ); + } + if invalid == 0 || invalid == 1 { + let len = blocks as usize * 64 * 9; + run!( + i64, + f64, + 9, + 64, + 64, + blocks, + prepare_i64_f64_manual, + i64_f64_manual_async_owned, + (0..len).map(|i| ((i * 11) % 23) as i64 - 11).collect(), + f64_values(len), + if invalid == 1 { 0 } else { 575 }, + |v: i64| v as u64, + f64::to_bits + ); + } + if invalid == 0 { + for count in [0, 32] { + run!( + u32, + i16, + 1, + 32, + 32, + 2, + prepare_u32_i16_auto, + u32_i16_auto_async_owned, + (0..64).map(|i| 1 + (i * 5 % 13) as u32).collect(), + (0..64).map(|i| i as i16 - 32).collect(), + count, + |v: u32| v as u64, + |v: i16| v as u16 as u64 + ); + } + } + assert_eq!( + (cases, sorted, preserved, sources, flags_checked, mutated), + (4, 7004, 3584, 3584, 1792, 3584) + ); + println!( + "merge_sort_pairs: {cases} cases, {sorted} sorted key/value observations, {preserved} preserved values, {sources} source values, {flags_checked} validity flags, {mutated} retained mutated values passed" + ); + Ok(()) +}