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/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 ` +- :doc:`Rust Development ` + - :doc:`Contributor Docs ` - :doc:`Maintainer Docs ` 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/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/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/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 new file mode 100644 index 000000000000..e1fe66597cb7 --- /dev/null +++ b/docs/rust/index.rst @@ -0,0 +1,155 @@ +.. 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 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/scan/src/main.rs + :language: rust + :lines: 12-34 + :dedent: 4 + +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/scan/helper/src/lib.rs + :language: rust + :lines: 20-38 + :dedent: 4 + +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 +------------------- + +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 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`, :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. +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. +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. 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, Radix Rank and Merge Sort preserve their +inputs. +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/scan + 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 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 +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 + scan + reduce + reduce_batched + exchange + shuffle + adjacent_difference + discontinuity + histogram + 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 +``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/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/docs/rust/payload.rst b/docs/rust/payload.rst new file mode 100644 index 000000000000..df414908189b --- /dev/null +++ b/docs/rust/payload.rst @@ -0,0 +1,81 @@ +.. 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 ``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 +------------------------------ + +``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..46b8f60ec784 --- /dev/null +++ b/docs/rust/provider.rst @@ -0,0 +1,172 @@ +.. 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. + +``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 +------------------ + +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/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/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/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/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/docs/rust/shuffle.rst b/docs/rust/shuffle.rst new file mode 100644 index 000000000000..a0325fb9622d --- /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 +``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 +``rust/``; ``shuffle``, ``ShuffleOptions`` and ``ShuffleMode`` document the +full signatures and selectors. diff --git a/rust/.cargo/config.toml b/rust/.cargo/config.toml new file mode 100644 index 000000000000..d59034c2bf7e --- /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-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/.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..b4f1a410d583 --- /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=732bec8e2166c33b9b851860363c81baa475203e#732bec8e2166c33b9b851860363c81baa475203e" +dependencies = [ + "cuda-macros", +] + +[[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 = "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=732bec8e2166c33b9b851860363c81baa475203e#732bec8e2166c33b9b851860363c81baa475203e" +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=732bec8e2166c33b9b851860363c81baa475203e#732bec8e2166c33b9b851860363c81baa475203e" + +[[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/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/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/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/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/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/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 new file mode 100644 index 000000000000..a92b7241aa8a --- /dev/null +++ b/rust/cuda-coop-codegen/src/lib.rs @@ -0,0 +1,1260 @@ +// 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 Load/Store, Scan, Reduce, batched +//! Reduce, Exchange, Shuffle, Adjacent Difference, Discontinuity, Histogram +//! 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. +//! +//! # 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. +//! +//! [`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. +//! [`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. +//! [`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. +//! [`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. +//! [`HistogramRequest`] counts bounded integer samples in a complete 1D Block, +//! with independent counter type/output extent and initialized zero padding. +//! [`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. [`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 +//! [`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 adjacent_difference; +mod discontinuity; +mod exchange; +mod headers; +mod histogram; +mod merge_sort; +mod radix_sort; +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 histogram::{HistogramAlgorithm, HistogramRecipe, HistogramRequest}; +pub use merge_sort::{ + MergeSortKeysRecipe, MergeSortKeysRequest, MergeSortPairsRecipe, MergeSortPairsRequest, +}; +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}; +pub use shuffle::{ShuffleRecipe, ShuffleRequest}; +#[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 + } +} + +/// 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), + /// 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), + /// 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), + /// 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), + /// 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), + /// 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 { + 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(), + Self::Exchange(request) => request.block_shape(), + Self::Shuffle(request) => request.block_shape(), + 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(), + Self::MergeSortKeys(request) => request.block_shape(), + Self::MergeSortPairs(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), + Self::Exchange(request) => request.source(symbol), + Self::Shuffle(request) => request.source(symbol), + 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), + Self::MergeSortKeys(request) => request.source(symbol), + Self::MergeSortPairs(request) => request.source(symbol), + } + } + + fn token(self) -> String { + match self { + 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(), + 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(), + 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(), + Self::MergeSortKeys(request) => request.recipe.token().into(), + Self::MergeSortPairs(request) => request.recipe.token(request.value_dtype), + } + } + + fn scratch_free(self) -> bool { + match self { + Self::LoadStore(request) => request.algorithm.scratch_free(), + Self::Scan(_) + | Self::Reduce(_) + | Self::ReduceBatched(_) + | Self::Exchange(_) + | Self::Shuffle(_) + | Self::AdjacentDifference(_) + | Self::Discontinuity(_) + | Self::Histogram(_) + | Self::RadixSortKeys(_) + | Self::RadixSortPairs(_) + | Self::RadixRankKeys(_) + | Self::MergeSortKeys(_) + | Self::MergeSortPairs(_) => false, + } + } + + fn mutates_payload(self) -> bool { + match self { + Self::LoadStore(request) => { + !request.operation.is_load() && !request.algorithm.scratch_free() + } + // These families write fresh outputs, never rearranging inputs. + Self::Scan(_) + | Self::Reduce(_) + | Self::ReduceBatched(_) + | Self::Exchange(_) + | Self::Shuffle(_) + | Self::AdjacentDifference(_) + | Self::Discontinuity(_) + | Self::Histogram(_) + | Self::RadixSortKeys(_) + | Self::RadixSortPairs(_) + | Self::RadixRankKeys(_) + | Self::MergeSortKeys(_) + | Self::MergeSortPairs(_) => false, + } + } +} + +/// 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 +/// +/// The memory-tile and payload-layout requirements below describe Load/Store. +/// Scan, Reduce, batched Reduce, Exchange, Shuffle, Adjacent Difference, +/// 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`], [`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. +/// [`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. +/// 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 { + let requests: Vec<_> = requests + .iter() + .copied() + .map(NativeRequest::LoadStore) + .collect(); + self.compile_native(&requests, target) + } + + /// 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/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. + /// + /// # 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 + { + 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] + .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")); + } + } + 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"), + include_bytes!("scan.rs"), + include_bytes!("reduce.rs"), + include_bytes!("reduce_batched.rs"), + include_bytes!("exchange.rs"), + include_bytes!("shuffle.rs"), + include_bytes!("adjacent_difference.rs"), + include_bytes!("discontinuity.rs"), + include_bytes!("histogram.rs"), + ] { + hash_field(&mut inputs, bytes); + } + for option in &options { + hash_field(&mut inputs, option.as_bytes()); + } + let mut batch_identity = inputs.clone(); + let mut source = String::from(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 (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)) + } + (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"), + }; + 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); + let suffix: String = identity.iter().map(|byte| format!("{byte:02x}")).collect(); + 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)); + 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.scratch_free(), + mutates_payload: request.mutates_payload(), + }); + } + 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/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/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..cdae849cc8b0 --- /dev/null +++ b/rust/cuda-coop-codegen/src/oxide.rs @@ -0,0 +1,2986 @@ +// Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +//! 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 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 +//! 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`], and partial-output +//! 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. [`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. +//! [`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. 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, MergeSortKeysRecipe, MergeSortKeysRequest, + MergeSortPairsRecipe, MergeSortPairsRequest, 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}; +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, +} + +/// 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.dtype != Scalar::Signed(32) || self.warp_width.is_some() { + return Err(Error::message( + "callable recipes require i32 Block payloads", + )); + } + 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 recipes require a C identifier and the (const environment*, i32, i32) -> i32 callback ABI", + )); + } + 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", + )), + } + } +} + +#[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 { + 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() { + 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() { + 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() { + 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() { + 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 + .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()?; + 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()?; + 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, + "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).map(NativeRequest::LoadStore), + None => Ok(NativeRequest::LoadStore(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: 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)), + 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)), + 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![ + 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![ + 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![ + 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![ + 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![ + 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 = + 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-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()); + } + _ => {} + } + 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}"))) +} + +/// 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}"))) +} + +/// Register full-input i32 Block Reduce with a compiler-admitted Rust callable. +/// +/// The compiler supplies the concrete thunk and environment through protocol +/// version two. Input validity and root-only output use the payload ABI. +/// +/// # Errors +/// Rejects a nonabsolute executable or a path that cannot be encoded as JSON. +pub fn reduce_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: "reduce_callable_payload_validity_bits_v1", + algorithm: "raking", + }) + .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 +/// 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}"))) +} + +/// Encode a Warp batched-reduction recipe for the existing provider executable. +/// +/// Dtype, exact Block geometry, Warp width and batch count arrive with +/// specialization. Each lane's output capacity is derived as ceil(batches/width). +/// There are four typed payload pointers and 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 reduce_batched_registration( + executable: &Path, + recipe: ReduceBatchedRecipe, +) -> 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 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}"))) +} + +/// 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, +/// 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. +/// 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}"))) +} + +/// 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 { + 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, 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) + .read_to_end(&mut bytes)?; + 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 { + 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_native(&requests, &request.target)?; + 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( + "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, + // 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()?, + scratch_free: entry.scratch_free, + mutates_payload: entry.mutates_payload, + }) + }) + .collect::>>()?; + let header = serde_json::to_vec(&BatchResponse { + 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(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)?; + output.write_all(&artifact.ltoir)?; + Ok(()) +} + +#[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(()) + } + + 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 [ + 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 [ + (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::*; + 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] { + 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 [ + ( + 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() { + 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() { + 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() { + // 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(NativeRequest::LoadStore(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 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()); + } + } + + #[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/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/src/reduce.rs b/rust/cuda-coop-codegen/src/reduce.rs new file mode 100644 index 000000000000..1b81bb3c7b6d --- /dev/null +++ b/rust/cuda-coop-codegen/src/reduce.rs @@ -0,0 +1,395 @@ +// 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 { + pub(crate) 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", + } + } + + pub(crate) 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")), + } + } + + pub(crate) 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 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 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/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 new file mode 100644 index 000000000000..a4abfdcd687f --- /dev/null +++ b/rust/cuda-coop-codegen/src/scan.rs @@ -0,0 +1,355 @@ +// 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 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 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/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/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..68fefa9fcd62 --- /dev/null +++ b/rust/cuda-coop-codegen/tests/native.rs @@ -0,0 +1,2189 @@ +// 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::{ + AdjacentDifferenceRecipe, AdjacentDifferenceRequest, Algorithm, DiscontinuityRecipe, + DiscontinuityRequest, Dtype, ExchangeMode, ExchangeRecipe, ExchangeRequest, HistogramAlgorithm, + HistogramRecipe, HistogramRequest, Layout, MergeSortKeysRecipe, MergeSortKeysRequest, + MergeSortPairsRecipe, MergeSortPairsRequest, 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_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> { + 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> { + 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> { + 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> { + 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> { + 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_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> { + 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> { + 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> { + 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)] { + 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..8580098be28d --- /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 = "732bec8e2166c33b9b851860363c81baa475203e" } + +[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..21e2fc33e68c --- /dev/null +++ b/rust/cuda-coop/build.rs @@ -0,0 +1,314 @@ +// 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 cooperative operations. +use cuda_coop_codegen::{ + AdjacentDifferenceRecipe, Algorithm, DiscontinuityRecipe, Dtype, ExchangeMode, ExchangeRecipe, + HistogramAlgorithm, MergeSortKeysRecipe, MergeSortPairsRecipe, Operation, RadixRankKeysRecipe, + RadixSortKeysRecipe, RadixSortPairsRecipe, ReduceAlgorithm, ReduceBatchedLayout, + ReduceBatchedRecipe, ReduceForm, ReduceOperator, ReduceRecipe, ScanAlgorithm, ScanOperation, + ScanOperator, ScanRecipe, ShuffleRecipe, 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-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}"); + 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"), + (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"), + ] { + 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"), + (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"); + 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"), + (ReduceOperator::Min, "MIN"), + (ReduceOperator::Max, "MAX"), + (ReduceOperator::BitAnd, "BIT_AND"), + (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"), + (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"), + ( + 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"), + (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/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/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/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/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..785c01edc7dc --- /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=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/load_store/Cargo.toml b/rust/cuda-coop/examples/load_store/Cargo.toml new file mode 100644 index 000000000000..4b4a0d8c327b --- /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 = "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/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/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/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/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/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/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/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..34ea363f2141 --- /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=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/scan/Cargo.toml b/rust/cuda-coop/examples/scan/Cargo.toml new file mode 100644 index 000000000000..f63c89373ef8 --- /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 = "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/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/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/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/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/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/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/group.rs b/rust/cuda-coop/src/group.rs new file mode 100644 index 000000000000..67dcb34dc0ec --- /dev/null +++ b/rust/cuda-coop/src/group.rs @@ -0,0 +1,587 @@ +// 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 PREFIX: bool = !E && N >= 17 && N <= 31; + 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 if Self::PREFIX { + library::warp_prefix_member_v1::() + } 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/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 new file mode 100644 index 000000000000..1d2763bc3315 --- /dev/null +++ b/rust/cuda-coop/src/lib.rs @@ -0,0 +1,156 @@ +// 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 `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. +//! [`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. +//! [`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. +//! [`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. +//! [`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. +//! [`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. +//! [`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. +//! [`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. + +#![no_std] +#![feature(generic_const_exprs)] +#![allow(incomplete_features)] + +mod adjacent_difference; +mod alignment; +mod discontinuity; +mod exchange; +mod group; +mod histogram; +mod load_store; +mod merge_sort; +mod radix_sort; +mod reduce; +mod reduce_batched; +mod scan; +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, +}; +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, +}; +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, + radix_sort_pairs, +}; +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, + 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}; + +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/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 new file mode 100644 index 000000000000..316633b80e96 --- /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; + +/// Select ascending ordering. +#[derive(Clone, Copy, Debug, Default)] +pub struct Ascending; + +/// Select descending ordering. +#[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/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..e1fae6677c02 --- /dev/null +++ b/rust/cuda-coop/src/reduce/boundary.rs @@ -0,0 +1,476 @@ +// 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"), +); + +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, + 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 Selection::::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/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/src/scan.rs b/rust/cuda-coop/src/scan.rs new file mode 100644 index 000000000000..22742cff9b1d --- /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 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 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`. +/// +/// [`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 algorithm; default for Scan and an explicit choice for Reduce." + ), + (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/src/shuffle.rs b/rust/cuda-coop/src/shuffle.rs new file mode 100644 index 000000000000..a63807d93e04 --- /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 `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. +/// +/// ```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/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/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 new file mode 100644 index 000000000000..f54a08a50103 --- /dev/null +++ b/rust/cuda-coop/tests/check_load_store.py @@ -0,0 +1,2188 @@ +# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. +# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +"""Compile public collective participation, result and resource-epoch contrasts. + +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 +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", + 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( + "--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") +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) +(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>", + ) + + +# 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, + ), +) + +# 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", + ), +) + +# 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", + ), +) + +# 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, + ), +) + +# 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, + ), +) + +# 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", + ), +) + +# 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", + ), +) + +# 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: + 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", + "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_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 == "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)", + "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/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(()) +} 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/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(()) +} 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..39bf1e33848f --- /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=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/load_store/Cargo.toml b/rust/cuda-coop/tests/load_store/Cargo.toml new file mode 100644 index 000000000000..52b9a20dfad2 --- /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 = "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/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/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(()) +} 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/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(()) +} 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/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(()) +} 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..0989d39428e1 --- /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=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/scan/Cargo.toml b/rust/cuda-coop/tests/scan/Cargo.toml new file mode 100644 index 000000000000..a4c015982eae --- /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 = "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/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(()) +} 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(()) +} 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"