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/index.rst b/docs/rust/index.rst new file mode 100644 index 000000000000..9f10b3ff44ef --- /dev/null +++ b/docs/rust/index.rst @@ -0,0 +1,136 @@ +.. 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` and :doc:`discontinuity`. +Scan returns fresh per-group prefixes; +Reduce returns a total only at each group's root. Batched reduction returns a fresh +payload whose extent and valid slots follow its Warp width and output layout. +Exchange converts between blocked and striped per-thread layouts. +Shuffle shifts blocked values by one item, leaving the exposed edge invalid. +Discontinuity returns head flags, tail flags or both as fresh ``i32`` payloads. +Scan, Reduce, batched reduction, Exchange, Shuffle, Adjacent Difference and +Discontinuity preserve their input. Additional primitive families are under +development. The crate is unpublished and its interface can change. + +Build and run +------------- + +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 + provider + +Generate the API reference with ``cargo doc --workspace --no-deps`` from +``rust/``; ``cuda_coop`` contains the public item documentation. The Rust +workspace is separate from CCCL's CMake build, and existing C++ users do not +need a Rust toolchain. diff --git a/docs/rust/load_store.rst b/docs/rust/load_store.rst new file mode 100644 index 000000000000..32e7b88e9c34 --- /dev/null +++ b/docs/rust/load_store.rst @@ -0,0 +1,365 @@ +.. Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. +.. SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +Block and Warp Load/Store +========================= + +``cuda_coop::load`` fills an existing ``ThreadData`` and returns unit. +``cuda_coop::store`` writes its values and returns unit. Both take the group +first. Complete one-, two- and three-dimensional blocks support Direct, Striped, +Vectorize, Transpose, WarpTranspose and WarpTransposeTimesliced. The two warp +transpose algorithms require a total block thread count divisible by 32. +Striped Load gives each thread elements separated by the block size; the other +algorithms produce blocked payloads, where each thread owns consecutive elements. +Physical Warp and logical lane groups of widths 1, 2, 4, 8, 16 and 32 support +Direct, Striped, Vectorize and Transpose with implicit storage when the exact +block size is a multiple of 32. WarpTranspose and WarpTransposeTimesliced are +Block algorithms. + +Use ``use cuda_coop as coop;`` and keep the per-thread item count a const generic +named ``ITEMS_PER_THREAD``. The payload element type is inferred from its source. +The supported types are ``i8``, ``u8``, ``i16``, ``u16``, ``i32``, ``u32``, +``i64``, ``u64``, ``f32`` and ``f64``. + +A complete program +------------------ + +The example is an independent Cargo workspace with a separate generic helper +crate. It copies three block tiles with each of the six algorithms at item counts +1 and 4, including a short last tile. It also uses explicit automatic scratch in +a repeated helper call. The consumer has no build script, native shim, provider, +or scratch-size argument. +Its ``cuda-coop`` dependency supplies the native specialization during the normal +Oxide build. + +Use the Oxide revision selected by the manifests and ``1.99.1``. +The initial native configuration is Linux with CUDA 13.3 and a supported NVIDIA +GPU. From a complete CCCL checkout, with that Oxide toolchain installed: + +.. code-block:: shell + + cd rust/cuda-coop/examples/load_store + CUDA_HOME=/usr/local/cuda-13.3 CUDA_PATH=/usr/local/cuda-13.3 \ + cargo oxide run --arch sm_120 --materialize-cubin + +Choose the architecture for your GPU. Materialization produces the native device +artifact during the Cargo build. No separate compilation or linking step is +required. The exact block launch contract is checked when preparing the launch; +``auto_shared`` uses the compiler's inferred requirement, which is zero for +Direct, Striped and Vectorize Load/Store and includes the planned shared storage +for the three transposing algorithms. + +.. literalinclude:: ../../rust/cuda-coop/examples/load_store/src/main.rs + :language: rust + :start-after: // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +The reusable helper preserves the source element type. Its algorithm-generic +form borrows the payload mutably because the transposing Stores permute it. + +.. literalinclude:: ../../rust/cuda-coop/examples/load_store/helper/src/lib.rs + :language: rust + :start-after: use cuda_coop as coop; + +Counts, offsets, and payload validity +------------------------------------- + +``LoadOptions::default()`` and ``StoreOptions::default()`` select the full tile +at element offset zero. Named setters select a partial count, a typed Load +default, and an offset. The count describes the whole group's prefix, not a +per-thread count. It must fit the tile and ``i32::MAX``. Offsets are element +indices from the start of the allocation; supply each block's origin explicitly. +Counts, defaults, offsets, and their presence must be uniform across the +participating group; sibling groups may choose different values. + +For block shape ``[X,Y,Z]``, thread rank is x-fastest: +``t = x + X * (y + Y * z)``. A blocked payload slot ``i`` addresses +``offset + t * ITEMS_PER_THREAD + i``; a Striped slot addresses +``offset + t + X*Y*Z * i``. The exact shape is part of the kernel +launch contract; another shape with the same thread count is rejected. +Use ``LaunchConfig2D`` or ``LaunchConfig3D`` for the corresponding launch domain. +Declare the matching kernel contract, for example +``#[launch_contract(domain = 3, block = (8, 4, 4), dynamic_shared = inferred)]``. +The block origin in a multidimensional grid uses x-fastest block rank as well; +for a collective carrier, include that rank times its ownership grain +``block_size * CAPACITY`` in the element offset (see `Collective ownership and +completion`_). Read-only slice loads have no ownership grain. + +A masked Load without ``oob_default`` resets every out-of-prefix slot to +unwritten, even if initialized before the Load. ``get`` returns ``None`` for an +unwritten slot; indexing it traps. A default requires an explicit count and has +the source's numeric type. Filled Load initializes every slot, including the +missing tail. Partial Store writes only its selected prefix and requires only +those slots to be initialized. + +Direct, Striped and Vectorize Store accept ``&ThreadData``, a shared reborrow of ``&mut ThreadData``, +or a scalar numeric value. A scalar contributes one item per thread. All three preserve every payload value +and validity bit, including +initialized slots outside their write domain. + +Physical and logical Warp Load/Store +------------------------------------ + +Pass ``this_warp()`` through the same group-first API. No block dimensions or +extra payload type annotations are needed at the call site: + +.. code-block:: rust + + let warp = coop::this_warp(); + let mut values = coop::ThreadData::<_, ITEMS_PER_THREAD>::new(); + coop::load(warp, input, &mut values, + coop::LoadOptions::default().offset(block_origin).valid_items(count)); + coop::store(warp, output, &values, + coop::StoreOptions::default().offset(block_origin).valid_items(count)); + +Use ``this_warp().group_by::<8>()`` for an eight-lane group. Both exhaustive +and non-exhaustive forms support the six power-of-two widths; every lane is a +member for these widths. Let ``W`` be the selected width (32 for ``this_warp()``), +``t = x + X * (y + Y * z)``, ``q = t / W``, and ``lane = t % W``. +A blocked Warp slot ``i`` addresses +``offset + q * W * ITEMS_PER_THREAD + lane * ITEMS_PER_THREAD + i``; +a Striped slot addresses +``offset + q * W * ITEMS_PER_THREAD + lane + W * i``. +The operation adds this Warp origin; the caller still supplies the block +origin. The exact launch contract must contain only complete physical warps. +Direct, Striped and Vectorize use no shared scratch and accept shared payload +borrows for Store. Striped selects the striped mapping above. Full Vectorize +can use wide accesses when the actual offset-adjusted pointer is aligned; it +falls back to Direct otherwise. Counted Vectorize calls use Direct accesses. +Transpose uses disjoint shared storage for each group. The compiler reserves +one padded instance per group from the provider's actual layout; launch metadata +includes every instance. Its mutable Store permutes both values and validity. +Select it with ``.algorithm(coop::Transpose)`` on the same Load/Store options. +Explicit ``TempStorage`` is currently supported for Block calls only. + +Planned library calls cannot share a root with an unconstrained native dynamic +shared-memory view. This also applies to zero-scratch Block Direct calls with +an explicit numeric launch contract. Declare a fixed native view with +``DynamicSharedArray::get_fixed`` when composing it with planned storage. + +A collective carrier keeps its declared ``CAPACITY`` even when an operation +uses a different item count. Each group owns ``W * CAPACITY`` elements of its +block region. For example, with 64 threads, ``CAPACITY=4``, and one item per +thread, warp one owns ``[128,256)``. Caller offset zero gives the common +operation origin 32 and traps for a nonempty Store; caller offset 96 gives +origin 128 and is legal. Warp zero can keep offset zero. Supply the appropriate +block origin as well. The operation neither changes the offset nor repacks the +carrier to fit unrelated operations. Read-only slice loads have no carrier +grain restriction. With logical ``W=8``, ``CAPACITY=4`` and a scalar Store, +group one owns ``[32,64)``: offset zero gives origin 8 and traps, while offset 24 +gives origin 32. Carrier ownership does not change with an operation's item count. + +A group descriptor alone does not prove participation. Checked calls require +matching group-uniform controls and supported structured control flow. Ordinary +lane arithmetic is not inferred to be Warp-uniform merely because it divides +by a group width. Canonical ``group.rank_at::()`` for a logical +lane group and ``this_block().rank_at::()`` supply checked rank +queries. + +Conflicting accesses within a Warp tile can be completed by the matching +``cuda_device::cooperative_groups::WarpTile::sync()``. Construct it with +``this_thread_block().tiled_partition::()`` and import that module's +``ThreadGroup`` trait. Every lane in the synchronization tile must participate; +a narrower tile does not complete a wider access. A containing tile can +complete narrower accesses. Transitions involving Block accesses still need +``cuda_device::thread::sync_threads()``. Raw mask barriers are unsupported in checked +collective roots, and Warp synchronization does not release Block scratch. Direct adds no +implicit Warp barrier. + +Reusing dirty implicit Transpose storage inserts a matching Warp barrier before +the next call and completes conflicting accesses covered by that tile. First +use, scratch-free calls, and an operation in a different storage pool do not +provide this completion. A canonical containing barrier also completes narrower +implicit pools; a narrower barrier leaves wider and Block storage pending. +These barriers never end live explicit reservations or imply completion after +the final call. + +The public logical group width is ``usize`` and the canonical WarpTile width is +``u32``. An ordinary generic helper can take the group and tile as separate +arguments, or leave synchronization to its caller; no generic const-expression +feature or width cast is required. The compiler checks participation and whether +the chosen completion scope suffices for the accesses. + +Generic helpers can state ``G: coop::LoadStoreGroup`` alongside +``A: coop::Algorithm`` and ``S: coop::TempStorageLike``. Calls infer this sealed +bound from their group, algorithm and options; Block helpers retain their +existing algorithm and storage choices. + +Selecting an algorithm +---------------------- + +Select the algorithm through the trailing options. Selection is a Rust type, +so each call compiles one algorithm; ordinary defaults continue to select Direct. + +.. code-block:: rust + + coop::load(block, input, &mut values, + coop::LoadOptions::default().algorithm(coop::Transpose)); + coop::store(block, output, &mut values, + coop::StoreOptions::default().valid_items(count).algorithm(coop::Transpose)); + +Striped uses CUB's direct striped accesses without shared scratch. Select +``.algorithm(coop::Striped)`` on both operations to load and store a striped +payload. Its Store accepts a shared payload borrow and preserves the complete +payload, including a zero-count Store and initialized slots outside its prefix. + +.. code-block:: rust + + coop::load(block, input, &mut values, + coop::LoadOptions::default().algorithm(coop::Striped)); + coop::store(block, output, &values, + coop::StoreOptions::default().algorithm(coop::Striped)); + +Vectorize preserves blocked payload order and uses no shared scratch. Full +Load and Store can use wide accesses when the numeric type, item count and +actual offset-adjusted pointer permit them. CUB falls back to Direct access +when the pointer is not sufficiently aligned; ordinary scalar alignment is +enough to call the operation. A counted overload uses Direct access even when +its count equals the tile size. Filled Load is also counted. The example helper +uses full calls for complete tiles and counted calls for its final short tile. +Its Store accepts a shared payload borrow and preserves values and validity. + +Transpose uses CUB's coalesced accesses and shared exchange within the selected group. Its Store +requires a mutable payload borrow and permutes every value and initialized bit +from blocked into striped order. For a group of ``B`` threads and ``I`` items +per thread, destination lane/item ``(t, i)`` receives the old blocked slot at +linear index ``B*i+t``. This mutation occurs even for a zero-count Store; a +partial count only controls the global writes. Initialized slots outside the +written prefix remain initialized at their permuted positions. Scalar Store +uses a one-item temporary and accepts every algorithm. + +WarpTranspose and WarpTransposeTimesliced use the same blocked Load layout. +Their Store requires a mutable payload and permutes every slot within its +physical warp: ``(t, i)`` receives old blocked slot +``q = (t / 32) * (32 * I) + 32 * i + t % 32``. Both values and validity follow +that mapping, including zero-count and repeated Stores. The timesliced variant +reuses one warp's CUB exchange storage across successive warps; its counted +Store retains a separate validity plane for all warps. Both algorithms still +require the entire block to participate, and their internal warp barriers do +not complete a conflicting full-block memory access. For shape ``[X,Y,Z]``, +``X*Y*Z`` must be divisible by 32; X alone need not be. + +The compiler allocates the provider's complete shared storage, including the +validity plane when needed, and inserts completion before reusing managed +scratch. This does not add an unconditional trailing barrier. ``auto_shared`` +uses the inferred total for the kernel. + +Explicit temporary storage +-------------------------- + +Construct ``TempStorage`` inside the kernel and pass a mutable borrow through +the options. Default storage shares one slice between its uses and requires +an explicit full-block barrier before reuse. Omitting ``temp_storage`` requests +compiler-managed reuse synchronization instead. + +.. code-block:: rust + + let mut scratch = coop::TempStorage::new().auto_sync(); + coop::load(block, input, &mut values, + coop::LoadOptions::default().algorithm(coop::Transpose) + .temp_storage(&mut scratch)); + coop::store(block, output, &mut values, + coop::StoreOptions::default().algorithm(coop::Transpose) + .temp_storage(&mut scratch)); + +The consuming ``auto_sync()`` builder inserts a trailing block barrier after +each operation that actually uses scratch. Direct, Striped and Vectorize +Load/Store receive no barrier or allocation from this policy. Options carrying a mutable storage borrow are +consumed by the operation; construct fresh options with an ordinary reborrow +for each call, including inside loops. + +``exclusive()`` gives separate static primitive and reservation sites separate +slices. Repeated execution of one site still reuses its slice. Each static +constructor site denotes one descriptor, including in a loop; different +constructor sites receive independent storage. Independent scratch does not +complete conflicting accesses to the same global-memory carrier. + +``with_capacity::()`` specifies a positive capacity for the complete +descriptor plan, including exclusive slices and padding. An undersized request +is a compilation error. ``with_alignment::()`` selects a positive +power-of-two minimum alignment for the allocation base; each exclusive slice +also satisfies its own operation or reservation requirement. Omitted capacity +and alignment are inferred. Consuming builders retain the original descriptor +identity and must not change its policy after use. No destructor synchronizes. + +Typed reservations +------------------ + +Descriptors without auto-sync support ``reserve::()`` and +``reserve_aligned::()``. ``N`` is a positive element count, and the +aligned form requests a positive power-of-two minimum for that view. Only the +ten numeric types listed above are supported. Select the element type in the reservation, for example +``scratch.reserve::()``. + +Reservations return non-Copy tokens with ``as_mut_ptr()`` and consuming +``release()`` methods. They expose no initialized Rust value, reference, slice, +or safe indexing operation. The token does not hold a descriptor-long mutable +borrow, allowing simultaneous exclusive reservations. + +All element access is unsafe. The caller must establish bounds, initialization +before reads, aliasing, non-conflicting accesses, and phase synchronization. +Saved pointers must not be accessed after release, during overlapping reuse, +or after kernel exit. Complete asynchronous accesses before ending their view. +The compiler checks planned origins and live overlap; it does not prove arbitrary +unsafe pointer code. + +To reuse a shared region, release its reservation, then synchronize the full +block before the next overlapping operation. Release does not emit a barrier or +complete global-memory accesses. A barrier before release is insufficient, and +dropping or forgetting the token does not release it. Typed reservations cannot +be combined with automatic synchronization on their descriptor. + +Collective ownership and completion +----------------------------------- + +Writable kernel parameters use the canonical spelling +``DisjointSlice>``. The capacity defines each block's +ownership grain as ``block_size * CAPACITY`` elements; it is independent of the +operation's per-thread item count. The carrier keeps the buffer lifetime and +collective ownership without giving each lane an overlapping mutable slice. +The compiler's generated host launch consumes or borrows the actual device +buffers until the operation completes. + +For a grid of ``G`` blocks with ownership grain ``C``, block rank ``r`` owns +``[r*C, (r+1)*C)`` in each period of ``G*C`` elements. An offset may start inside +that region, and its nonempty selected range must stay inside the same region. +This also supports grid-stride traversal through later periods. + +The descriptor returned by ``this_block()`` is a query value. Safe collective +calls additionally require a supported Oxide build, an exact launch contract, +uniform control flow, and matching group participation. Compilation diagnoses +lane-varying controls or conflicting resource use. Runtime checks trap invalid +bounds, ownership ranges, counts, or unwritten values selected for Store. +Failures are not transactional across lanes. + +Use ``cuda_device::thread::sync_threads()`` between conflicting accesses to a +carrier: Store followed by Load or Store, or Load followed by Store. Pass +``&mut carrier`` as the Load source. The compiler requires completion, including reuse in loops. A genuine planned +scratch-reuse barrier can complete earlier accesses; code should not rely on an +algorithm merely being named Transpose to provide a trailing barrier. Ending a Rust borrow or dropping a value is not +completion. The host example calls ``sync()`` before inspecting its output. + +Host development +---------------- + +Host payload tests and rustdoc do not require a GPU or CUDA installation: + +.. code-block:: shell + + cd rust + cargo test --workspace + cargo clippy --workspace --all-targets -- -D warnings + RUSTDOCFLAGS="-D warnings" cargo doc --workspace --no-deps + +The codegen crate is only a host build dependency. Native headers come from the +same complete CCCL checkout as that dependency; a Python installation or the +consumer's working directory does not select headers. Native builds use +Oxide's selected toolkit. Cargo tracks the provider and declared native inputs; +replacing them with unchanged or older timestamps may require an explicit +rebuild. Crate publication is not supported yet. + +The rustdoc entries for ``load``, ``store``, ``LoadOptions``, ``StoreOptions``, +``LoadSource``, ``StoreValue``, ``Algorithm`` and the six algorithm selectors +specify their full generic signatures and +supported argument forms. Native correctness and artifact checks are separate +from host compilation. Performance qualification for masked payload access is +ongoing; the current implementation makes no negligible-overhead claim. diff --git a/docs/rust/payload.rst b/docs/rust/payload.rst new file mode 100644 index 000000000000..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/reduce.rst b/docs/rust/reduce.rst new file mode 100644 index 000000000000..8eb3c6acc240 --- /dev/null +++ b/docs/rust/reduce.rst @@ -0,0 +1,131 @@ +.. 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 a separate generic helper, 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. + +.. literalinclude:: ../../rust/cuda-coop/examples/reduce/src/main.rs + :language: rust + :start-after: // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +The helper fills the missing input tail with an explicit zero before summing. +Every thread calls the reduction. Only afterwards does the kernel inspect the +result and obtain a block-output witness. + +.. 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. Custom operators are not yet +supported. A generic helper can carry ``Op: ReduceOperator`` and pass the +actual operator value through ``.binary_op(op)``. + +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..8861258fbb13 --- /dev/null +++ b/rust/Cargo.lock @@ -0,0 +1,284 @@ +# This file is automatically @generated by Cargo. +# It is not intended for manual editing. +version = 4 + +[[package]] +name = "block-buffer" +version = "0.10.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3078c7629b62d3f0439517fa394996acacc5cbc91c5a20d8c658e77abd503a71" +dependencies = [ + "generic-array", +] + +[[package]] +name = "cfg-if" +version = "1.0.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4e7648175b45a9a48536d676f68d918270699102aa8dab5496df06904c914600" + +[[package]] +name = "cpufeatures" +version = "0.2.17" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "59ed5838eebb26a2bb2e58f6d5b5316989ae9d08bab10e0e6d103e656d1b0280" +dependencies = [ + "libc", +] + +[[package]] +name = "crypto-common" +version = "0.1.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "78c8292055d1c1df0cce5d180393dc8cce0abec0a7102adb6c7b1eef6016d60a" +dependencies = [ + "generic-array", + "typenum", +] + +[[package]] +name = "cuda-coop" +version = "0.0.0" +dependencies = [ + "cuda-coop-codegen", + "cuda-device", +] + +[[package]] +name = "cuda-coop-codegen" +version = "0.0.0" +dependencies = [ + "libloading", + "serde", + "serde_json", + "sha2", +] + +[[package]] +name = "cuda-device" +version = "0.2.1" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "cuda-macros", +] + +[[package]] +name = "cuda-macros" +version = "0.2.1" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "oxide-tile-export", + "proc-macro2", + "quote", + "reserved-oxide-symbols", + "serde", + "serde_json", + "syn 2.0.119", +] + +[[package]] +name = "digest" +version = "0.10.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9ed9a281f7bc9b7576e61468ba615a66a5c8cfdff42420a70aa82701a3b1e292" +dependencies = [ + "block-buffer", + "crypto-common", +] + +[[package]] +name = "generic-array" +version = "0.14.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "85649ca51fd72272d7821adaf274ad91c288277713d9c18820d8499a7ff69e9a" +dependencies = [ + "typenum", + "version_check", +] + +[[package]] +name = "itoa" +version = "1.0.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8f42a60cbdf9a97f5d2305f08a87dc4e09308d1276d28c869c684d7777685682" + +[[package]] +name = "libc" +version = "0.2.191" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "65cd87c9a6d60536580b316ad985726104730a5151c99845f2490428ef91240a" + +[[package]] +name = "libloading" +version = "0.8.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d7c4b02199fee7c5d21a5ae7d8cfa79a6ef5bb2fc834d6e9058e89c825efdc55" +dependencies = [ + "cfg-if", + "windows-link", +] + +[[package]] +name = "memchr" +version = "2.8.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cf8baf1c55e62ffcace7a9f06f4bd9cd3f0c4beb022d3b367256b91b87513d98" + +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + +[[package]] +name = "proc-macro2" +version = "1.0.107" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "985e7ec9bb745e6ce6535b544d84d6cd6f7ad8bd711c398938ae983b91a766d9" +dependencies = [ + "unicode-ident", +] + +[[package]] +name = "quote" +version = "1.0.47" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1fbf4db142a473a8d80c26bbf18454ed458bf8d26c8219c331daecfdbd079001" +dependencies = [ + "proc-macro2", +] + +[[package]] +name = "reserved-oxide-symbols" +version = "0.2.1" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" + +[[package]] +name = "serde" +version = "1.0.229" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4148590afebada386688f18773da617792bf2ef03ffc1e4cbd2b1d45b023e0ba" +dependencies = [ + "serde_core", + "serde_derive", +] + +[[package]] +name = "serde_core" +version = "1.0.229" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "67dca2c9c51e58a4791a4b1ed58308b39c64224d349a935ab5039aa360942a48" +dependencies = [ + "serde_derive", +] + +[[package]] +name = "serde_derive" +version = "1.0.229" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e7a5d71263a5a7d47b41f6b3f06ba276f10cc18b0931f1799f710578e2309348" +dependencies = [ + "proc-macro2", + "quote", + "syn 3.0.9", +] + +[[package]] +name = "serde_json" +version = "1.0.154" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e7e9cc8b1b85264074fbcc02a88680c4096b1e47df8f739dceb03bf482f04bd6" +dependencies = [ + "itoa", + "memchr", + "serde", + "serde_core", + "zmij", +] + +[[package]] +name = "sha2" +version = "0.10.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a7507d819769d01a365ab707794a4084392c824f54a7a6a7862f8c3d0892b283" +dependencies = [ + "cfg-if", + "cpufeatures", + "digest", +] + +[[package]] +name = "syn" +version = "2.0.119" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "872831b642d1a07999a962a351ed35b955ea2cfc8f3862091e2a240a84f17297" +dependencies = [ + "proc-macro2", + "quote", + "unicode-ident", +] + +[[package]] +name = "syn" +version = "3.0.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d78c8dee4c7bf0e14673097256fed6142ce9d3b85a408189d07482442145823b" +dependencies = [ + "proc-macro2", + "quote", + "unicode-ident", +] + +[[package]] +name = "thiserror" +version = "2.0.21" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "09e52cb86a36cede5cb101bf8908837b3e4c6e5e59fe7fd85c23fb56200d189e" +dependencies = [ + "thiserror-impl", +] + +[[package]] +name = "thiserror-impl" +version = "2.0.21" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fe5197923287db20a58125f0bc85c062f7f2c892de97b18c356f9efb14b28524" +dependencies = [ + "proc-macro2", + "quote", + "syn 3.0.9", +] + +[[package]] +name = "typenum" +version = "1.20.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b6f5e870be6c3b371b77fe0ee0bafb859fa4964b4404c27de1d380043c4dda20" + +[[package]] +name = "unicode-ident" +version = "1.0.27" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a2c754d6c33795a1c324727428e5a7dedb5b06195f9890bdbcba760d3e246563" + +[[package]] +name = "version_check" +version = "0.9.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0b928f33d975fc6ad9f86c8f283853ad26bdd5b10b7f1542aa2fa15e2289105a" + +[[package]] +name = "windows-link" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f0805222e57f7521d6a62e36fa9163bc891acd422f971defe97d64e70d0a4fe5" + +[[package]] +name = "zmij" +version = "1.0.23" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "29666d0abbfad1e3dc4dcf6144730dd3a3ab225bbbdac83319345b1b44ccfc1b" diff --git a/rust/Cargo.toml b/rust/Cargo.toml new file mode 100644 index 000000000000..c6fd608c3716 --- /dev/null +++ b/rust/Cargo.toml @@ -0,0 +1,18 @@ +# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. +# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +[workspace] +members = ["cuda-coop", "cuda-coop-codegen"] +resolver = "3" + +[workspace.package] +edition = "2024" +license = "Apache-2.0 WITH LLVM-exception" +rust-version = "1.85" + +[workspace.lints.rust] +missing_docs = "deny" +unsafe_op_in_unsafe_fn = "deny" + +[workspace.lints.rustdoc] +broken_intra_doc_links = "deny" diff --git a/rust/cuda-coop-codegen/Cargo.toml b/rust/cuda-coop-codegen/Cargo.toml new file mode 100644 index 000000000000..7a92fcaaded4 --- /dev/null +++ b/rust/cuda-coop-codegen/Cargo.toml @@ -0,0 +1,19 @@ +# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. +# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +[package] +name = "cuda-coop-codegen" +version = "0.0.0" +edition.workspace = true +license.workspace = true +rust-version.workspace = true +publish = false + +[dependencies] +libloading = "0.8" +serde = { version = "1", features = ["derive"] } +serde_json = "1" +sha2 = "0.10" + +[lints] +workspace = true diff --git a/rust/cuda-coop-codegen/examples/specialize.rs b/rust/cuda-coop-codegen/examples/specialize.rs new file mode 100644 index 000000000000..3a2cfe26364b --- /dev/null +++ b/rust/cuda-coop-codegen/examples/specialize.rs @@ -0,0 +1,35 @@ +// Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +//! Compile a host-side provider without a CUDA context or a device consumer. + +use cuda_coop_codegen::{Algorithm, Operation, Provider, Request}; +use std::{env, path::Path, time::Instant}; + +fn main() -> Result<(), Box> { + let arguments: Vec<_> = env::args().collect(); + if arguments.len() != 4 { + return Err("usage: specialize CUDA_TOOLKIT COMPUTE_TARGET OUTPUT_LTOIR".into()); + } + // This example is built inside the selected source checkout. Consumers + // supply their own explicit source root instead of relying on their cwd. + let checkout = Path::new(env!("CARGO_MANIFEST_DIR")).join("../.."); + let started = Instant::now(); + let provider = Provider::from_checkout(checkout, &arguments[1])?; + let initialized = started.elapsed(); + let request = Request::new(Operation::LoadFilled, Algorithm::Transpose, 128, 4)?; + let started = Instant::now(); + let artifact = provider.compile(&[request], &arguments[2])?; + let compiled = started.elapsed(); + std::fs::write(&arguments[3], &artifact.ltoir)?; + println!("symbol: {}", artifact.specializations[0].symbol); + println!( + "CUB layout: {:?}; scratch: {:?}; scratch free: {}", + artifact.specializations[0].cub_layout, + artifact.specializations[0].scratch, + artifact.specializations[0].scratch_free + ); + println!("initialization: {initialized:?}; compilation: {compiled:?}"); + print!("{}", artifact.compiler_log); + Ok(()) +} diff --git a/rust/cuda-coop-codegen/src/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/lib.rs b/rust/cuda-coop-codegen/src/lib.rs new file mode 100644 index 000000000000..3b255258c71f --- /dev/null +++ b/rust/cuda-coop-codegen/src/lib.rs @@ -0,0 +1,1156 @@ +// 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 and Discontinuity with NVRTC. +//! It returns linkable LTO-IR and the exact C++ scratch layout from the +//! **same compilation**, without a CUDA context, GPU or Python dependency. +//! It is a compiler building block, not a device API or a package release. +//! +//! # 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. +//! [`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 reduce; +mod reduce_batched; +mod scan; +mod shuffle; +pub use adjacent_difference::{AdjacentDifferenceRecipe, AdjacentDifferenceRequest}; +pub use discontinuity::{DiscontinuityRecipe, DiscontinuityRequest}; +pub use exchange::{ExchangeMode, ExchangeRecipe, ExchangeRequest}; +pub use reduce::{ReduceAlgorithm, ReduceForm, ReduceOperator, ReduceRecipe, ReduceRequest}; +pub use reduce_batched::{ReduceBatchedLayout, ReduceBatchedRecipe, ReduceBatchedRequest}; +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), +} + +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(), + } + } + + 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), + } + } + + 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(), + } + } + + 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(_) => 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(_) => 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 and +/// Discontinuity use lane-private input and fresh output. Their ABI and validity contracts are +/// documented by [`ScanRequest`], [`ReduceRequest`], [`ReduceBatchedRequest`], +/// [`ExchangeRequest`], [`ShuffleRequest`], [`AdjacentDifferenceRequest`] and +/// [`DiscontinuityRequest`]. +/// +/// Every thread in the selected group executes the provider in converged control +/// flow. The group is the exact block, or one aligned Warp of the selected width. +/// 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 + /// scratch describes one group and includes the actual CUB TempStorage, even + /// when its C++ empty type occupies one byte. The caller plans replication. + /// + /// # Errors + /// Rejects empty/duplicate requests and malformed targets as [`Self::compile`] + /// does; native compilation and layout decoding errors retain their diagnostics. + pub fn compile_native(&self, requests: &[NativeRequest], target: &str) -> Result { + if !target.starts_with("compute_") + || !target[8..].chars().all(|c| c.is_ascii_alphanumeric()) + || target.len() <= 8 + { + return Err(Error::message( + "target must be an explicit NVRTC architecture such as compute_120", + )); + } + if requests.is_empty() { + return Err(Error::message("at least one specialization is required")); + } + for (index, request) in requests.iter().enumerate() { + if requests[..index].contains(request) { + return Err(Error::message("duplicate specialization in one batch")); + } + } + let options = vec![ + "--std=c++17".to_owned(), + format!("--gpu-architecture={target}"), + "--device-c".to_owned(), + "-dlto".to_owned(), + "--no-source-include".to_owned(), + ]; + let mut inputs = Sha256::new(); + for bytes in [ + b"cuda-coop-codegen-v1".as_slice(), + &self.headers.digest, + &self.compiler.identity, + include_bytes!("lib.rs"), + include_bytes!("headers.rs"), + include_bytes!("nvrtc.rs"), + include_bytes!("block_io.cuh"), + 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"), + ] { + hash_field(&mut inputs, bytes); + } + for option in &options { + hash_field(&mut inputs, option.as_bytes()); + } + let mut batch_identity = inputs.clone(); + let mut source = String::from(include_str!("block_io.cuh")); + source.push_str("\ntemplate \nextern __device__ unsigned char layout_probe;\n"); + let mut expressions = Vec::with_capacity(requests.len() * 2); + let mut specializations = Vec::with_capacity(requests.len()); + for &request in requests { + let mut hash = inputs.clone(); + for axis in request.block_shape() { + hash_field(&mut hash, &axis.to_le_bytes()); + } + hash_field( + &mut hash, + request.source("cuda_coop_specialization").as_bytes(), + ); + let identity: [u8; 32] = hash.finalize().into(); + hash_field(&mut batch_identity, &identity); + let suffix: String = identity.iter().map(|byte| format!("{byte:02x}")).collect(); + let symbol = format!("cuda_coop_{}_{}", request.token(), suffix); + source.push_str(&request.source(&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/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..a4c9db979998 --- /dev/null +++ b/rust/cuda-coop-codegen/src/oxide.rs @@ -0,0 +1,1722 @@ +// Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +//! Version-one host transport for the Oxide library-operation marker. +//! +//! A dependency build script can serve as its already-built provider executable: +//! normal Cargo execution embeds [`registration`] in a device crate constant; +//! `--oxide-library-request-v1` execution passes stdin/stdout to [`serve`]; +//! `--oxide-library-inputs-v1` passes them to [`serve_inputs`] during macro expansion. +//! Use `cuda_device::library::registration!` in the getter to track these inputs. +//! The compiler reads the registration getter's constant MIR and invokes this +//! host process. This module never invokes Cargo or executes device Rust code. +//! The protocol supports raw i32 and typed validity 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. +//! Device calls retain all participation, validity and scratch obligations of +//! [`crate::Specialization`]. No allocation or synchronization policy is supplied. + +use crate::{ + AdjacentDifferenceRecipe, AdjacentDifferenceRequest, Algorithm, Artifact, DiscontinuityRecipe, + DiscontinuityRequest, Dtype, Error, ExchangeRecipe, ExchangeRequest, Layout, NativeRequest, + Operation, Provider, ReduceBatchedRecipe, ReduceBatchedRequest, ReduceForm, 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, +} + +#[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("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::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-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}"))) +} + +/// 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 a Block Adjacent Difference recipe for the existing provider executable. +/// +/// Dtype, exact Block geometry and payload extent arrive with specialization. +/// Four typed payload pointers describe preserved input and fully initialized +/// fresh output with packed validity. Partial recipes append an `i32` count; +/// recipes with a tile predecessor or successor append that value in the exact +/// input dtype, after the count when present. Scratch is inserted by the +/// existing library-call ABI. See [`AdjacentDifferenceRequest`] for capacities, +/// boundary behavior and participation requirements. +/// +/// # Errors +/// Rejects a nonabsolute executable or a path that cannot be encoded as JSON. +pub fn adjacent_difference_registration( + executable: &Path, + recipe: AdjacentDifferenceRecipe, +) -> Result { + if !executable.is_absolute() { + return Err(Error::message("provider executable path must be absolute")); + } + serde_json::to_string(&Registration { + version: VERSION, + provider: "cuda_coop", + executable, + operation: recipe.token(), + algorithm: "block", + }) + .map_err(|error| Error::message(format!("cannot encode provider registration: {error}"))) +} + +/// Encode a Block Discontinuity recipe for the existing provider executable. +/// +/// Input dtype, exact Block geometry and payload extent arrive with +/// specialization. Single-result recipes use four typed payload pointers; +/// paired recipes use six, ordered input, heads, tails, input validity, heads +/// validity, tails validity. Outputs are i32 regardless of input dtype. +/// Optional predecessor and successor controls follow, in that order, with +/// exactly the input type. Scratch is inserted at native parameter two by the +/// existing ABI. See [`DiscontinuityRequest`] for initialization and reuse. +/// +/// # Errors +/// Rejects a nonabsolute executable or a path that cannot be encoded as JSON. +pub fn discontinuity_registration( + executable: &Path, + recipe: DiscontinuityRecipe, +) -> Result { + if !executable.is_absolute() { + return Err(Error::message("provider executable path must be absolute")); + } + serde_json::to_string(&Registration { + version: VERSION, + provider: "cuda_coop", + executable, + operation: recipe.token(), + algorithm: "block", + }) + .map_err(|error| Error::message(format!("cannot encode provider registration: {error}"))) +} + +#[derive(Deserialize)] +#[serde(deny_unknown_fields)] +struct InputRequest { + version: u32, + cuda_root: PathBuf, + target: String, +} + +#[derive(Serialize)] +struct InputResponse { + version: u32, + paths: Vec, + environment: Vec, +} + +/// Declare native inputs for Oxide's registration macro without loading CUDA. +/// +/// `checkout` selects the same CCCL tree as [`serve`]. The bounded version-one +/// JSON request contains `cuda_root` and `target`; the response lists header +/// directories and NVRTC libraries. The macro tracks files and directories +/// recursively before rustc emits dep-info, including new header files. +/// No GPU, native compiler invocation or library loading occurs here. +/// Cargo's timestamp freshness rules apply; replacements with unchanged or older +/// timestamps can remain undetected. Rebuild after restoring or replacing inputs +/// this way, including a symlink switched to an older installation. +/// +/// # Errors +/// +/// Rejects malformed, oversized or unsupported requests, nonabsolute toolkit +/// paths, invalid targets, and inaccessible checkout paths. +pub fn serve_inputs(checkout: &Path, input: impl Read, mut output: impl Write) -> Result<()> { + let mut bytes = Vec::new(); + input + .take((MAX_HEADER + 1) as u64) + .read_to_end(&mut bytes)?; + if bytes.len() > MAX_HEADER { + return Err(Error::message("provider input request exceeds 1 MiB")); + } + let request: InputRequest = serde_json::from_slice(&bytes) + .map_err(|error| Error::message(format!("invalid provider input request: {error}")))?; + if request.version != VERSION + || !request.cuda_root.is_absolute() + || !request.target.starts_with("compute_") + || request.target.len() <= 8 + || !request.target[8..] + .chars() + .all(|character| character.is_ascii_alphanumeric()) + { + return Err(Error::message( + "provider inputs require version 1, an absolute toolkit path and compute_XX target", + )); + } + let checkout = checkout + .canonicalize() + .map_err(|error| Error::input(checkout, error))?; + let paths = crate::headers::Headers::roots(&checkout, &request.cuda_root) + .into_iter() + .map(|(path, _)| path) + .chain(crate::nvrtc::Compiler::input_paths(&request.cuda_root)) + .collect(); + let response = serde_json::to_vec(&InputResponse { + version: VERSION, + paths, + environment: Vec::new(), + }) + .map_err(|error| Error::message(format!("cannot encode provider inputs: {error}")))?; + if response.len() > MAX_HEADER { + return Err(Error::message("provider input response exceeds 1 MiB")); + } + output.write_all(&response)?; + Ok(()) +} + +/// Compile one versioned stdin request and write its framed native response. +/// +/// `checkout` selects the CCCL checkout whose headers the provider captures. +/// The request supplies an explicit toolkit path and target. Stdout contains +/// only the frame; the caller reports errors on stderr and exits nonzero. +/// The 24-byte prefix is `OXDLIB01`, JSON header length, and raw LTO length +/// (both lengths are little-endian u64), followed by the two byte buffers. +/// Requests and response headers are limited to 1 MiB; LTO is limited to 256 MiB. +/// Runtime arguments are a scalar/raw-pointer tuple; scratch is inserted at +/// parameter two and uses a generic byte pointer after shared address conversion. +/// +/// # Errors +/// +/// Rejects malformed/version-mismatched requests, unsupported types/recipes, +/// invalid geometry, and duplicate specializations. Native compiler errors +/// retain NVRTC diagnostics. No response is written until compilation succeeds. +pub fn serve(checkout: &Path, input: impl Read, mut output: impl Write) -> Result<()> { + let mut bytes = Vec::new(); + input + .take((MAX_HEADER + 1) as u64) + .read_to_end(&mut bytes)?; + if bytes.len() > MAX_HEADER { + return Err(Error::message("provider request exceeds 1 MiB")); + } + let request: BatchRequest = serde_json::from_slice(&bytes) + .map_err(|error| Error::message(format!("invalid provider request: {error}")))?; + if request.version != VERSION { + return Err(Error::message(format!( + "unsupported provider protocol {}; expected {VERSION}", + request.version + ))); + } + let requests: Vec<_> = request + .requests + .iter() + .map(SpecializationRequest::native) + .collect::>()?; + let artifact = Provider::from_checkout(checkout, &request.cuda_root)? + .compile_native(&requests, &request.target)?; + write_response(&requests, artifact, &mut output) +} + +fn write_response( + requests: &[NativeRequest], + artifact: Artifact, + output: &mut impl Write, +) -> Result<()> { + if artifact.ltoir.is_empty() || artifact.ltoir.len() > MAX_LTO { + return Err(Error::message( + "native LTO buffer must contain 1..=256 MiB bytes", + )); + } + let specializations = artifact + .specializations + .into_iter() + .zip(requests) + .enumerate() + .map(|(index, (entry, request))| { + Ok(SpecializationResponse { + request_index: u32::try_from(index) + .map_err(|_| Error::message("request index exceeds u32"))?, + identity: entry.identity, + symbol: entry.symbol, + arguments: arguments(*request), + cub_layout: entry.cub_layout.try_into()?, + scratch: entry.scratch.try_into()?, + scratch_free: entry.scratch_free, + mutates_payload: entry.mutates_payload, + }) + }) + .collect::>>()?; + let header = serde_json::to_vec(&BatchResponse { + version: VERSION, + batch_identity: artifact.identity, + specializations, + }) + .map_err(|error| Error::message(format!("cannot encode provider response: {error}")))?; + if header.len() > MAX_HEADER { + return Err(Error::message("provider response header exceeds 1 MiB")); + } + output.write_all(MAGIC)?; + output.write_all(&(header.len() as u64).to_le_bytes())?; + output.write_all(&(artifact.ltoir.len() as u64).to_le_bytes())?; + output.write_all(&header)?; + output.write_all(&artifact.ltoir)?; + Ok(()) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::{ + ReduceAlgorithm, ReduceBatchedLayout, ReduceOperator, ScanAlgorithm, ScanOperation, + ScanOperator, + }; + + #[test] + fn discontinuity_registration_preserves_mixed_output_abi() { + use DiscontinuityRecipe::*; + for (recipe, paired, neighbors) in [ + (Heads, false, 0), + (HeadsPredecessor, false, 1), + (Tails, false, 0), + (TailsSuccessor, false, 1), + (HeadsAndTails, true, 0), + (HeadsAndTailsPredecessor, true, 1), + (HeadsAndTailsSuccessor, true, 1), + (HeadsAndTailsBoth, true, 2), + ] { + let registration: serde_json::Value = serde_json::from_str( + &discontinuity_registration(Path::new("/provider"), recipe).unwrap(), + ) + .unwrap(); + let mut request = SpecializationRequest { + operation: registration["operation"].as_str().unwrap().into(), + algorithm: registration["algorithm"].as_str().unwrap().into(), + dtype: Scalar::Float(64), + block: [4, 3, 2], + items_per_thread: 4, + warp_width: None, + }; + let native = request.native().unwrap(); + assert_eq!( + native, + NativeRequest::Discontinuity( + DiscontinuityRequest::block(Dtype::F64, recipe, [4, 3, 2], 4).unwrap() + ) + ); + let mut expected = vec![ + Argument::ConstPointer(Scalar::Float(64)), + Argument::MutPointer(Scalar::Signed(32)), + ]; + if paired { + expected.push(Argument::MutPointer(Scalar::Signed(32))); + } + expected.extend([ + Argument::ConstPointer(Scalar::Unsigned(8)), + Argument::MutPointer(Scalar::Unsigned(8)), + ]); + if paired { + expected.push(Argument::MutPointer(Scalar::Unsigned(8))); + } + expected.extend((0..neighbors).map(|_| Argument::Value(Scalar::Float(64)))); + assert_eq!(arguments(native), expected); + assert!(!native.scratch_free() && !native.mutates_payload()); + assert_eq!(registration["version"], 1); + assert_eq!(registration["provider"], "cuda_coop"); + request.warp_width = Some(32); + assert!(request.native().is_err()); + } + assert!(discontinuity_registration(Path::new("relative"), Heads).is_err()); + assert!( + DiscontinuityRecipe::parse("discontinuity_heads_successor_validity_bits_v1", "block") + .is_err() + ); + assert!( + DiscontinuityRecipe::parse("discontinuity_tails_predecessor_validity_bits_v1", "block") + .is_err() + ); + assert!(DiscontinuityRecipe::parse(Heads.token(), "warp").is_err()); + assert!(DiscontinuityRecipe::parse(Heads.token(), "sort").is_err()); + } + + #[test] + fn adjacent_difference_registration_preserves_typed_controls() { + use AdjacentDifferenceRecipe::*; + for (recipe, controls) in [ + (LeftFull, vec![]), + (LeftFullPredecessor, vec![Scalar::Float(64)]), + (LeftPartial, vec![Scalar::Signed(32)]), + ( + LeftPartialPredecessor, + vec![Scalar::Signed(32), Scalar::Float(64)], + ), + (RightFull, vec![]), + (RightFullSuccessor, vec![Scalar::Float(64)]), + (RightPartial, vec![Scalar::Signed(32)]), + ] { + let registration: serde_json::Value = serde_json::from_str( + &adjacent_difference_registration(Path::new("/provider"), recipe).unwrap(), + ) + .unwrap(); + let mut request = SpecializationRequest { + operation: registration["operation"].as_str().unwrap().into(), + algorithm: registration["algorithm"].as_str().unwrap().into(), + dtype: Scalar::Float(64), + block: [4, 3, 2], + items_per_thread: 9, + warp_width: None, + }; + let native = request.native().unwrap(); + assert_eq!( + native, + NativeRequest::AdjacentDifference( + AdjacentDifferenceRequest::block(Dtype::F64, recipe, [4, 3, 2], 9).unwrap() + ) + ); + let mut expected = vec![ + Argument::ConstPointer(Scalar::Float(64)), + Argument::MutPointer(Scalar::Float(64)), + Argument::ConstPointer(Scalar::Unsigned(8)), + Argument::MutPointer(Scalar::Unsigned(8)), + ]; + expected.extend(controls.into_iter().map(Argument::Value)); + assert_eq!(arguments(native), expected); + assert!(!native.scratch_free() && !native.mutates_payload()); + assert_eq!(registration["version"], 1); + assert_eq!(registration["provider"], "cuda_coop"); + request.warp_width = Some(32); + assert!(request.native().is_err()); + } + assert!(adjacent_difference_registration(Path::new("relative"), LeftFull).is_err()); + assert!( + AdjacentDifferenceRecipe::parse( + "adjacent_difference_right_partial_successor_validity_bits_v1", + "block" + ) + .is_err() + ); + assert!(AdjacentDifferenceRecipe::parse(LeftFull.token(), "warp").is_err()); + assert!(AdjacentDifferenceRecipe::parse(RightFull.token(), "transpose").is_err()); + } + + #[test] + fn shuffle_registration_preserves_scope_extent_and_pointer_abi() { + for recipe in [ShuffleRecipe::Down, ShuffleRecipe::Up] { + 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/reduce.rs b/rust/cuda-coop-codegen/src/reduce.rs new file mode 100644 index 000000000000..37226c221974 --- /dev/null +++ b/rust/cuda-coop-codegen/src/reduce.rs @@ -0,0 +1,373 @@ +// 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 { + 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 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 call = 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 +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; + 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..7c2024768da8 --- /dev/null +++ b/rust/cuda-coop-codegen/tests/native.rs @@ -0,0 +1,1571 @@ +// 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, Layout, + NativeRequest, Operation, Provider, ReduceAlgorithm, ReduceBatchedLayout, ReduceBatchedRecipe, + ReduceBatchedRequest, ReduceForm, ReduceOperator, ReduceRequest, Request, ScanAlgorithm, + ScanOperation, ScanOperator, ScanRequest, ShuffleRecipe, ShuffleRequest, +}; +use std::{env, path::Path, time::Instant}; + +#[test] +#[ignore = "requires CUDA_HOME and a Linux CUDA toolkit with NVRTC LTO support"] +fn native_mixed_discontinuity_recipes() -> Result<(), Box> { + use DiscontinuityRecipe::*; + let toolkit = env::var("CUDA_HOME")?; + let target = env::var("CUDA_COOP_TEST_TARGET").unwrap_or_else(|_| "compute_120".into()); + let checkout = Path::new(env!("CARGO_MANIFEST_DIR")).join("../.."); + let provider = Provider::from_checkout(&checkout, &toolkit)?; + let mut requests = Vec::new(); + // Eight overloads and all input types; every result uses i32. N9 crosses + // the bitmap boundary, including independent heads/tails validity planes. + for (dtype, block, items, recipe) in [ + (Dtype::F64, [4, 3, 2], 4, HeadsAndTailsBoth), + (Dtype::I8, [1, 1, 1], 1, Heads), + (Dtype::U8, [1, 1, 1], 1, Tails), + (Dtype::I16, [33, 1, 1], 4, HeadsPredecessor), + (Dtype::U16, [4, 3, 2], 4, TailsSuccessor), + (Dtype::I32, [33, 1, 1], 9, HeadsAndTails), + (Dtype::U32, [64, 1, 1], 4, HeadsAndTailsPredecessor), + (Dtype::I64, [64, 1, 1], 1, HeadsAndTailsSuccessor), + (Dtype::U64, [4, 3, 2], 9, Heads), + (Dtype::F32, [128, 1, 1], 4, HeadsAndTailsBoth), + ] { + requests.push(NativeRequest::Discontinuity(DiscontinuityRequest::block( + dtype, recipe, block, items, + )?)); + } + requests.extend([ + NativeRequest::LoadStore(Request::new_typed( + Dtype::I32, + Operation::LoadFullValidity, + Algorithm::Direct, + 64, + 4, + )?), + NativeRequest::Scan(ScanRequest::warp( + Dtype::I32, + ScanOperation::ExclusiveSum, + [64, 1, 1], + 8, + )?), + NativeRequest::Reduce(ReduceRequest::warp( + Dtype::I32, + ReduceOperator::Sum, + ReduceForm::Payload, + [64, 1, 1], + 8, + 4, + )?), + NativeRequest::ReduceBatched(ReduceBatchedRequest::warp( + Dtype::I32, + ReduceBatchedRecipe::new(ReduceOperator::Sum, ReduceBatchedLayout::Striped), + [64, 1, 1], + 8, + 4, + )?), + NativeRequest::Exchange(ExchangeRequest::warp( + Dtype::I32, + ExchangeRecipe::Warp(ExchangeMode::StripedToBlocked), + [64, 1, 1], + 32, + 4, + )?), + NativeRequest::Shuffle(ShuffleRequest::block( + Dtype::I32, + ShuffleRecipe::Down, + [64, 1, 1], + 4, + )?), + NativeRequest::AdjacentDifference(AdjacentDifferenceRequest::block( + Dtype::I32, + AdjacentDifferenceRecipe::LeftPartialPredecessor, + [64, 1, 1], + 4, + )?), + ]); + let artifact = provider.compile_native(&requests, &target)?; + assert_eq!(artifact.specializations.len(), 17); + assert!(!artifact.ltoir.is_empty()); + let mut identities = std::collections::HashSet::new(); + let mut symbols = std::collections::HashSet::new(); + for (request, entry) in requests.iter().zip(&artifact.specializations) { + assert!(identities.insert(entry.identity)); + assert!(symbols.insert(&entry.symbol)); + if matches!(request, NativeRequest::LoadStore(_)) { + assert!(entry.scratch_free); + assert_eq!(entry.scratch.size, 0); + } else { + assert_eq!(entry.cub_layout, entry.scratch); + assert!(entry.scratch.size > 0 && entry.scratch.alignment.is_power_of_two()); + assert!(!entry.scratch_free && !entry.mutates_payload); + } + eprintln!( + "{request:?}: CUB {:?}, scratch {:?}", + entry.cub_layout, entry.scratch + ); + } + // Actual CUB TempStorage contains both endpoint planes, even for Heads. + assert_eq!( + artifact.specializations[0].scratch, + Layout { + size: 384, + alignment: 8 + } + ); + assert_eq!( + artifact.specializations[1].scratch, + Layout { + size: 2, + alignment: 1 + } + ); + assert!( + provider + .compile_native(&[requests[0], requests[0]], &target) + .unwrap_err() + .to_string() + .contains("duplicate") + ); + eprintln!( + "{} mixed Discontinuity/Adjacent Difference/Load/Scan/Reduce/batched Reduce/Exchange/Shuffle specializations, {} LTO bytes", + requests.len(), + artifact.ltoir.len() + ); + Ok(()) +} + +#[test] +#[ignore = "requires CUDA_HOME and a Linux CUDA toolkit with NVRTC LTO support"] +fn native_mixed_adjacent_difference_recipes() -> Result<(), Box> { + use AdjacentDifferenceRecipe::*; + let toolkit = env::var("CUDA_HOME")?; + let target = env::var("CUDA_COOP_TEST_TARGET").unwrap_or_else(|_| "compute_120".into()); + let checkout = Path::new(env!("CARGO_MANIFEST_DIR")).join("../.."); + let provider = Provider::from_checkout(&checkout, &toolkit)?; + let mut requests = Vec::new(); + // Distribute all types and seven overloads over small, odd and XYZ Blocks; + // N9 crosses the packed-validity byte boundary without a Cartesian matrix. + for (dtype, block, items, recipe) in [ + (Dtype::I8, [1, 1, 1], 1, LeftFull), + (Dtype::U8, [1, 1, 1], 1, RightFull), + (Dtype::I16, [33, 1, 1], 1, LeftFullPredecessor), + (Dtype::U16, [4, 3, 2], 4, RightFullSuccessor), + (Dtype::I32, [33, 1, 1], 9, LeftPartial), + (Dtype::U32, [64, 1, 1], 4, RightPartial), + (Dtype::I64, [4, 3, 2], 4, LeftPartialPredecessor), + (Dtype::U64, [64, 1, 1], 1, LeftFull), + (Dtype::F32, [128, 1, 1], 4, RightFull), + (Dtype::F64, [4, 3, 2], 9, LeftPartialPredecessor), + ] { + requests.push(NativeRequest::AdjacentDifference( + AdjacentDifferenceRequest::block(dtype, recipe, block, items)?, + )); + } + requests.extend([ + NativeRequest::LoadStore(Request::new_typed( + Dtype::I32, + Operation::LoadFullValidity, + Algorithm::Direct, + 64, + 4, + )?), + NativeRequest::Scan(ScanRequest::warp( + Dtype::I32, + ScanOperation::ExclusiveSum, + [64, 1, 1], + 8, + )?), + NativeRequest::Reduce(ReduceRequest::warp( + Dtype::I32, + ReduceOperator::Sum, + ReduceForm::Payload, + [64, 1, 1], + 8, + 4, + )?), + NativeRequest::ReduceBatched(ReduceBatchedRequest::warp( + Dtype::I32, + ReduceBatchedRecipe::new(ReduceOperator::Sum, ReduceBatchedLayout::Striped), + [64, 1, 1], + 8, + 4, + )?), + NativeRequest::Exchange(ExchangeRequest::warp( + Dtype::I32, + ExchangeRecipe::Warp(ExchangeMode::StripedToBlocked), + [64, 1, 1], + 32, + 4, + )?), + NativeRequest::Shuffle(ShuffleRequest::block( + Dtype::I32, + ShuffleRecipe::Down, + [64, 1, 1], + 4, + )?), + ]); + let artifact = provider.compile_native(&requests, &target)?; + assert_eq!(artifact.specializations.len(), 16); + assert!(!artifact.ltoir.is_empty()); + let mut identities = std::collections::HashSet::new(); + let mut symbols = std::collections::HashSet::new(); + for (request, entry) in requests.iter().zip(&artifact.specializations) { + assert!(identities.insert(entry.identity)); + assert!(symbols.insert(&entry.symbol)); + if matches!(request, NativeRequest::LoadStore(_)) { + assert!(entry.scratch_free); + assert_eq!(entry.scratch.size, 0); + } else { + assert_eq!(entry.cub_layout, entry.scratch); + assert!(entry.scratch.size > 0 && entry.scratch.alignment.is_power_of_two()); + assert!(!entry.scratch_free && !entry.mutates_payload); + } + eprintln!( + "{request:?}: CUB {:?}, scratch {:?}", + entry.cub_layout, entry.scratch + ); + } + // Both endpoint planes belong to CUB TempStorage, even for one direction. + assert_eq!( + artifact.specializations[0].scratch, + Layout { + size: 2, + alignment: 1 + } + ); + assert_eq!( + artifact.specializations[9].scratch, + Layout { + size: 384, + alignment: 8 + } + ); + assert!( + provider + .compile_native(&[requests[0], requests[0]], &target) + .unwrap_err() + .to_string() + .contains("duplicate") + ); + eprintln!( + "{} mixed Adjacent Difference/Load/Scan/Reduce/batched Reduce/Exchange/Shuffle specializations, {} LTO bytes", + requests.len(), + artifact.ltoir.len() + ); + Ok(()) +} + +#[test] +#[ignore = "requires CUDA_HOME and a Linux CUDA toolkit with NVRTC LTO support"] +fn native_mixed_shuffle_recipes() -> Result<(), Box> { + 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..45682aa419bd --- /dev/null +++ b/rust/cuda-coop/Cargo.toml @@ -0,0 +1,18 @@ +# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. +# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +[package] +name = "cuda-coop" +version = "0.0.0" +edition.workspace = true +license.workspace = true +publish = false + +[dependencies] +cuda-device = { git = "https://github.com/tpn/cuda-rust.git", rev = "96caed46c776eb49226137eac849e795b83b7ce0" } + +[build-dependencies] +cuda-coop-codegen = { path = "../cuda-coop-codegen" } + +[lints] +workspace = true diff --git a/rust/cuda-coop/build.rs b/rust/cuda-coop/build.rs new file mode 100644 index 000000000000..70ef835e9f9a --- /dev/null +++ b/rust/cuda-coop/build.rs @@ -0,0 +1,214 @@ +// 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, ExchangeMode, ExchangeRecipe, + Operation, 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-inputs-v1") + ) { + return; + } + let executable = std::env::current_exe().expect("cuda.coop provider executable"); + 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..be2013f83518 --- /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=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "libnvvm-sys", + "nvjitlink-sys", + "serde", + "sha2", + "thiserror 2.0.21", +] + +[[package]] +name = "cuda-async" +version = "0.5.0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "anyhow", + "cuda-bindings", + "cuda-core", + "dashmap", + "futures", + "half", + "once_cell", + "rustc-hash 2.1.3", + "thiserror 1.0.69", +] + +[[package]] +name = "cuda-bindings" +version = "0.5.0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "bindgen", + "libloading", + "prettyplease", + "proc-macro2", + "quote", + "syn 2.0.119", +] + +[[package]] +name = "cuda-coop" +version = "0.0.0" +dependencies = [ + "cuda-coop-codegen", + "cuda-device", +] + +[[package]] +name = "cuda-coop-codegen" +version = "0.0.0" +dependencies = [ + "libloading", + "serde", + "serde_json", + "sha2", +] + +[[package]] +name = "cuda-core" +version = "0.5.0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "anyhow", + "cuda-bindings", + "cuda-core-derive", + "half", + "oxide-artifacts", +] + +[[package]] +name = "cuda-core-derive" +version = "0.5.0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", +] + +[[package]] +name = "cuda-device" +version = "0.2.1" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "cuda-macros", +] + +[[package]] +name = "cuda-host" +version = "0.2.1" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "cuda-artifact-finalizer", + "cuda-async", + "cuda-core", + "cuda-device", + "cuda-macros", + "half", + "libc", + "oxide-artifacts", + "ptx-parse", + "sha2", + "thiserror 2.0.21", +] + +[[package]] +name = "cuda-macros" +version = "0.2.1" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "oxide-tile-export", + "proc-macro2", + "quote", + "reserved-oxide-symbols", + "serde", + "serde_json", + "syn 2.0.119", +] + +[[package]] +name = "cuda-target-spec" +version = "0.1.0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" + +[[package]] +name = "dashmap" +version = "6.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e6361d5c062261c78a176addb82d4c821ae42bed6089de0e12603cd25de2059c" +dependencies = [ + "cfg-if", + "crossbeam-utils", + "hashbrown 0.14.5", + "lock_api", + "once_cell", + "parking_lot_core", +] + +[[package]] +name = "digest" +version = "0.10.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9ed9a281f7bc9b7576e61468ba615a66a5c8cfdff42420a70aa82701a3b1e292" +dependencies = [ + "block-buffer", + "crypto-common", +] + +[[package]] +name = "either" +version = "1.19.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0e9c71c2167ca323c882b99918929403426e2373ea17242ff5653e0d5e1058be" + +[[package]] +name = "equivalent" +version = "1.0.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "00d174d5400e5e8fd687ad1049e2f578285fa914201b1af7e8b112a4546bd826" + +[[package]] +name = "errno" +version = "0.3.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "39cab71617ae0d63f51a36d69f866391735b51691dbda63cf6f96d042b63efeb" +dependencies = [ + "libc", + "windows-sys 0.61.2", +] + +[[package]] +name = "foldhash" +version = "0.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d9c4f5dac5e15c24eb999c26181a6ca40b39fe946cbe4c263c7209467bc83af2" + +[[package]] +name = "futures" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9a31d2a3fbaaeb2af2368bbdd904aa8e812d3c04a1ee10d3171f52d556e5d0a3" +dependencies = [ + "futures-channel", + "futures-core", + "futures-executor", + "futures-io", + "futures-sink", + "futures-task", + "futures-util", +] + +[[package]] +name = "futures-channel" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b1f9e3d69d39e4862ffed03ed071a76f9a13ba1d9109d355b0f0aa6b15e393c4" +dependencies = [ + "futures-core", + "futures-sink", +] + +[[package]] +name = "futures-core" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "92d699e522242e69e3003b94ecc1f960f3a5e015aa7c5d7486e65ad01dd94f5e" + +[[package]] +name = "futures-executor" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "031b47cf1a3c6cc8bc2fc76cd437f521619387907d469316e7c0bc278f1f5432" +dependencies = [ + "futures-core", + "futures-task", + "futures-util", +] + +[[package]] +name = "futures-io" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "53c0fa8157de1303bfffdaa1cc2a673bfffb60102f76b0ef4441659124373fed" + +[[package]] +name = "futures-macro" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9fb9654ba8355388abeb8dcb4fc62f511300867002afc858860463bdd9fe0c44" +dependencies = [ + "proc-macro2", + "quote", + "syn 3.0.9", +] + +[[package]] +name = "futures-sink" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1944426bf7d03f1d14f708785e4b33efd750b36d48a157b836b3efc15ede8e1d" + +[[package]] +name = "futures-task" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cd417de3d1d015fc3bfd2b1ea46dfc7bab72ef86f1cc7cc9c78e728b34a6d1fd" + +[[package]] +name = "futures-util" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0d50a92467f8ba5dd6e3ee5d4bd04d73ab2e4e1c44474a0674821dfce14b79bc" +dependencies = [ + "futures-channel", + "futures-core", + "futures-io", + "futures-macro", + "futures-sink", + "futures-task", + "memchr", + "pin-project-lite", + "slab", +] + +[[package]] +name = "generic-array" +version = "0.14.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "85649ca51fd72272d7821adaf274ad91c288277713d9c18820d8499a7ff69e9a" +dependencies = [ + "typenum", + "version_check", +] + +[[package]] +name = "glob" +version = "0.3.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e4eba85ea1d0a966a983acd07deee566e67395d2d96b6fb39e62b5a833f1eb0b" + +[[package]] +name = "half" +version = "2.7.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" +dependencies = [ + "cfg-if", + "crunchy", + "zerocopy", +] + +[[package]] +name = "hashbrown" +version = "0.14.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e5274423e17b7c9fc20b6e7e208532f9b19825d82dfd615708b70edd83df41f1" + +[[package]] +name = "hashbrown" +version = "0.15.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9229cfe53dfd69f0609a49f65461bd93001ea1ef889cd5529dd176593f5338a1" +dependencies = [ + "foldhash", +] + +[[package]] +name = "hashbrown" +version = "0.17.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ed5909b6e89a2db4456e54cd5f673791d7eca6732202bbf2a9cc504fe2f9b84a" + +[[package]] +name = "home" +version = "0.5.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cc627f471c528ff0c4a49e1d5e60450c8f6461dd6d10ba9dcd3a61d3dff7728d" +dependencies = [ + "windows-sys 0.61.2", +] + +[[package]] +name = "indexmap" +version = "2.14.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cc4e190f5d26ca7051642629da2c52fc03bde85a03197c99408dcd291734c855" +dependencies = [ + "equivalent", + "hashbrown 0.17.1", +] + +[[package]] +name = "itertools" +version = "0.12.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ba291022dbbd398a455acf126c1e341954079855bc60dfdda641363bd6922569" +dependencies = [ + "either", +] + +[[package]] +name = "itoa" +version = "1.0.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8f42a60cbdf9a97f5d2305f08a87dc4e09308d1276d28c869c684d7777685682" + +[[package]] +name = "lazy_static" +version = "1.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "20870f649af7073d53e38067b2a84312175d56ea15217e1b15bc83506ec50afb" + +[[package]] +name = "lazycell" +version = "1.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "830d08ce1d1d941e6b30645f1a0eb5643013d835ce3779a5fc208261dbe10f55" + +[[package]] +name = "libc" +version = "0.2.191" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "65cd87c9a6d60536580b316ad985726104730a5151c99845f2490428ef91240a" + +[[package]] +name = "libloading" +version = "0.8.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d7c4b02199fee7c5d21a5ae7d8cfa79a6ef5bb2fc834d6e9058e89c825efdc55" +dependencies = [ + "cfg-if", + "windows-link", +] + +[[package]] +name = "libnvvm-sys" +version = "0.2.1" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "cuda-target-spec", + "libloading", + "thiserror 2.0.21", +] + +[[package]] +name = "linux-raw-sys" +version = "0.4.15" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d26c52dbd32dccf2d10cac7725f8eae5296885fb5703b261f7d0a0739ec807ab" + +[[package]] +name = "lock_api" +version = "0.4.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "224399e74b87b5f3557511d98dff8b14089b3dadafcab6bb93eab67d3aace965" +dependencies = [ + "scopeguard", +] + +[[package]] +name = "log" +version = "0.4.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9f8bd3e56ce4dfc153cf470fffbfa98c7620958b312ca5c3a4b8d5181fd13c6" + +[[package]] +name = "memchr" +version = "2.8.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cf8baf1c55e62ffcace7a9f06f4bd9cd3f0c4beb022d3b367256b91b87513d98" + +[[package]] +name = "minimal-lexical" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "68354c5c6bd36d73ff3feceb05efa59b6acb7626617f4962be322a825e61f79a" + +[[package]] +name = "nom" +version = "7.1.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d273983c5a657a70a3e8f2a01329822f3b8c8172b73826411a55751e404a0a4a" +dependencies = [ + "memchr", + "minimal-lexical", +] + +[[package]] +name = "nvjitlink-sys" +version = "0.2.1" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "libloading", + "thiserror 2.0.21", +] + +[[package]] +name = "object" +version = "0.36.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "62948e14d923ea95ea2c7c86c71013138b66525b86bdc08d2dcc262bdb497b87" +dependencies = [ + "crc32fast", + "hashbrown 0.15.5", + "indexmap", + "memchr", +] + +[[package]] +name = "once_cell" +version = "1.21.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9f7c3e4beb33f85d45ae3e3a1792185706c8e16d043238c593331cc7cd313b50" + +[[package]] +name = "oxide-artifacts" +version = "0.2.1" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "object", +] + +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "serde", + "serde_json", + "thiserror 2.0.21", +] + +[[package]] +name = "parking_lot_core" +version = "0.9.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2621685985a2ebf1c516881c026032ac7deafcda1a2c9b7850dc81e3dfcb64c1" +dependencies = [ + "cfg-if", + "libc", + "redox_syscall", + "smallvec", + "windows-link", +] + +[[package]] +name = "pin-project-lite" +version = "0.2.17" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a89322df9ebe1c1578d689c92318e070967d1042b512afbe49518723f4e6d5cd" + +[[package]] +name = "prettyplease" +version = "0.2.37" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "479ca8adacdd7ce8f1fb39ce9ecccbfe93a3f1344b3d0d97f20bc0196208f62b" +dependencies = [ + "proc-macro2", + "syn 2.0.119", +] + +[[package]] +name = "proc-macro2" +version = "1.0.107" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "985e7ec9bb745e6ce6535b544d84d6cd6f7ad8bd711c398938ae983b91a766d9" +dependencies = [ + "unicode-ident", +] + +[[package]] +name = "ptx-parse" +version = "0.1.0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" + +[[package]] +name = "quote" +version = "1.0.47" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1fbf4db142a473a8d80c26bbf18454ed458bf8d26c8219c331daecfdbd079001" +dependencies = [ + "proc-macro2", +] + +[[package]] +name = "redox_syscall" +version = "0.5.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ed2bf2547551a7053d6fdfafda3f938979645c44812fbfcda098faae3f1a362d" +dependencies = [ + "bitflags", +] + +[[package]] +name = "regex" +version = "1.13.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f020237b6c8eed93db2e2cb53c00c60a8e1bc73da7d073199a1180401450218d" +dependencies = [ + "aho-corasick", + "memchr", + "regex-automata", + "regex-syntax", +] + +[[package]] +name = "regex-automata" +version = "0.4.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ad8553b9b26413251cbf30e620595c7a41b3887f03da04579c0e6b0d6a06b4b2" +dependencies = [ + "aho-corasick", + "memchr", + "regex-syntax", +] + +[[package]] +name = "regex-syntax" +version = "0.8.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d6f6ff9a378485b298a5286656da665ba74413d36db0979633275d2e708145d4" + +[[package]] +name = "reserved-oxide-symbols" +version = "0.2.1" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" + +[[package]] +name = "rustc-hash" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "08d43f7aa6b08d49f382cde6a7982047c3426db949b1424bc4b7ec9ae12c6ce2" + +[[package]] +name = "rustc-hash" +version = "2.1.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6b1e7f9a428571be2dc5bc0505c13fb6bf936822b894ec87abf8a08a4e51742d" + +[[package]] +name = "rustix" +version = "0.38.44" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fdb5bc1ae2baa591800df16c9ca78619bf65c0488b41b96ccec5d11220d8c154" +dependencies = [ + "bitflags", + "errno", + "libc", + "linux-raw-sys", + "windows-sys 0.59.0", +] + +[[package]] +name = "scopeguard" +version = "1.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "94143f37725109f92c262ed2cf5e59bce7498c01bcc1502d7b9afe439a4e9f49" + +[[package]] +name = "serde" +version = "1.0.229" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4148590afebada386688f18773da617792bf2ef03ffc1e4cbd2b1d45b023e0ba" +dependencies = [ + "serde_core", + "serde_derive", +] + +[[package]] +name = "serde_core" +version = "1.0.229" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "67dca2c9c51e58a4791a4b1ed58308b39c64224d349a935ab5039aa360942a48" +dependencies = [ + "serde_derive", +] + +[[package]] +name = "serde_derive" +version = "1.0.229" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e7a5d71263a5a7d47b41f6b3f06ba276f10cc18b0931f1799f710578e2309348" +dependencies = [ + "proc-macro2", + "quote", + "syn 3.0.9", +] + +[[package]] +name = "serde_json" +version = "1.0.154" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e7e9cc8b1b85264074fbcc02a88680c4096b1e47df8f739dceb03bf482f04bd6" +dependencies = [ + "itoa", + "memchr", + "serde", + "serde_core", + "zmij", +] + +[[package]] +name = "sha2" +version = "0.10.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a7507d819769d01a365ab707794a4084392c824f54a7a6a7862f8c3d0892b283" +dependencies = [ + "cfg-if", + "cpufeatures", + "digest", +] + +[[package]] +name = "shlex" +version = "1.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0fda2ff0d084019ba4d7c6f371c95d8fd75ce3524c3cb8fb653a3023f6323e64" + +[[package]] +name = "slab" +version = "0.4.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0c790de23124f9ab44544d7ac05d60440adc586479ce501c1d6d7da3cd8c9cf5" + +[[package]] +name = "smallvec" +version = "1.16.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5b3dc8af474f516a851ff4bd12db780f948b9250ad37211e4eec0bccea54e01b" + +[[package]] +name = "syn" +version = "2.0.119" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "872831b642d1a07999a962a351ed35b955ea2cfc8f3862091e2a240a84f17297" +dependencies = [ + "proc-macro2", + "quote", + "unicode-ident", +] + +[[package]] +name = "syn" +version = "3.0.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d78c8dee4c7bf0e14673097256fed6142ce9d3b85a408189d07482442145823b" +dependencies = [ + "proc-macro2", + "quote", + "unicode-ident", +] + +[[package]] +name = "thiserror" +version = "1.0.69" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b6aaf5339b578ea85b50e080feb250a3e8ae8cfcdff9a461c9ec2904bc923f52" +dependencies = [ + "thiserror-impl 1.0.69", +] + +[[package]] +name = "thiserror" +version = "2.0.21" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "09e52cb86a36cede5cb101bf8908837b3e4c6e5e59fe7fd85c23fb56200d189e" +dependencies = [ + "thiserror-impl 2.0.21", +] + +[[package]] +name = "thiserror-impl" +version = "1.0.69" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4fee6c4efc90059e10f81e6d42c60a18f76588c3d74cb83a0b242a2b6c7504c1" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", +] + +[[package]] +name = "thiserror-impl" +version = "2.0.21" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fe5197923287db20a58125f0bc85c062f7f2c892de97b18c356f9efb14b28524" +dependencies = [ + "proc-macro2", + "quote", + "syn 3.0.9", +] + +[[package]] +name = "typenum" +version = "1.20.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b6f5e870be6c3b371b77fe0ee0bafb859fa4964b4404c27de1d380043c4dda20" + +[[package]] +name = "unicode-ident" +version = "1.0.27" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a2c754d6c33795a1c324727428e5a7dedb5b06195f9890bdbcba760d3e246563" + +[[package]] +name = "version_check" +version = "0.9.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0b928f33d975fc6ad9f86c8f283853ad26bdd5b10b7f1542aa2fa15e2289105a" + +[[package]] +name = "which" +version = "4.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "87ba24419a2078cd2b0f2ede2691b6c66d8e47836da3b6db8265ebad47afbfc7" +dependencies = [ + "either", + "home", + "once_cell", + "rustix", +] + +[[package]] +name = "windows-link" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f0805222e57f7521d6a62e36fa9163bc891acd422f971defe97d64e70d0a4fe5" + +[[package]] +name = "windows-sys" +version = "0.59.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1e38bc4d79ed67fd075bcc251a1c39b32a1776bbe92e5bef1f0bf1f8c531853b" +dependencies = [ + "windows-targets", +] + +[[package]] +name = "windows-sys" +version = "0.61.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ae137229bcbd6cdf0f7b80a31df61766145077ddf49416a728b02cb3921ff3fc" +dependencies = [ + "windows-link", +] + +[[package]] +name = "windows-targets" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9b724f72796e036ab90c1021d4780d4d3d648aca59e491e6b98e725b84e99973" +dependencies = [ + "windows_aarch64_gnullvm", + "windows_aarch64_msvc", + "windows_i686_gnu", + "windows_i686_gnullvm", + "windows_i686_msvc", + "windows_x86_64_gnu", + "windows_x86_64_gnullvm", + "windows_x86_64_msvc", +] + +[[package]] +name = "windows_aarch64_gnullvm" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "32a4622180e7a0ec044bb555404c800bc9fd9ec262ec147edd5989ccd0c02cd3" + +[[package]] +name = "windows_aarch64_msvc" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "09ec2a7bb152e2252b53fa7803150007879548bc709c039df7627cabbd05d469" + +[[package]] +name = "windows_i686_gnu" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8e9b5ad5ab802e97eb8e295ac6720e509ee4c243f69d781394014ebfe8bbfa0b" + +[[package]] +name = "windows_i686_gnullvm" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0eee52d38c090b3caa76c563b86c3a4bd71ef1a819287c19d586d7334ae8ed66" + +[[package]] +name = "windows_i686_msvc" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "240948bc05c5e7c6dabba28bf89d89ffce3e303022809e73deaefe4f6ec56c66" + +[[package]] +name = "windows_x86_64_gnu" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "147a5c80aabfbf0c7d901cb5895d1de30ef2907eb21fbbab29ca94c5b08b1a78" + +[[package]] +name = "windows_x86_64_gnullvm" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "24d5b23dc417412679681396f2b49f3de8c1473deb516bd34410872eff51ed0d" + +[[package]] +name = "windows_x86_64_msvc" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "589f6da84c646204747d1270a2a5661ea66ed1cced2631d546fdfb155959f9ec" + +[[package]] +name = "zerocopy" +version = "0.8.63" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e5fe1f8f1b06191a00962174c61aa5005e0bb391a6d80d07e24d115c01a92ed8" +dependencies = [ + "zerocopy-derive", +] + +[[package]] +name = "zerocopy-derive" +version = "0.8.63" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "863ad3ac83293fb4d740aedbfdc9240dd8d1a50c1099acd76ce80ce7c7230c7f" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", +] + +[[package]] +name = "zmij" +version = "1.0.23" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "29666d0abbfad1e3dc4dcf6144730dd3a3ab225bbbdac83319345b1b44ccfc1b" diff --git a/rust/cuda-coop/examples/adjacent_difference/Cargo.toml b/rust/cuda-coop/examples/adjacent_difference/Cargo.toml new file mode 100644 index 000000000000..dc5776b3ffc1 --- /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 = "96caed46c776eb49226137eac849e795b83b7ce0" } +cuda-host = { git = "https://github.com/tpn/cuda-rust.git", rev = "96caed46c776eb49226137eac849e795b83b7ce0", features = ["async"] } +cuda-core = { git = "https://github.com/tpn/cuda-rust.git", rev = "96caed46c776eb49226137eac849e795b83b7ce0" } +cuda-async = { git = "https://github.com/tpn/cuda-rust.git", rev = "96caed46c776eb49226137eac849e795b83b7ce0" } diff --git a/rust/cuda-coop/examples/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..f62d83043ab3 --- /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=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "libnvvm-sys", + "nvjitlink-sys", + "serde", + "sha2", + "thiserror 2.0.21", +] + +[[package]] +name = "cuda-async" +version = "0.5.0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "anyhow", + "cuda-bindings", + "cuda-core", + "dashmap", + "futures", + "half", + "once_cell", + "rustc-hash 2.1.3", + "thiserror 1.0.69", +] + +[[package]] +name = "cuda-bindings" +version = "0.5.0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "bindgen", + "libloading", + "prettyplease", + "proc-macro2", + "quote", + "syn 2.0.119", +] + +[[package]] +name = "cuda-coop" +version = "0.0.0" +dependencies = [ + "cuda-coop-codegen", + "cuda-device", +] + +[[package]] +name = "cuda-coop-codegen" +version = "0.0.0" +dependencies = [ + "libloading", + "serde", + "serde_json", + "sha2", +] + +[[package]] +name = "cuda-core" +version = "0.5.0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "anyhow", + "cuda-bindings", + "cuda-core-derive", + "half", + "oxide-artifacts", +] + +[[package]] +name = "cuda-core-derive" +version = "0.5.0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", +] + +[[package]] +name = "cuda-device" +version = "0.2.1" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "cuda-macros", +] + +[[package]] +name = "cuda-host" +version = "0.2.1" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "cuda-artifact-finalizer", + "cuda-async", + "cuda-core", + "cuda-device", + "cuda-macros", + "half", + "libc", + "oxide-artifacts", + "ptx-parse", + "sha2", + "thiserror 2.0.21", +] + +[[package]] +name = "cuda-macros" +version = "0.2.1" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "oxide-tile-export", + "proc-macro2", + "quote", + "reserved-oxide-symbols", + "serde", + "serde_json", + "syn 2.0.119", +] + +[[package]] +name = "cuda-target-spec" +version = "0.1.0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" + +[[package]] +name = "dashmap" +version = "6.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e6361d5c062261c78a176addb82d4c821ae42bed6089de0e12603cd25de2059c" +dependencies = [ + "cfg-if", + "crossbeam-utils", + "hashbrown 0.14.5", + "lock_api", + "once_cell", + "parking_lot_core", +] + +[[package]] +name = "digest" +version = "0.10.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9ed9a281f7bc9b7576e61468ba615a66a5c8cfdff42420a70aa82701a3b1e292" +dependencies = [ + "block-buffer", + "crypto-common", +] + +[[package]] +name = "either" +version = "1.19.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0e9c71c2167ca323c882b99918929403426e2373ea17242ff5653e0d5e1058be" + +[[package]] +name = "equivalent" +version = "1.0.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "00d174d5400e5e8fd687ad1049e2f578285fa914201b1af7e8b112a4546bd826" + +[[package]] +name = "errno" +version = "0.3.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "39cab71617ae0d63f51a36d69f866391735b51691dbda63cf6f96d042b63efeb" +dependencies = [ + "libc", + "windows-sys 0.61.2", +] + +[[package]] +name = "foldhash" +version = "0.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d9c4f5dac5e15c24eb999c26181a6ca40b39fe946cbe4c263c7209467bc83af2" + +[[package]] +name = "futures" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9a31d2a3fbaaeb2af2368bbdd904aa8e812d3c04a1ee10d3171f52d556e5d0a3" +dependencies = [ + "futures-channel", + "futures-core", + "futures-executor", + "futures-io", + "futures-sink", + "futures-task", + "futures-util", +] + +[[package]] +name = "futures-channel" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b1f9e3d69d39e4862ffed03ed071a76f9a13ba1d9109d355b0f0aa6b15e393c4" +dependencies = [ + "futures-core", + "futures-sink", +] + +[[package]] +name = "futures-core" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "92d699e522242e69e3003b94ecc1f960f3a5e015aa7c5d7486e65ad01dd94f5e" + +[[package]] +name = "futures-executor" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "031b47cf1a3c6cc8bc2fc76cd437f521619387907d469316e7c0bc278f1f5432" +dependencies = [ + "futures-core", + "futures-task", + "futures-util", +] + +[[package]] +name = "futures-io" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "53c0fa8157de1303bfffdaa1cc2a673bfffb60102f76b0ef4441659124373fed" + +[[package]] +name = "futures-macro" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9fb9654ba8355388abeb8dcb4fc62f511300867002afc858860463bdd9fe0c44" +dependencies = [ + "proc-macro2", + "quote", + "syn 3.0.9", +] + +[[package]] +name = "futures-sink" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1944426bf7d03f1d14f708785e4b33efd750b36d48a157b836b3efc15ede8e1d" + +[[package]] +name = "futures-task" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cd417de3d1d015fc3bfd2b1ea46dfc7bab72ef86f1cc7cc9c78e728b34a6d1fd" + +[[package]] +name = "futures-util" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0d50a92467f8ba5dd6e3ee5d4bd04d73ab2e4e1c44474a0674821dfce14b79bc" +dependencies = [ + "futures-channel", + "futures-core", + "futures-io", + "futures-macro", + "futures-sink", + "futures-task", + "memchr", + "pin-project-lite", + "slab", +] + +[[package]] +name = "generic-array" +version = "0.14.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "85649ca51fd72272d7821adaf274ad91c288277713d9c18820d8499a7ff69e9a" +dependencies = [ + "typenum", + "version_check", +] + +[[package]] +name = "glob" +version = "0.3.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e4eba85ea1d0a966a983acd07deee566e67395d2d96b6fb39e62b5a833f1eb0b" + +[[package]] +name = "half" +version = "2.7.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" +dependencies = [ + "cfg-if", + "crunchy", + "zerocopy", +] + +[[package]] +name = "hashbrown" +version = "0.14.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e5274423e17b7c9fc20b6e7e208532f9b19825d82dfd615708b70edd83df41f1" + +[[package]] +name = "hashbrown" +version = "0.15.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9229cfe53dfd69f0609a49f65461bd93001ea1ef889cd5529dd176593f5338a1" +dependencies = [ + "foldhash", +] + +[[package]] +name = "hashbrown" +version = "0.17.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ed5909b6e89a2db4456e54cd5f673791d7eca6732202bbf2a9cc504fe2f9b84a" + +[[package]] +name = "home" +version = "0.5.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cc627f471c528ff0c4a49e1d5e60450c8f6461dd6d10ba9dcd3a61d3dff7728d" +dependencies = [ + "windows-sys 0.61.2", +] + +[[package]] +name = "indexmap" +version = "2.14.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cc4e190f5d26ca7051642629da2c52fc03bde85a03197c99408dcd291734c855" +dependencies = [ + "equivalent", + "hashbrown 0.17.1", +] + +[[package]] +name = "itertools" +version = "0.12.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ba291022dbbd398a455acf126c1e341954079855bc60dfdda641363bd6922569" +dependencies = [ + "either", +] + +[[package]] +name = "itoa" +version = "1.0.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8f42a60cbdf9a97f5d2305f08a87dc4e09308d1276d28c869c684d7777685682" + +[[package]] +name = "lazy_static" +version = "1.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "20870f649af7073d53e38067b2a84312175d56ea15217e1b15bc83506ec50afb" + +[[package]] +name = "lazycell" +version = "1.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "830d08ce1d1d941e6b30645f1a0eb5643013d835ce3779a5fc208261dbe10f55" + +[[package]] +name = "libc" +version = "0.2.191" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "65cd87c9a6d60536580b316ad985726104730a5151c99845f2490428ef91240a" + +[[package]] +name = "libloading" +version = "0.8.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d7c4b02199fee7c5d21a5ae7d8cfa79a6ef5bb2fc834d6e9058e89c825efdc55" +dependencies = [ + "cfg-if", + "windows-link", +] + +[[package]] +name = "libnvvm-sys" +version = "0.2.1" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "cuda-target-spec", + "libloading", + "thiserror 2.0.21", +] + +[[package]] +name = "linux-raw-sys" +version = "0.4.15" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d26c52dbd32dccf2d10cac7725f8eae5296885fb5703b261f7d0a0739ec807ab" + +[[package]] +name = "lock_api" +version = "0.4.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "224399e74b87b5f3557511d98dff8b14089b3dadafcab6bb93eab67d3aace965" +dependencies = [ + "scopeguard", +] + +[[package]] +name = "log" +version = "0.4.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9f8bd3e56ce4dfc153cf470fffbfa98c7620958b312ca5c3a4b8d5181fd13c6" + +[[package]] +name = "memchr" +version = "2.8.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cf8baf1c55e62ffcace7a9f06f4bd9cd3f0c4beb022d3b367256b91b87513d98" + +[[package]] +name = "minimal-lexical" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "68354c5c6bd36d73ff3feceb05efa59b6acb7626617f4962be322a825e61f79a" + +[[package]] +name = "nom" +version = "7.1.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d273983c5a657a70a3e8f2a01329822f3b8c8172b73826411a55751e404a0a4a" +dependencies = [ + "memchr", + "minimal-lexical", +] + +[[package]] +name = "nvjitlink-sys" +version = "0.2.1" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "libloading", + "thiserror 2.0.21", +] + +[[package]] +name = "object" +version = "0.36.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "62948e14d923ea95ea2c7c86c71013138b66525b86bdc08d2dcc262bdb497b87" +dependencies = [ + "crc32fast", + "hashbrown 0.15.5", + "indexmap", + "memchr", +] + +[[package]] +name = "once_cell" +version = "1.21.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9f7c3e4beb33f85d45ae3e3a1792185706c8e16d043238c593331cc7cd313b50" + +[[package]] +name = "oxide-artifacts" +version = "0.2.1" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "object", +] + +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "serde", + "serde_json", + "thiserror 2.0.21", +] + +[[package]] +name = "parking_lot_core" +version = "0.9.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2621685985a2ebf1c516881c026032ac7deafcda1a2c9b7850dc81e3dfcb64c1" +dependencies = [ + "cfg-if", + "libc", + "redox_syscall", + "smallvec", + "windows-link", +] + +[[package]] +name = "pin-project-lite" +version = "0.2.17" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a89322df9ebe1c1578d689c92318e070967d1042b512afbe49518723f4e6d5cd" + +[[package]] +name = "prettyplease" +version = "0.2.37" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "479ca8adacdd7ce8f1fb39ce9ecccbfe93a3f1344b3d0d97f20bc0196208f62b" +dependencies = [ + "proc-macro2", + "syn 2.0.119", +] + +[[package]] +name = "proc-macro2" +version = "1.0.107" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "985e7ec9bb745e6ce6535b544d84d6cd6f7ad8bd711c398938ae983b91a766d9" +dependencies = [ + "unicode-ident", +] + +[[package]] +name = "ptx-parse" +version = "0.1.0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" + +[[package]] +name = "quote" +version = "1.0.47" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1fbf4db142a473a8d80c26bbf18454ed458bf8d26c8219c331daecfdbd079001" +dependencies = [ + "proc-macro2", +] + +[[package]] +name = "redox_syscall" +version = "0.5.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ed2bf2547551a7053d6fdfafda3f938979645c44812fbfcda098faae3f1a362d" +dependencies = [ + "bitflags", +] + +[[package]] +name = "regex" +version = "1.13.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f020237b6c8eed93db2e2cb53c00c60a8e1bc73da7d073199a1180401450218d" +dependencies = [ + "aho-corasick", + "memchr", + "regex-automata", + "regex-syntax", +] + +[[package]] +name = "regex-automata" +version = "0.4.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ad8553b9b26413251cbf30e620595c7a41b3887f03da04579c0e6b0d6a06b4b2" +dependencies = [ + "aho-corasick", + "memchr", + "regex-syntax", +] + +[[package]] +name = "regex-syntax" +version = "0.8.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d6f6ff9a378485b298a5286656da665ba74413d36db0979633275d2e708145d4" + +[[package]] +name = "reserved-oxide-symbols" +version = "0.2.1" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" + +[[package]] +name = "rustc-hash" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "08d43f7aa6b08d49f382cde6a7982047c3426db949b1424bc4b7ec9ae12c6ce2" + +[[package]] +name = "rustc-hash" +version = "2.1.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6b1e7f9a428571be2dc5bc0505c13fb6bf936822b894ec87abf8a08a4e51742d" + +[[package]] +name = "rustix" +version = "0.38.44" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fdb5bc1ae2baa591800df16c9ca78619bf65c0488b41b96ccec5d11220d8c154" +dependencies = [ + "bitflags", + "errno", + "libc", + "linux-raw-sys", + "windows-sys 0.59.0", +] + +[[package]] +name = "scopeguard" +version = "1.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "94143f37725109f92c262ed2cf5e59bce7498c01bcc1502d7b9afe439a4e9f49" + +[[package]] +name = "serde" +version = "1.0.229" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4148590afebada386688f18773da617792bf2ef03ffc1e4cbd2b1d45b023e0ba" +dependencies = [ + "serde_core", + "serde_derive", +] + +[[package]] +name = "serde_core" +version = "1.0.229" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "67dca2c9c51e58a4791a4b1ed58308b39c64224d349a935ab5039aa360942a48" +dependencies = [ + "serde_derive", +] + +[[package]] +name = "serde_derive" +version = "1.0.229" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e7a5d71263a5a7d47b41f6b3f06ba276f10cc18b0931f1799f710578e2309348" +dependencies = [ + "proc-macro2", + "quote", + "syn 3.0.9", +] + +[[package]] +name = "serde_json" +version = "1.0.154" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e7e9cc8b1b85264074fbcc02a88680c4096b1e47df8f739dceb03bf482f04bd6" +dependencies = [ + "itoa", + "memchr", + "serde", + "serde_core", + "zmij", +] + +[[package]] +name = "sha2" +version = "0.10.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a7507d819769d01a365ab707794a4084392c824f54a7a6a7862f8c3d0892b283" +dependencies = [ + "cfg-if", + "cpufeatures", + "digest", +] + +[[package]] +name = "shlex" +version = "1.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0fda2ff0d084019ba4d7c6f371c95d8fd75ce3524c3cb8fb653a3023f6323e64" + +[[package]] +name = "slab" +version = "0.4.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0c790de23124f9ab44544d7ac05d60440adc586479ce501c1d6d7da3cd8c9cf5" + +[[package]] +name = "smallvec" +version = "1.16.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5b3dc8af474f516a851ff4bd12db780f948b9250ad37211e4eec0bccea54e01b" + +[[package]] +name = "syn" +version = "2.0.119" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "872831b642d1a07999a962a351ed35b955ea2cfc8f3862091e2a240a84f17297" +dependencies = [ + "proc-macro2", + "quote", + "unicode-ident", +] + +[[package]] +name = "syn" +version = "3.0.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d78c8dee4c7bf0e14673097256fed6142ce9d3b85a408189d07482442145823b" +dependencies = [ + "proc-macro2", + "quote", + "unicode-ident", +] + +[[package]] +name = "thiserror" +version = "1.0.69" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b6aaf5339b578ea85b50e080feb250a3e8ae8cfcdff9a461c9ec2904bc923f52" +dependencies = [ + "thiserror-impl 1.0.69", +] + +[[package]] +name = "thiserror" +version = "2.0.21" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "09e52cb86a36cede5cb101bf8908837b3e4c6e5e59fe7fd85c23fb56200d189e" +dependencies = [ + "thiserror-impl 2.0.21", +] + +[[package]] +name = "thiserror-impl" +version = "1.0.69" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4fee6c4efc90059e10f81e6d42c60a18f76588c3d74cb83a0b242a2b6c7504c1" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", +] + +[[package]] +name = "thiserror-impl" +version = "2.0.21" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fe5197923287db20a58125f0bc85c062f7f2c892de97b18c356f9efb14b28524" +dependencies = [ + "proc-macro2", + "quote", + "syn 3.0.9", +] + +[[package]] +name = "typenum" +version = "1.20.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b6f5e870be6c3b371b77fe0ee0bafb859fa4964b4404c27de1d380043c4dda20" + +[[package]] +name = "unicode-ident" +version = "1.0.27" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a2c754d6c33795a1c324727428e5a7dedb5b06195f9890bdbcba760d3e246563" + +[[package]] +name = "version_check" +version = "0.9.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0b928f33d975fc6ad9f86c8f283853ad26bdd5b10b7f1542aa2fa15e2289105a" + +[[package]] +name = "which" +version = "4.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "87ba24419a2078cd2b0f2ede2691b6c66d8e47836da3b6db8265ebad47afbfc7" +dependencies = [ + "either", + "home", + "once_cell", + "rustix", +] + +[[package]] +name = "windows-link" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f0805222e57f7521d6a62e36fa9163bc891acd422f971defe97d64e70d0a4fe5" + +[[package]] +name = "windows-sys" +version = "0.59.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1e38bc4d79ed67fd075bcc251a1c39b32a1776bbe92e5bef1f0bf1f8c531853b" +dependencies = [ + "windows-targets", +] + +[[package]] +name = "windows-sys" +version = "0.61.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ae137229bcbd6cdf0f7b80a31df61766145077ddf49416a728b02cb3921ff3fc" +dependencies = [ + "windows-link", +] + +[[package]] +name = "windows-targets" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9b724f72796e036ab90c1021d4780d4d3d648aca59e491e6b98e725b84e99973" +dependencies = [ + "windows_aarch64_gnullvm", + "windows_aarch64_msvc", + "windows_i686_gnu", + "windows_i686_gnullvm", + "windows_i686_msvc", + "windows_x86_64_gnu", + "windows_x86_64_gnullvm", + "windows_x86_64_msvc", +] + +[[package]] +name = "windows_aarch64_gnullvm" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "32a4622180e7a0ec044bb555404c800bc9fd9ec262ec147edd5989ccd0c02cd3" + +[[package]] +name = "windows_aarch64_msvc" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "09ec2a7bb152e2252b53fa7803150007879548bc709c039df7627cabbd05d469" + +[[package]] +name = "windows_i686_gnu" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8e9b5ad5ab802e97eb8e295ac6720e509ee4c243f69d781394014ebfe8bbfa0b" + +[[package]] +name = "windows_i686_gnullvm" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0eee52d38c090b3caa76c563b86c3a4bd71ef1a819287c19d586d7334ae8ed66" + +[[package]] +name = "windows_i686_msvc" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "240948bc05c5e7c6dabba28bf89d89ffce3e303022809e73deaefe4f6ec56c66" + +[[package]] +name = "windows_x86_64_gnu" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "147a5c80aabfbf0c7d901cb5895d1de30ef2907eb21fbbab29ca94c5b08b1a78" + +[[package]] +name = "windows_x86_64_gnullvm" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "24d5b23dc417412679681396f2b49f3de8c1473deb516bd34410872eff51ed0d" + +[[package]] +name = "windows_x86_64_msvc" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "589f6da84c646204747d1270a2a5661ea66ed1cced2631d546fdfb155959f9ec" + +[[package]] +name = "zerocopy" +version = "0.8.63" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e5fe1f8f1b06191a00962174c61aa5005e0bb391a6d80d07e24d115c01a92ed8" +dependencies = [ + "zerocopy-derive", +] + +[[package]] +name = "zerocopy-derive" +version = "0.8.63" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "863ad3ac83293fb4d740aedbfdc9240dd8d1a50c1099acd76ce80ce7c7230c7f" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", +] + +[[package]] +name = "zmij" +version = "1.0.23" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "29666d0abbfad1e3dc4dcf6144730dd3a3ab225bbbdac83319345b1b44ccfc1b" diff --git a/rust/cuda-coop/examples/discontinuity/Cargo.toml b/rust/cuda-coop/examples/discontinuity/Cargo.toml new file mode 100644 index 000000000000..791cccdaab0e --- /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 = "96caed46c776eb49226137eac849e795b83b7ce0" } +cuda-host = { git = "https://github.com/tpn/cuda-rust.git", rev = "96caed46c776eb49226137eac849e795b83b7ce0", features = ["async"] } +cuda-core = { git = "https://github.com/tpn/cuda-rust.git", rev = "96caed46c776eb49226137eac849e795b83b7ce0" } +cuda-async = { git = "https://github.com/tpn/cuda-rust.git", rev = "96caed46c776eb49226137eac849e795b83b7ce0" } diff --git a/rust/cuda-coop/examples/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..17bac02d6bdf --- /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=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "libnvvm-sys", + "nvjitlink-sys", + "serde", + "sha2", + "thiserror 2.0.21", +] + +[[package]] +name = "cuda-async" +version = "0.5.0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "anyhow", + "cuda-bindings", + "cuda-core", + "dashmap", + "futures", + "half", + "once_cell", + "rustc-hash 2.1.3", + "thiserror 1.0.69", +] + +[[package]] +name = "cuda-bindings" +version = "0.5.0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "bindgen", + "libloading", + "prettyplease", + "proc-macro2", + "quote", + "syn 2.0.119", +] + +[[package]] +name = "cuda-coop" +version = "0.0.0" +dependencies = [ + "cuda-coop-codegen", + "cuda-device", +] + +[[package]] +name = "cuda-coop-codegen" +version = "0.0.0" +dependencies = [ + "libloading", + "serde", + "serde_json", + "sha2", +] + +[[package]] +name = "cuda-core" +version = "0.5.0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "anyhow", + "cuda-bindings", + "cuda-core-derive", + "half", + "oxide-artifacts", +] + +[[package]] +name = "cuda-core-derive" +version = "0.5.0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", +] + +[[package]] +name = "cuda-device" +version = "0.2.1" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "cuda-macros", +] + +[[package]] +name = "cuda-host" +version = "0.2.1" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "cuda-artifact-finalizer", + "cuda-async", + "cuda-core", + "cuda-device", + "cuda-macros", + "half", + "libc", + "oxide-artifacts", + "ptx-parse", + "sha2", + "thiserror 2.0.21", +] + +[[package]] +name = "cuda-macros" +version = "0.2.1" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "oxide-tile-export", + "proc-macro2", + "quote", + "reserved-oxide-symbols", + "serde", + "serde_json", + "syn 2.0.119", +] + +[[package]] +name = "cuda-target-spec" +version = "0.1.0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" + +[[package]] +name = "dashmap" +version = "6.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e6361d5c062261c78a176addb82d4c821ae42bed6089de0e12603cd25de2059c" +dependencies = [ + "cfg-if", + "crossbeam-utils", + "hashbrown 0.14.5", + "lock_api", + "once_cell", + "parking_lot_core", +] + +[[package]] +name = "digest" +version = "0.10.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9ed9a281f7bc9b7576e61468ba615a66a5c8cfdff42420a70aa82701a3b1e292" +dependencies = [ + "block-buffer", + "crypto-common", +] + +[[package]] +name = "either" +version = "1.19.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0e9c71c2167ca323c882b99918929403426e2373ea17242ff5653e0d5e1058be" + +[[package]] +name = "equivalent" +version = "1.0.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "00d174d5400e5e8fd687ad1049e2f578285fa914201b1af7e8b112a4546bd826" + +[[package]] +name = "errno" +version = "0.3.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "39cab71617ae0d63f51a36d69f866391735b51691dbda63cf6f96d042b63efeb" +dependencies = [ + "libc", + "windows-sys 0.61.2", +] + +[[package]] +name = "foldhash" +version = "0.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d9c4f5dac5e15c24eb999c26181a6ca40b39fe946cbe4c263c7209467bc83af2" + +[[package]] +name = "futures" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9a31d2a3fbaaeb2af2368bbdd904aa8e812d3c04a1ee10d3171f52d556e5d0a3" +dependencies = [ + "futures-channel", + "futures-core", + "futures-executor", + "futures-io", + "futures-sink", + "futures-task", + "futures-util", +] + +[[package]] +name = "futures-channel" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b1f9e3d69d39e4862ffed03ed071a76f9a13ba1d9109d355b0f0aa6b15e393c4" +dependencies = [ + "futures-core", + "futures-sink", +] + +[[package]] +name = "futures-core" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "92d699e522242e69e3003b94ecc1f960f3a5e015aa7c5d7486e65ad01dd94f5e" + +[[package]] +name = "futures-executor" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "031b47cf1a3c6cc8bc2fc76cd437f521619387907d469316e7c0bc278f1f5432" +dependencies = [ + "futures-core", + "futures-task", + "futures-util", +] + +[[package]] +name = "futures-io" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "53c0fa8157de1303bfffdaa1cc2a673bfffb60102f76b0ef4441659124373fed" + +[[package]] +name = "futures-macro" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9fb9654ba8355388abeb8dcb4fc62f511300867002afc858860463bdd9fe0c44" +dependencies = [ + "proc-macro2", + "quote", + "syn 3.0.9", +] + +[[package]] +name = "futures-sink" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1944426bf7d03f1d14f708785e4b33efd750b36d48a157b836b3efc15ede8e1d" + +[[package]] +name = "futures-task" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cd417de3d1d015fc3bfd2b1ea46dfc7bab72ef86f1cc7cc9c78e728b34a6d1fd" + +[[package]] +name = "futures-util" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0d50a92467f8ba5dd6e3ee5d4bd04d73ab2e4e1c44474a0674821dfce14b79bc" +dependencies = [ + "futures-channel", + "futures-core", + "futures-io", + "futures-macro", + "futures-sink", + "futures-task", + "memchr", + "pin-project-lite", + "slab", +] + +[[package]] +name = "generic-array" +version = "0.14.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "85649ca51fd72272d7821adaf274ad91c288277713d9c18820d8499a7ff69e9a" +dependencies = [ + "typenum", + "version_check", +] + +[[package]] +name = "glob" +version = "0.3.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e4eba85ea1d0a966a983acd07deee566e67395d2d96b6fb39e62b5a833f1eb0b" + +[[package]] +name = "half" +version = "2.7.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" +dependencies = [ + "cfg-if", + "crunchy", + "zerocopy", +] + +[[package]] +name = "hashbrown" +version = "0.14.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e5274423e17b7c9fc20b6e7e208532f9b19825d82dfd615708b70edd83df41f1" + +[[package]] +name = "hashbrown" +version = "0.15.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9229cfe53dfd69f0609a49f65461bd93001ea1ef889cd5529dd176593f5338a1" +dependencies = [ + "foldhash", +] + +[[package]] +name = "hashbrown" +version = "0.17.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ed5909b6e89a2db4456e54cd5f673791d7eca6732202bbf2a9cc504fe2f9b84a" + +[[package]] +name = "home" +version = "0.5.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cc627f471c528ff0c4a49e1d5e60450c8f6461dd6d10ba9dcd3a61d3dff7728d" +dependencies = [ + "windows-sys 0.61.2", +] + +[[package]] +name = "indexmap" +version = "2.14.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cc4e190f5d26ca7051642629da2c52fc03bde85a03197c99408dcd291734c855" +dependencies = [ + "equivalent", + "hashbrown 0.17.1", +] + +[[package]] +name = "itertools" +version = "0.12.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ba291022dbbd398a455acf126c1e341954079855bc60dfdda641363bd6922569" +dependencies = [ + "either", +] + +[[package]] +name = "itoa" +version = "1.0.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8f42a60cbdf9a97f5d2305f08a87dc4e09308d1276d28c869c684d7777685682" + +[[package]] +name = "lazy_static" +version = "1.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "20870f649af7073d53e38067b2a84312175d56ea15217e1b15bc83506ec50afb" + +[[package]] +name = "lazycell" +version = "1.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "830d08ce1d1d941e6b30645f1a0eb5643013d835ce3779a5fc208261dbe10f55" + +[[package]] +name = "libc" +version = "0.2.191" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "65cd87c9a6d60536580b316ad985726104730a5151c99845f2490428ef91240a" + +[[package]] +name = "libloading" +version = "0.8.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d7c4b02199fee7c5d21a5ae7d8cfa79a6ef5bb2fc834d6e9058e89c825efdc55" +dependencies = [ + "cfg-if", + "windows-link", +] + +[[package]] +name = "libnvvm-sys" +version = "0.2.1" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "cuda-target-spec", + "libloading", + "thiserror 2.0.21", +] + +[[package]] +name = "linux-raw-sys" +version = "0.4.15" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d26c52dbd32dccf2d10cac7725f8eae5296885fb5703b261f7d0a0739ec807ab" + +[[package]] +name = "lock_api" +version = "0.4.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "224399e74b87b5f3557511d98dff8b14089b3dadafcab6bb93eab67d3aace965" +dependencies = [ + "scopeguard", +] + +[[package]] +name = "log" +version = "0.4.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9f8bd3e56ce4dfc153cf470fffbfa98c7620958b312ca5c3a4b8d5181fd13c6" + +[[package]] +name = "memchr" +version = "2.8.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cf8baf1c55e62ffcace7a9f06f4bd9cd3f0c4beb022d3b367256b91b87513d98" + +[[package]] +name = "minimal-lexical" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "68354c5c6bd36d73ff3feceb05efa59b6acb7626617f4962be322a825e61f79a" + +[[package]] +name = "nom" +version = "7.1.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d273983c5a657a70a3e8f2a01329822f3b8c8172b73826411a55751e404a0a4a" +dependencies = [ + "memchr", + "minimal-lexical", +] + +[[package]] +name = "nvjitlink-sys" +version = "0.2.1" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "libloading", + "thiserror 2.0.21", +] + +[[package]] +name = "object" +version = "0.36.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "62948e14d923ea95ea2c7c86c71013138b66525b86bdc08d2dcc262bdb497b87" +dependencies = [ + "crc32fast", + "hashbrown 0.15.5", + "indexmap", + "memchr", +] + +[[package]] +name = "once_cell" +version = "1.21.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9f7c3e4beb33f85d45ae3e3a1792185706c8e16d043238c593331cc7cd313b50" + +[[package]] +name = "oxide-artifacts" +version = "0.2.1" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "object", +] + +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "serde", + "serde_json", + "thiserror 2.0.21", +] + +[[package]] +name = "parking_lot_core" +version = "0.9.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2621685985a2ebf1c516881c026032ac7deafcda1a2c9b7850dc81e3dfcb64c1" +dependencies = [ + "cfg-if", + "libc", + "redox_syscall", + "smallvec", + "windows-link", +] + +[[package]] +name = "pin-project-lite" +version = "0.2.17" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a89322df9ebe1c1578d689c92318e070967d1042b512afbe49518723f4e6d5cd" + +[[package]] +name = "prettyplease" +version = "0.2.37" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "479ca8adacdd7ce8f1fb39ce9ecccbfe93a3f1344b3d0d97f20bc0196208f62b" +dependencies = [ + "proc-macro2", + "syn 2.0.119", +] + +[[package]] +name = "proc-macro2" +version = "1.0.107" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "985e7ec9bb745e6ce6535b544d84d6cd6f7ad8bd711c398938ae983b91a766d9" +dependencies = [ + "unicode-ident", +] + +[[package]] +name = "ptx-parse" +version = "0.1.0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" + +[[package]] +name = "quote" +version = "1.0.47" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1fbf4db142a473a8d80c26bbf18454ed458bf8d26c8219c331daecfdbd079001" +dependencies = [ + "proc-macro2", +] + +[[package]] +name = "redox_syscall" +version = "0.5.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ed2bf2547551a7053d6fdfafda3f938979645c44812fbfcda098faae3f1a362d" +dependencies = [ + "bitflags", +] + +[[package]] +name = "regex" +version = "1.13.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f020237b6c8eed93db2e2cb53c00c60a8e1bc73da7d073199a1180401450218d" +dependencies = [ + "aho-corasick", + "memchr", + "regex-automata", + "regex-syntax", +] + +[[package]] +name = "regex-automata" +version = "0.4.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ad8553b9b26413251cbf30e620595c7a41b3887f03da04579c0e6b0d6a06b4b2" +dependencies = [ + "aho-corasick", + "memchr", + "regex-syntax", +] + +[[package]] +name = "regex-syntax" +version = "0.8.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d6f6ff9a378485b298a5286656da665ba74413d36db0979633275d2e708145d4" + +[[package]] +name = "reserved-oxide-symbols" +version = "0.2.1" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" + +[[package]] +name = "rustc-hash" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "08d43f7aa6b08d49f382cde6a7982047c3426db949b1424bc4b7ec9ae12c6ce2" + +[[package]] +name = "rustc-hash" +version = "2.1.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6b1e7f9a428571be2dc5bc0505c13fb6bf936822b894ec87abf8a08a4e51742d" + +[[package]] +name = "rustix" +version = "0.38.44" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fdb5bc1ae2baa591800df16c9ca78619bf65c0488b41b96ccec5d11220d8c154" +dependencies = [ + "bitflags", + "errno", + "libc", + "linux-raw-sys", + "windows-sys 0.59.0", +] + +[[package]] +name = "scopeguard" +version = "1.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "94143f37725109f92c262ed2cf5e59bce7498c01bcc1502d7b9afe439a4e9f49" + +[[package]] +name = "serde" +version = "1.0.229" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4148590afebada386688f18773da617792bf2ef03ffc1e4cbd2b1d45b023e0ba" +dependencies = [ + "serde_core", + "serde_derive", +] + +[[package]] +name = "serde_core" +version = "1.0.229" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "67dca2c9c51e58a4791a4b1ed58308b39c64224d349a935ab5039aa360942a48" +dependencies = [ + "serde_derive", +] + +[[package]] +name = "serde_derive" +version = "1.0.229" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e7a5d71263a5a7d47b41f6b3f06ba276f10cc18b0931f1799f710578e2309348" +dependencies = [ + "proc-macro2", + "quote", + "syn 3.0.9", +] + +[[package]] +name = "serde_json" +version = "1.0.154" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e7e9cc8b1b85264074fbcc02a88680c4096b1e47df8f739dceb03bf482f04bd6" +dependencies = [ + "itoa", + "memchr", + "serde", + "serde_core", + "zmij", +] + +[[package]] +name = "sha2" +version = "0.10.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a7507d819769d01a365ab707794a4084392c824f54a7a6a7862f8c3d0892b283" +dependencies = [ + "cfg-if", + "cpufeatures", + "digest", +] + +[[package]] +name = "shlex" +version = "1.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0fda2ff0d084019ba4d7c6f371c95d8fd75ce3524c3cb8fb653a3023f6323e64" + +[[package]] +name = "slab" +version = "0.4.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0c790de23124f9ab44544d7ac05d60440adc586479ce501c1d6d7da3cd8c9cf5" + +[[package]] +name = "smallvec" +version = "1.16.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5b3dc8af474f516a851ff4bd12db780f948b9250ad37211e4eec0bccea54e01b" + +[[package]] +name = "syn" +version = "2.0.119" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "872831b642d1a07999a962a351ed35b955ea2cfc8f3862091e2a240a84f17297" +dependencies = [ + "proc-macro2", + "quote", + "unicode-ident", +] + +[[package]] +name = "syn" +version = "3.0.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d78c8dee4c7bf0e14673097256fed6142ce9d3b85a408189d07482442145823b" +dependencies = [ + "proc-macro2", + "quote", + "unicode-ident", +] + +[[package]] +name = "thiserror" +version = "1.0.69" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b6aaf5339b578ea85b50e080feb250a3e8ae8cfcdff9a461c9ec2904bc923f52" +dependencies = [ + "thiserror-impl 1.0.69", +] + +[[package]] +name = "thiserror" +version = "2.0.21" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "09e52cb86a36cede5cb101bf8908837b3e4c6e5e59fe7fd85c23fb56200d189e" +dependencies = [ + "thiserror-impl 2.0.21", +] + +[[package]] +name = "thiserror-impl" +version = "1.0.69" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4fee6c4efc90059e10f81e6d42c60a18f76588c3d74cb83a0b242a2b6c7504c1" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", +] + +[[package]] +name = "thiserror-impl" +version = "2.0.21" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fe5197923287db20a58125f0bc85c062f7f2c892de97b18c356f9efb14b28524" +dependencies = [ + "proc-macro2", + "quote", + "syn 3.0.9", +] + +[[package]] +name = "typenum" +version = "1.20.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b6f5e870be6c3b371b77fe0ee0bafb859fa4964b4404c27de1d380043c4dda20" + +[[package]] +name = "unicode-ident" +version = "1.0.27" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a2c754d6c33795a1c324727428e5a7dedb5b06195f9890bdbcba760d3e246563" + +[[package]] +name = "version_check" +version = "0.9.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0b928f33d975fc6ad9f86c8f283853ad26bdd5b10b7f1542aa2fa15e2289105a" + +[[package]] +name = "which" +version = "4.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "87ba24419a2078cd2b0f2ede2691b6c66d8e47836da3b6db8265ebad47afbfc7" +dependencies = [ + "either", + "home", + "once_cell", + "rustix", +] + +[[package]] +name = "windows-link" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f0805222e57f7521d6a62e36fa9163bc891acd422f971defe97d64e70d0a4fe5" + +[[package]] +name = "windows-sys" +version = "0.59.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1e38bc4d79ed67fd075bcc251a1c39b32a1776bbe92e5bef1f0bf1f8c531853b" +dependencies = [ + "windows-targets", +] + +[[package]] +name = "windows-sys" +version = "0.61.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ae137229bcbd6cdf0f7b80a31df61766145077ddf49416a728b02cb3921ff3fc" +dependencies = [ + "windows-link", +] + +[[package]] +name = "windows-targets" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9b724f72796e036ab90c1021d4780d4d3d648aca59e491e6b98e725b84e99973" +dependencies = [ + "windows_aarch64_gnullvm", + "windows_aarch64_msvc", + "windows_i686_gnu", + "windows_i686_gnullvm", + "windows_i686_msvc", + "windows_x86_64_gnu", + "windows_x86_64_gnullvm", + "windows_x86_64_msvc", +] + +[[package]] +name = "windows_aarch64_gnullvm" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "32a4622180e7a0ec044bb555404c800bc9fd9ec262ec147edd5989ccd0c02cd3" + +[[package]] +name = "windows_aarch64_msvc" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "09ec2a7bb152e2252b53fa7803150007879548bc709c039df7627cabbd05d469" + +[[package]] +name = "windows_i686_gnu" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8e9b5ad5ab802e97eb8e295ac6720e509ee4c243f69d781394014ebfe8bbfa0b" + +[[package]] +name = "windows_i686_gnullvm" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0eee52d38c090b3caa76c563b86c3a4bd71ef1a819287c19d586d7334ae8ed66" + +[[package]] +name = "windows_i686_msvc" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "240948bc05c5e7c6dabba28bf89d89ffce3e303022809e73deaefe4f6ec56c66" + +[[package]] +name = "windows_x86_64_gnu" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "147a5c80aabfbf0c7d901cb5895d1de30ef2907eb21fbbab29ca94c5b08b1a78" + +[[package]] +name = "windows_x86_64_gnullvm" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "24d5b23dc417412679681396f2b49f3de8c1473deb516bd34410872eff51ed0d" + +[[package]] +name = "windows_x86_64_msvc" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "589f6da84c646204747d1270a2a5661ea66ed1cced2631d546fdfb155959f9ec" + +[[package]] +name = "zerocopy" +version = "0.8.63" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e5fe1f8f1b06191a00962174c61aa5005e0bb391a6d80d07e24d115c01a92ed8" +dependencies = [ + "zerocopy-derive", +] + +[[package]] +name = "zerocopy-derive" +version = "0.8.63" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "863ad3ac83293fb4d740aedbfdc9240dd8d1a50c1099acd76ce80ce7c7230c7f" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", +] + +[[package]] +name = "zmij" +version = "1.0.23" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "29666d0abbfad1e3dc4dcf6144730dd3a3ab225bbbdac83319345b1b44ccfc1b" diff --git a/rust/cuda-coop/examples/exchange/Cargo.toml b/rust/cuda-coop/examples/exchange/Cargo.toml new file mode 100644 index 000000000000..41575ebcbc2f --- /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 = "96caed46c776eb49226137eac849e795b83b7ce0" } +cuda-host = { git = "https://github.com/tpn/cuda-rust.git", rev = "96caed46c776eb49226137eac849e795b83b7ce0", features = ["async"] } +cuda-core = { git = "https://github.com/tpn/cuda-rust.git", rev = "96caed46c776eb49226137eac849e795b83b7ce0" } +cuda-async = { git = "https://github.com/tpn/cuda-rust.git", rev = "96caed46c776eb49226137eac849e795b83b7ce0" } diff --git a/rust/cuda-coop/examples/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/load_store/.gitignore b/rust/cuda-coop/examples/load_store/.gitignore new file mode 100644 index 000000000000..74bffa8af4a4 --- /dev/null +++ b/rust/cuda-coop/examples/load_store/.gitignore @@ -0,0 +1,7 @@ +# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. +# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception +/target/ +/.oxide-artifacts/ +/*.ll +/*.options +/*.target diff --git a/rust/cuda-coop/examples/load_store/Cargo.lock b/rust/cuda-coop/examples/load_store/Cargo.lock new file mode 100644 index 000000000000..98b133dcd9dd --- /dev/null +++ b/rust/cuda-coop/examples/load_store/Cargo.lock @@ -0,0 +1,1019 @@ +# This file is automatically @generated by Cargo. +# It is not intended for manual editing. +version = 4 + +[[package]] +name = "aho-corasick" +version = "1.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c982642fa9e8606056828ee9a8505737230110bb1099153c79efe865c59d12ba" +dependencies = [ + "memchr", +] + +[[package]] +name = "anyhow" +version = "1.0.104" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "330a5ed07fa54e4702c9d6c4174f74427fc0ef6e214bbd677ae50a5099946470" + +[[package]] +name = "bindgen" +version = "0.69.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "271383c67ccabffb7381723dea0672a673f292304fcb45c01cc648c7a8d58088" +dependencies = [ + "bitflags", + "cexpr", + "clang-sys", + "itertools", + "lazy_static", + "lazycell", + "log", + "prettyplease", + "proc-macro2", + "quote", + "regex", + "rustc-hash 1.1.0", + "shlex", + "syn 2.0.119", + "which", +] + +[[package]] +name = "bitflags" +version = "2.13.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3ded4057c258ba199e2d26386d3af3780957ecaee6c4ef4041c6b4b8b97c0b06" + +[[package]] +name = "block-buffer" +version = "0.10.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3078c7629b62d3f0439517fa394996acacc5cbc91c5a20d8c658e77abd503a71" +dependencies = [ + "generic-array", +] + +[[package]] +name = "cexpr" +version = "0.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6fac387a98bb7c37292057cffc56d62ecb629900026402633ae9160df93a8766" +dependencies = [ + "nom", +] + +[[package]] +name = "cfg-if" +version = "1.0.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4e7648175b45a9a48536d676f68d918270699102aa8dab5496df06904c914600" + +[[package]] +name = "clang-sys" +version = "1.9.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "157a8ba7b480713b56f4c09fd13fc3e0a22a5dfab8097ba61cbc5feef950788a" +dependencies = [ + "glob", + "libc", + "libloading", +] + +[[package]] +name = "coop-load-store" +version = "0.0.0" +dependencies = [ + "coop-load-store-helper", + "cuda-async", + "cuda-coop", + "cuda-core", + "cuda-device", + "cuda-host", +] + +[[package]] +name = "coop-load-store-helper" +version = "0.0.0" +dependencies = [ + "cuda-coop", +] + +[[package]] +name = "cpufeatures" +version = "0.2.17" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "59ed5838eebb26a2bb2e58f6d5b5316989ae9d08bab10e0e6d103e656d1b0280" +dependencies = [ + "libc", +] + +[[package]] +name = "crc32fast" +version = "1.5.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0ba106e844045ddc02c20768ac4b0b7cb45f2cc2ecbd833ac9f5dacb183c2055" +dependencies = [ + "cfg-if", +] + +[[package]] +name = "crossbeam-utils" +version = "0.8.23" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a31eee39dddec8330830986fcd7625edb5a24ec90ea038215273bbc3adb08ac6" + +[[package]] +name = "crunchy" +version = "0.2.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "460fbee9c2c2f33933d720630a6a0bac33ba7053db5344fac858d4b8952d77d5" + +[[package]] +name = "crypto-common" +version = "0.1.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "78c8292055d1c1df0cce5d180393dc8cce0abec0a7102adb6c7b1eef6016d60a" +dependencies = [ + "generic-array", + "typenum", +] + +[[package]] +name = "cuda-artifact-finalizer" +version = "0.2.1" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "libnvvm-sys", + "nvjitlink-sys", + "serde", + "sha2", + "thiserror 2.0.21", +] + +[[package]] +name = "cuda-async" +version = "0.5.0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "anyhow", + "cuda-bindings", + "cuda-core", + "dashmap", + "futures", + "half", + "once_cell", + "rustc-hash 2.1.3", + "thiserror 1.0.69", +] + +[[package]] +name = "cuda-bindings" +version = "0.5.0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "bindgen", + "libloading", + "prettyplease", + "proc-macro2", + "quote", + "syn 2.0.119", +] + +[[package]] +name = "cuda-coop" +version = "0.0.0" +dependencies = [ + "cuda-coop-codegen", + "cuda-device", +] + +[[package]] +name = "cuda-coop-codegen" +version = "0.0.0" +dependencies = [ + "libloading", + "serde", + "serde_json", + "sha2", +] + +[[package]] +name = "cuda-core" +version = "0.5.0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "anyhow", + "cuda-bindings", + "cuda-core-derive", + "half", + "oxide-artifacts", +] + +[[package]] +name = "cuda-core-derive" +version = "0.5.0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", +] + +[[package]] +name = "cuda-device" +version = "0.2.1" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "cuda-macros", +] + +[[package]] +name = "cuda-host" +version = "0.2.1" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "cuda-artifact-finalizer", + "cuda-async", + "cuda-core", + "cuda-device", + "cuda-macros", + "half", + "libc", + "oxide-artifacts", + "ptx-parse", + "sha2", + "thiserror 2.0.21", +] + +[[package]] +name = "cuda-macros" +version = "0.2.1" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "oxide-tile-export", + "proc-macro2", + "quote", + "reserved-oxide-symbols", + "serde", + "serde_json", + "syn 2.0.119", +] + +[[package]] +name = "cuda-target-spec" +version = "0.1.0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" + +[[package]] +name = "dashmap" +version = "6.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e6361d5c062261c78a176addb82d4c821ae42bed6089de0e12603cd25de2059c" +dependencies = [ + "cfg-if", + "crossbeam-utils", + "hashbrown 0.14.5", + "lock_api", + "once_cell", + "parking_lot_core", +] + +[[package]] +name = "digest" +version = "0.10.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9ed9a281f7bc9b7576e61468ba615a66a5c8cfdff42420a70aa82701a3b1e292" +dependencies = [ + "block-buffer", + "crypto-common", +] + +[[package]] +name = "either" +version = "1.19.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0e9c71c2167ca323c882b99918929403426e2373ea17242ff5653e0d5e1058be" + +[[package]] +name = "equivalent" +version = "1.0.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "00d174d5400e5e8fd687ad1049e2f578285fa914201b1af7e8b112a4546bd826" + +[[package]] +name = "errno" +version = "0.3.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "39cab71617ae0d63f51a36d69f866391735b51691dbda63cf6f96d042b63efeb" +dependencies = [ + "libc", + "windows-sys 0.61.2", +] + +[[package]] +name = "foldhash" +version = "0.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d9c4f5dac5e15c24eb999c26181a6ca40b39fe946cbe4c263c7209467bc83af2" + +[[package]] +name = "futures" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9a31d2a3fbaaeb2af2368bbdd904aa8e812d3c04a1ee10d3171f52d556e5d0a3" +dependencies = [ + "futures-channel", + "futures-core", + "futures-executor", + "futures-io", + "futures-sink", + "futures-task", + "futures-util", +] + +[[package]] +name = "futures-channel" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b1f9e3d69d39e4862ffed03ed071a76f9a13ba1d9109d355b0f0aa6b15e393c4" +dependencies = [ + "futures-core", + "futures-sink", +] + +[[package]] +name = "futures-core" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "92d699e522242e69e3003b94ecc1f960f3a5e015aa7c5d7486e65ad01dd94f5e" + +[[package]] +name = "futures-executor" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "031b47cf1a3c6cc8bc2fc76cd437f521619387907d469316e7c0bc278f1f5432" +dependencies = [ + "futures-core", + "futures-task", + "futures-util", +] + +[[package]] +name = "futures-io" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "53c0fa8157de1303bfffdaa1cc2a673bfffb60102f76b0ef4441659124373fed" + +[[package]] +name = "futures-macro" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9fb9654ba8355388abeb8dcb4fc62f511300867002afc858860463bdd9fe0c44" +dependencies = [ + "proc-macro2", + "quote", + "syn 3.0.9", +] + +[[package]] +name = "futures-sink" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1944426bf7d03f1d14f708785e4b33efd750b36d48a157b836b3efc15ede8e1d" + +[[package]] +name = "futures-task" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cd417de3d1d015fc3bfd2b1ea46dfc7bab72ef86f1cc7cc9c78e728b34a6d1fd" + +[[package]] +name = "futures-util" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0d50a92467f8ba5dd6e3ee5d4bd04d73ab2e4e1c44474a0674821dfce14b79bc" +dependencies = [ + "futures-channel", + "futures-core", + "futures-io", + "futures-macro", + "futures-sink", + "futures-task", + "memchr", + "pin-project-lite", + "slab", +] + +[[package]] +name = "generic-array" +version = "0.14.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "85649ca51fd72272d7821adaf274ad91c288277713d9c18820d8499a7ff69e9a" +dependencies = [ + "typenum", + "version_check", +] + +[[package]] +name = "glob" +version = "0.3.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e4eba85ea1d0a966a983acd07deee566e67395d2d96b6fb39e62b5a833f1eb0b" + +[[package]] +name = "half" +version = "2.7.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" +dependencies = [ + "cfg-if", + "crunchy", + "zerocopy", +] + +[[package]] +name = "hashbrown" +version = "0.14.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e5274423e17b7c9fc20b6e7e208532f9b19825d82dfd615708b70edd83df41f1" + +[[package]] +name = "hashbrown" +version = "0.15.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9229cfe53dfd69f0609a49f65461bd93001ea1ef889cd5529dd176593f5338a1" +dependencies = [ + "foldhash", +] + +[[package]] +name = "hashbrown" +version = "0.17.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ed5909b6e89a2db4456e54cd5f673791d7eca6732202bbf2a9cc504fe2f9b84a" + +[[package]] +name = "home" +version = "0.5.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cc627f471c528ff0c4a49e1d5e60450c8f6461dd6d10ba9dcd3a61d3dff7728d" +dependencies = [ + "windows-sys 0.61.2", +] + +[[package]] +name = "indexmap" +version = "2.14.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cc4e190f5d26ca7051642629da2c52fc03bde85a03197c99408dcd291734c855" +dependencies = [ + "equivalent", + "hashbrown 0.17.1", +] + +[[package]] +name = "itertools" +version = "0.12.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ba291022dbbd398a455acf126c1e341954079855bc60dfdda641363bd6922569" +dependencies = [ + "either", +] + +[[package]] +name = "itoa" +version = "1.0.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8f42a60cbdf9a97f5d2305f08a87dc4e09308d1276d28c869c684d7777685682" + +[[package]] +name = "lazy_static" +version = "1.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "20870f649af7073d53e38067b2a84312175d56ea15217e1b15bc83506ec50afb" + +[[package]] +name = "lazycell" +version = "1.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "830d08ce1d1d941e6b30645f1a0eb5643013d835ce3779a5fc208261dbe10f55" + +[[package]] +name = "libc" +version = "0.2.191" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "65cd87c9a6d60536580b316ad985726104730a5151c99845f2490428ef91240a" + +[[package]] +name = "libloading" +version = "0.8.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d7c4b02199fee7c5d21a5ae7d8cfa79a6ef5bb2fc834d6e9058e89c825efdc55" +dependencies = [ + "cfg-if", + "windows-link", +] + +[[package]] +name = "libnvvm-sys" +version = "0.2.1" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "cuda-target-spec", + "libloading", + "thiserror 2.0.21", +] + +[[package]] +name = "linux-raw-sys" +version = "0.4.15" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d26c52dbd32dccf2d10cac7725f8eae5296885fb5703b261f7d0a0739ec807ab" + +[[package]] +name = "lock_api" +version = "0.4.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "224399e74b87b5f3557511d98dff8b14089b3dadafcab6bb93eab67d3aace965" +dependencies = [ + "scopeguard", +] + +[[package]] +name = "log" +version = "0.4.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9f8bd3e56ce4dfc153cf470fffbfa98c7620958b312ca5c3a4b8d5181fd13c6" + +[[package]] +name = "memchr" +version = "2.8.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cf8baf1c55e62ffcace7a9f06f4bd9cd3f0c4beb022d3b367256b91b87513d98" + +[[package]] +name = "minimal-lexical" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "68354c5c6bd36d73ff3feceb05efa59b6acb7626617f4962be322a825e61f79a" + +[[package]] +name = "nom" +version = "7.1.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d273983c5a657a70a3e8f2a01329822f3b8c8172b73826411a55751e404a0a4a" +dependencies = [ + "memchr", + "minimal-lexical", +] + +[[package]] +name = "nvjitlink-sys" +version = "0.2.1" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "libloading", + "thiserror 2.0.21", +] + +[[package]] +name = "object" +version = "0.36.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "62948e14d923ea95ea2c7c86c71013138b66525b86bdc08d2dcc262bdb497b87" +dependencies = [ + "crc32fast", + "hashbrown 0.15.5", + "indexmap", + "memchr", +] + +[[package]] +name = "once_cell" +version = "1.21.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9f7c3e4beb33f85d45ae3e3a1792185706c8e16d043238c593331cc7cd313b50" + +[[package]] +name = "oxide-artifacts" +version = "0.2.1" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "object", +] + +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "serde", + "serde_json", + "thiserror 2.0.21", +] + +[[package]] +name = "parking_lot_core" +version = "0.9.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2621685985a2ebf1c516881c026032ac7deafcda1a2c9b7850dc81e3dfcb64c1" +dependencies = [ + "cfg-if", + "libc", + "redox_syscall", + "smallvec", + "windows-link", +] + +[[package]] +name = "pin-project-lite" +version = "0.2.17" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a89322df9ebe1c1578d689c92318e070967d1042b512afbe49518723f4e6d5cd" + +[[package]] +name = "prettyplease" +version = "0.2.37" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "479ca8adacdd7ce8f1fb39ce9ecccbfe93a3f1344b3d0d97f20bc0196208f62b" +dependencies = [ + "proc-macro2", + "syn 2.0.119", +] + +[[package]] +name = "proc-macro2" +version = "1.0.107" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "985e7ec9bb745e6ce6535b544d84d6cd6f7ad8bd711c398938ae983b91a766d9" +dependencies = [ + "unicode-ident", +] + +[[package]] +name = "ptx-parse" +version = "0.1.0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" + +[[package]] +name = "quote" +version = "1.0.47" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1fbf4db142a473a8d80c26bbf18454ed458bf8d26c8219c331daecfdbd079001" +dependencies = [ + "proc-macro2", +] + +[[package]] +name = "redox_syscall" +version = "0.5.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ed2bf2547551a7053d6fdfafda3f938979645c44812fbfcda098faae3f1a362d" +dependencies = [ + "bitflags", +] + +[[package]] +name = "regex" +version = "1.13.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f020237b6c8eed93db2e2cb53c00c60a8e1bc73da7d073199a1180401450218d" +dependencies = [ + "aho-corasick", + "memchr", + "regex-automata", + "regex-syntax", +] + +[[package]] +name = "regex-automata" +version = "0.4.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ad8553b9b26413251cbf30e620595c7a41b3887f03da04579c0e6b0d6a06b4b2" +dependencies = [ + "aho-corasick", + "memchr", + "regex-syntax", +] + +[[package]] +name = "regex-syntax" +version = "0.8.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d6f6ff9a378485b298a5286656da665ba74413d36db0979633275d2e708145d4" + +[[package]] +name = "reserved-oxide-symbols" +version = "0.2.1" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" + +[[package]] +name = "rustc-hash" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "08d43f7aa6b08d49f382cde6a7982047c3426db949b1424bc4b7ec9ae12c6ce2" + +[[package]] +name = "rustc-hash" +version = "2.1.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6b1e7f9a428571be2dc5bc0505c13fb6bf936822b894ec87abf8a08a4e51742d" + +[[package]] +name = "rustix" +version = "0.38.44" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fdb5bc1ae2baa591800df16c9ca78619bf65c0488b41b96ccec5d11220d8c154" +dependencies = [ + "bitflags", + "errno", + "libc", + "linux-raw-sys", + "windows-sys 0.59.0", +] + +[[package]] +name = "scopeguard" +version = "1.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "94143f37725109f92c262ed2cf5e59bce7498c01bcc1502d7b9afe439a4e9f49" + +[[package]] +name = "serde" +version = "1.0.229" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4148590afebada386688f18773da617792bf2ef03ffc1e4cbd2b1d45b023e0ba" +dependencies = [ + "serde_core", + "serde_derive", +] + +[[package]] +name = "serde_core" +version = "1.0.229" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "67dca2c9c51e58a4791a4b1ed58308b39c64224d349a935ab5039aa360942a48" +dependencies = [ + "serde_derive", +] + +[[package]] +name = "serde_derive" +version = "1.0.229" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e7a5d71263a5a7d47b41f6b3f06ba276f10cc18b0931f1799f710578e2309348" +dependencies = [ + "proc-macro2", + "quote", + "syn 3.0.9", +] + +[[package]] +name = "serde_json" +version = "1.0.154" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e7e9cc8b1b85264074fbcc02a88680c4096b1e47df8f739dceb03bf482f04bd6" +dependencies = [ + "itoa", + "memchr", + "serde", + "serde_core", + "zmij", +] + +[[package]] +name = "sha2" +version = "0.10.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a7507d819769d01a365ab707794a4084392c824f54a7a6a7862f8c3d0892b283" +dependencies = [ + "cfg-if", + "cpufeatures", + "digest", +] + +[[package]] +name = "shlex" +version = "1.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0fda2ff0d084019ba4d7c6f371c95d8fd75ce3524c3cb8fb653a3023f6323e64" + +[[package]] +name = "slab" +version = "0.4.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0c790de23124f9ab44544d7ac05d60440adc586479ce501c1d6d7da3cd8c9cf5" + +[[package]] +name = "smallvec" +version = "1.16.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5b3dc8af474f516a851ff4bd12db780f948b9250ad37211e4eec0bccea54e01b" + +[[package]] +name = "syn" +version = "2.0.119" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "872831b642d1a07999a962a351ed35b955ea2cfc8f3862091e2a240a84f17297" +dependencies = [ + "proc-macro2", + "quote", + "unicode-ident", +] + +[[package]] +name = "syn" +version = "3.0.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d78c8dee4c7bf0e14673097256fed6142ce9d3b85a408189d07482442145823b" +dependencies = [ + "proc-macro2", + "quote", + "unicode-ident", +] + +[[package]] +name = "thiserror" +version = "1.0.69" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b6aaf5339b578ea85b50e080feb250a3e8ae8cfcdff9a461c9ec2904bc923f52" +dependencies = [ + "thiserror-impl 1.0.69", +] + +[[package]] +name = "thiserror" +version = "2.0.21" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "09e52cb86a36cede5cb101bf8908837b3e4c6e5e59fe7fd85c23fb56200d189e" +dependencies = [ + "thiserror-impl 2.0.21", +] + +[[package]] +name = "thiserror-impl" +version = "1.0.69" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4fee6c4efc90059e10f81e6d42c60a18f76588c3d74cb83a0b242a2b6c7504c1" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", +] + +[[package]] +name = "thiserror-impl" +version = "2.0.21" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fe5197923287db20a58125f0bc85c062f7f2c892de97b18c356f9efb14b28524" +dependencies = [ + "proc-macro2", + "quote", + "syn 3.0.9", +] + +[[package]] +name = "typenum" +version = "1.20.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b6f5e870be6c3b371b77fe0ee0bafb859fa4964b4404c27de1d380043c4dda20" + +[[package]] +name = "unicode-ident" +version = "1.0.27" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a2c754d6c33795a1c324727428e5a7dedb5b06195f9890bdbcba760d3e246563" + +[[package]] +name = "version_check" +version = "0.9.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0b928f33d975fc6ad9f86c8f283853ad26bdd5b10b7f1542aa2fa15e2289105a" + +[[package]] +name = "which" +version = "4.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "87ba24419a2078cd2b0f2ede2691b6c66d8e47836da3b6db8265ebad47afbfc7" +dependencies = [ + "either", + "home", + "once_cell", + "rustix", +] + +[[package]] +name = "windows-link" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f0805222e57f7521d6a62e36fa9163bc891acd422f971defe97d64e70d0a4fe5" + +[[package]] +name = "windows-sys" +version = "0.59.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1e38bc4d79ed67fd075bcc251a1c39b32a1776bbe92e5bef1f0bf1f8c531853b" +dependencies = [ + "windows-targets", +] + +[[package]] +name = "windows-sys" +version = "0.61.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ae137229bcbd6cdf0f7b80a31df61766145077ddf49416a728b02cb3921ff3fc" +dependencies = [ + "windows-link", +] + +[[package]] +name = "windows-targets" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9b724f72796e036ab90c1021d4780d4d3d648aca59e491e6b98e725b84e99973" +dependencies = [ + "windows_aarch64_gnullvm", + "windows_aarch64_msvc", + "windows_i686_gnu", + "windows_i686_gnullvm", + "windows_i686_msvc", + "windows_x86_64_gnu", + "windows_x86_64_gnullvm", + "windows_x86_64_msvc", +] + +[[package]] +name = "windows_aarch64_gnullvm" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "32a4622180e7a0ec044bb555404c800bc9fd9ec262ec147edd5989ccd0c02cd3" + +[[package]] +name = "windows_aarch64_msvc" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "09ec2a7bb152e2252b53fa7803150007879548bc709c039df7627cabbd05d469" + +[[package]] +name = "windows_i686_gnu" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8e9b5ad5ab802e97eb8e295ac6720e509ee4c243f69d781394014ebfe8bbfa0b" + +[[package]] +name = "windows_i686_gnullvm" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0eee52d38c090b3caa76c563b86c3a4bd71ef1a819287c19d586d7334ae8ed66" + +[[package]] +name = "windows_i686_msvc" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "240948bc05c5e7c6dabba28bf89d89ffce3e303022809e73deaefe4f6ec56c66" + +[[package]] +name = "windows_x86_64_gnu" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "147a5c80aabfbf0c7d901cb5895d1de30ef2907eb21fbbab29ca94c5b08b1a78" + +[[package]] +name = "windows_x86_64_gnullvm" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "24d5b23dc417412679681396f2b49f3de8c1473deb516bd34410872eff51ed0d" + +[[package]] +name = "windows_x86_64_msvc" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "589f6da84c646204747d1270a2a5661ea66ed1cced2631d546fdfb155959f9ec" + +[[package]] +name = "zerocopy" +version = "0.8.63" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e5fe1f8f1b06191a00962174c61aa5005e0bb391a6d80d07e24d115c01a92ed8" +dependencies = [ + "zerocopy-derive", +] + +[[package]] +name = "zerocopy-derive" +version = "0.8.63" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "863ad3ac83293fb4d740aedbfdc9240dd8d1a50c1099acd76ce80ce7c7230c7f" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", +] + +[[package]] +name = "zmij" +version = "1.0.23" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "29666d0abbfad1e3dc4dcf6144730dd3a3ab225bbbdac83319345b1b44ccfc1b" diff --git a/rust/cuda-coop/examples/load_store/Cargo.toml b/rust/cuda-coop/examples/load_store/Cargo.toml new file mode 100644 index 000000000000..45b0a1df3f9b --- /dev/null +++ b/rust/cuda-coop/examples/load_store/Cargo.toml @@ -0,0 +1,19 @@ +# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. +# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +[package] +name = "coop-load-store" +version = "0.0.0" +edition = "2024" +publish = false + +[workspace] +members = ["helper"] + +[dependencies] +cuda-coop = { path = "../.." } +coop-load-store-helper = { path = "helper" } +cuda-device = { git = "https://github.com/tpn/cuda-rust.git", rev = "96caed46c776eb49226137eac849e795b83b7ce0" } +cuda-host = { git = "https://github.com/tpn/cuda-rust.git", rev = "96caed46c776eb49226137eac849e795b83b7ce0", features = ["async"] } +cuda-core = { git = "https://github.com/tpn/cuda-rust.git", rev = "96caed46c776eb49226137eac849e795b83b7ce0" } +cuda-async = { git = "https://github.com/tpn/cuda-rust.git", rev = "96caed46c776eb49226137eac849e795b83b7ce0" } diff --git a/rust/cuda-coop/examples/load_store/helper/Cargo.toml b/rust/cuda-coop/examples/load_store/helper/Cargo.toml new file mode 100644 index 000000000000..1c6b59094851 --- /dev/null +++ b/rust/cuda-coop/examples/load_store/helper/Cargo.toml @@ -0,0 +1,11 @@ +# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. +# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +[package] +name = "coop-load-store-helper" +version = "0.0.0" +edition = "2024" +publish = false + +[dependencies] +cuda-coop = { path = "../../.." } diff --git a/rust/cuda-coop/examples/load_store/helper/src/lib.rs b/rust/cuda-coop/examples/load_store/helper/src/lib.rs new file mode 100644 index 000000000000..76d573ce03c6 --- /dev/null +++ b/rust/cuda-coop/examples/load_store/helper/src/lib.rs @@ -0,0 +1,116 @@ +// Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +//! A reusable generic helper shared by downstream CUDA kernels. +#![no_std] +use cuda_coop as coop; + +/// Copy a partial block tile, preserving the source's element width. +pub fn copy( + block: coop::ThreadGroup, + input: &[T], + output: &mut coop::DisjointSlice<'_, T, coop::CollectiveTiles>, + offset: usize, + valid: usize, + default: T, +) { + copy_with_algorithm::( + block, + input, + output, + offset, + valid, + default, + coop::Direct, + ); +} + +/// Copy a partial tile using a compile-time algorithm chosen by the caller. +pub fn copy_with_algorithm( + block: coop::ThreadGroup, + input: &[T], + output: &mut coop::DisjointSlice<'_, T, coop::CollectiveTiles>, + offset: usize, + valid: usize, + default: T, + algorithm: A, +) { + let mut values = coop::ThreadData::<_, ITEMS_PER_THREAD>::new(); + // Full calls let Vectorize use wide accesses for complete tiles. Its counted + // calls retain CUB's guarded Direct path for the final partial tile. + if valid == block.count() as usize * ITEMS_PER_THREAD { + coop::load( + block, + input, + &mut values, + coop::LoadOptions::from(algorithm).offset(offset), + ); + coop::store( + block, + output, + &mut values, + coop::StoreOptions::from(algorithm).offset(offset), + ); + } else { + coop::load( + block, + input, + &mut values, + coop::LoadOptions::default() + .valid_items(valid) + .oob_default(default) + .offset(offset) + .algorithm(algorithm), + ); + coop::store( + block, + output, + &mut values, + coop::StoreOptions::default() + .valid_items(valid) + .offset(offset) + .algorithm(algorithm), + ); + } +} + +/// Copy a tile using caller-selected explicit storage and synchronization policy. +pub fn copy_with_storage< + T: coop::Numeric, + const ITEMS_PER_THREAD: usize, + const C: usize, + const ALIGN: usize, + const AUTO_SYNC: bool, + const EXCLUSIVE: bool, +>( + block: coop::ThreadGroup, + input: &[T], + output: &mut coop::DisjointSlice<'_, T, coop::CollectiveTiles>, + offset: usize, + valid: usize, + default: T, + scratch: &mut coop::TempStorage, +) { + let mut values = coop::ThreadData::<_, ITEMS_PER_THREAD>::new(); + coop::load( + block, + input, + &mut values, + coop::LoadOptions::default() + .valid_items(valid) + .oob_default(default) + .offset(offset) + .algorithm(coop::Transpose) + .temp_storage(scratch), + ); + coop::store( + block, + output, + &mut values, + coop::StoreOptions::default() + .valid_items(valid) + .offset(offset) + .algorithm(coop::Transpose) + .temp_storage(scratch), + ); +} diff --git a/rust/cuda-coop/examples/load_store/src/main.rs b/rust/cuda-coop/examples/load_store/src/main.rs new file mode 100644 index 000000000000..3255aeecde4c --- /dev/null +++ b/rust/cuda-coop/examples/load_store/src/main.rs @@ -0,0 +1,244 @@ +// Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +use cuda_coop as coop; +use cuda_device::{kernel, launch_contract}; +use cuda_host::cuda_module; + +#[cuda_module] +mod kernels { + use super::*; + + #[kernel] + #[launch_contract(domain = 1, block = (128, 1, 1), dynamic_shared = inferred)] + pub fn copy( + input: &[i16], + mut output: coop::DisjointSlice>, + ) { + let block = coop::this_block(); + let tile = block.count() as usize * ITEMS_PER_THREAD; + let offset = block.rank_at::() as usize * tile; + let remaining = if offset <= input.len() { + input.len() - offset + } else { + 0 + }; + let valid = remaining.min(tile); + coop_load_store_helper::copy::<_, ITEMS_PER_THREAD>( + block, + input, + &mut output, + offset, + valid, + -7_i16, + ); + } + #[kernel] + #[launch_contract(domain = 1, block = (128, 1, 1), dynamic_shared = inferred)] + pub fn transpose( + input: &[i16], + mut output: coop::DisjointSlice>, + ) { + let block = coop::this_block(); + let tile = block.count() as usize * ITEMS_PER_THREAD; + let offset = block.rank_at::() as usize * tile; + let remaining = if offset <= input.len() { + input.len() - offset + } else { + 0 + }; + let valid = remaining.min(tile); + coop_load_store_helper::copy_with_algorithm::<_, ITEMS_PER_THREAD, _>( + block, + input, + &mut output, + offset, + valid, + -7_i16, + coop::Transpose, + ); + } + #[kernel] + #[launch_contract(domain = 1, block = (128, 1, 1), dynamic_shared = inferred)] + pub fn striped( + input: &[i16], + mut output: coop::DisjointSlice>, + ) { + let block = coop::this_block(); + let tile = block.count() as usize * ITEMS_PER_THREAD; + let offset = block.rank_at::() as usize * tile; + let remaining = if offset <= input.len() { + input.len() - offset + } else { + 0 + }; + let valid = remaining.min(tile); + coop_load_store_helper::copy_with_algorithm::<_, ITEMS_PER_THREAD, _>( + block, + input, + &mut output, + offset, + valid, + -7_i16, + coop::Striped, + ); + } + #[kernel] + #[launch_contract(domain = 1, block = (128, 1, 1), dynamic_shared = inferred)] + pub fn vectorize( + input: &[i16], + mut output: coop::DisjointSlice>, + ) { + let block = coop::this_block(); + let tile = block.count() as usize * ITEMS_PER_THREAD; + let offset = block.rank_at::() as usize * tile; + let remaining = if offset <= input.len() { + input.len() - offset + } else { + 0 + }; + let valid = remaining.min(tile); + coop_load_store_helper::copy_with_algorithm::<_, ITEMS_PER_THREAD, _>( + block, + input, + &mut output, + offset, + valid, + -7_i16, + coop::Vectorize, + ); + } + #[kernel] + #[launch_contract(domain = 1, block = (128, 1, 1), dynamic_shared = inferred)] + pub fn warp_transpose( + input: &[i16], + mut output: coop::DisjointSlice>, + ) { + let block = coop::this_block(); + let tile = block.count() as usize * ITEMS_PER_THREAD; + let offset = block.rank_at::() as usize * tile; + let remaining = if offset <= input.len() { + input.len() - offset + } else { + 0 + }; + let valid = remaining.min(tile); + coop_load_store_helper::copy_with_algorithm::<_, ITEMS_PER_THREAD, _>( + block, + input, + &mut output, + offset, + valid, + -7_i16, + coop::WarpTranspose, + ); + } + #[kernel] + #[launch_contract(domain = 1, block = (128, 1, 1), dynamic_shared = inferred)] + pub fn warp_transpose_timesliced( + input: &[i16], + mut output: coop::DisjointSlice>, + ) { + let block = coop::this_block(); + let tile = block.count() as usize * ITEMS_PER_THREAD; + let offset = block.rank_at::() as usize * tile; + let remaining = if offset <= input.len() { + input.len() - offset + } else { + 0 + }; + let valid = remaining.min(tile); + coop_load_store_helper::copy_with_algorithm::<_, ITEMS_PER_THREAD, _>( + block, + input, + &mut output, + offset, + valid, + -7_i16, + coop::WarpTransposeTimesliced, + ); + } + #[kernel] + #[launch_contract(domain = 1, block = (128, 1, 1), dynamic_shared = inferred)] + pub fn explicit_storage( + input: &[i16], + mut output: coop::DisjointSlice>, + ) { + let block = coop::this_block(); + let tile = block.count() as usize * ITEMS_PER_THREAD; + let offset = block.rank_at::() as usize * tile; + let remaining = if offset <= input.len() { + input.len() - offset + } else { + 0 + }; + let valid = remaining.min(tile); + let mut scratch = coop::TempStorage::new().auto_sync(); + let mut iteration = 0; + while iteration < 2 { + coop_load_store_helper::copy_with_storage( + block, + input, + &mut output, + offset, + valid, + -7_i16, + &mut scratch, + ); + iteration += 1; + } + } +} + +fn main() -> Result<(), Box> { + use cuda_async::simt::device_context::{init_device_contexts, with_cuda_context}; + use cuda_async::simt::device_operation::DeviceOperation; + use cuda_core::{DeviceBuffer, LaunchConfig1D}; + + init_device_contexts(0, 1)?; + let context = with_cuda_context(0, |context| context.clone())?; + let stream = context.default_stream(); + // SAFETY: these typed declarations and embedded kernels compile together. + let module = unsafe { kernels::load_async(0)? }; + macro_rules! run { + ($items:literal, $prepare:ident, $launch:ident) => {{ + let input: Vec = (0..3 * 128 * $items - 7) + .map(|i| (i % 251) as i16) + .collect(); + let source = DeviceBuffer::from_host(&stream, &input)?; + let target = DeviceBuffer::from_host(&stream, &vec![-1_i16; input.len()])?; + let plan = module.$prepare::<$items>(LaunchConfig1D::auto_shared(3, 128))?; + let (_, target) = module + .$launch::<$items, _, _>(&plan, source, target) + .sync()?; + let mut output = vec![0_i16; input.len()]; + target.copy_to_host(&stream, &mut output)?; + stream.synchronize()?; + assert_eq!(output, input); + println!("items/thread={}: {} values copied", $items, output.len()); + }}; + } + run!(1, prepare_copy, copy_async_owned); + run!(4, prepare_copy, copy_async_owned); + run!(1, prepare_striped, striped_async_owned); + run!(4, prepare_striped, striped_async_owned); + run!(1, prepare_vectorize, vectorize_async_owned); + run!(4, prepare_vectorize, vectorize_async_owned); + run!(1, prepare_transpose, transpose_async_owned); + run!(4, prepare_transpose, transpose_async_owned); + run!(1, prepare_warp_transpose, warp_transpose_async_owned); + run!(4, prepare_warp_transpose, warp_transpose_async_owned); + run!( + 1, + prepare_warp_transpose_timesliced, + warp_transpose_timesliced_async_owned + ); + run!( + 4, + prepare_warp_transpose_timesliced, + warp_transpose_timesliced_async_owned + ); + run!(1, prepare_explicit_storage, explicit_storage_async_owned); + run!(4, prepare_explicit_storage, explicit_storage_async_owned); + Ok(()) +} diff --git a/rust/cuda-coop/examples/payload.rs b/rust/cuda-coop/examples/payload.rs new file mode 100644 index 000000000000..a7625da6c1ba --- /dev/null +++ b/rust/cuda-coop/examples/payload.rs @@ -0,0 +1,22 @@ +// Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +//! Construct and inspect a per-thread payload on the host. + +use cuda_coop as coop; + +fn payload(value: i16) -> coop::ThreadData { + let mut items = coop::ThreadData::<_, ITEMS_PER_THREAD>::new(); + items.write(0, value); + items +} + +fn main() { + let mut items = payload::<4>(7); + assert_eq!(items.get(0), Some(7_i16)); + assert_eq!(items.get(1), None); + items[0] += 2; + items.write(3, 15); + assert_eq!(items.get(0), Some(9)); + assert_eq!(items.get(3), Some(15)); +} diff --git a/rust/cuda-coop/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..28f967b5bf14 --- /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=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "libnvvm-sys", + "nvjitlink-sys", + "serde", + "sha2", + "thiserror 2.0.21", +] + +[[package]] +name = "cuda-async" +version = "0.5.0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "anyhow", + "cuda-bindings", + "cuda-core", + "dashmap", + "futures", + "half", + "once_cell", + "rustc-hash 2.1.3", + "thiserror 1.0.69", +] + +[[package]] +name = "cuda-bindings" +version = "0.5.0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "bindgen", + "libloading", + "prettyplease", + "proc-macro2", + "quote", + "syn 2.0.119", +] + +[[package]] +name = "cuda-coop" +version = "0.0.0" +dependencies = [ + "cuda-coop-codegen", + "cuda-device", +] + +[[package]] +name = "cuda-coop-codegen" +version = "0.0.0" +dependencies = [ + "libloading", + "serde", + "serde_json", + "sha2", +] + +[[package]] +name = "cuda-core" +version = "0.5.0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "anyhow", + "cuda-bindings", + "cuda-core-derive", + "half", + "oxide-artifacts", +] + +[[package]] +name = "cuda-core-derive" +version = "0.5.0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", +] + +[[package]] +name = "cuda-device" +version = "0.2.1" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "cuda-macros", +] + +[[package]] +name = "cuda-host" +version = "0.2.1" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "cuda-artifact-finalizer", + "cuda-async", + "cuda-core", + "cuda-device", + "cuda-macros", + "half", + "libc", + "oxide-artifacts", + "ptx-parse", + "sha2", + "thiserror 2.0.21", +] + +[[package]] +name = "cuda-macros" +version = "0.2.1" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "oxide-tile-export", + "proc-macro2", + "quote", + "reserved-oxide-symbols", + "serde", + "serde_json", + "syn 2.0.119", +] + +[[package]] +name = "cuda-target-spec" +version = "0.1.0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" + +[[package]] +name = "dashmap" +version = "6.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e6361d5c062261c78a176addb82d4c821ae42bed6089de0e12603cd25de2059c" +dependencies = [ + "cfg-if", + "crossbeam-utils", + "hashbrown 0.14.5", + "lock_api", + "once_cell", + "parking_lot_core", +] + +[[package]] +name = "digest" +version = "0.10.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9ed9a281f7bc9b7576e61468ba615a66a5c8cfdff42420a70aa82701a3b1e292" +dependencies = [ + "block-buffer", + "crypto-common", +] + +[[package]] +name = "either" +version = "1.19.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0e9c71c2167ca323c882b99918929403426e2373ea17242ff5653e0d5e1058be" + +[[package]] +name = "equivalent" +version = "1.0.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "00d174d5400e5e8fd687ad1049e2f578285fa914201b1af7e8b112a4546bd826" + +[[package]] +name = "errno" +version = "0.3.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "39cab71617ae0d63f51a36d69f866391735b51691dbda63cf6f96d042b63efeb" +dependencies = [ + "libc", + "windows-sys 0.61.2", +] + +[[package]] +name = "foldhash" +version = "0.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d9c4f5dac5e15c24eb999c26181a6ca40b39fe946cbe4c263c7209467bc83af2" + +[[package]] +name = "futures" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9a31d2a3fbaaeb2af2368bbdd904aa8e812d3c04a1ee10d3171f52d556e5d0a3" +dependencies = [ + "futures-channel", + "futures-core", + "futures-executor", + "futures-io", + "futures-sink", + "futures-task", + "futures-util", +] + +[[package]] +name = "futures-channel" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b1f9e3d69d39e4862ffed03ed071a76f9a13ba1d9109d355b0f0aa6b15e393c4" +dependencies = [ + "futures-core", + "futures-sink", +] + +[[package]] +name = "futures-core" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "92d699e522242e69e3003b94ecc1f960f3a5e015aa7c5d7486e65ad01dd94f5e" + +[[package]] +name = "futures-executor" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "031b47cf1a3c6cc8bc2fc76cd437f521619387907d469316e7c0bc278f1f5432" +dependencies = [ + "futures-core", + "futures-task", + "futures-util", +] + +[[package]] +name = "futures-io" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "53c0fa8157de1303bfffdaa1cc2a673bfffb60102f76b0ef4441659124373fed" + +[[package]] +name = "futures-macro" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9fb9654ba8355388abeb8dcb4fc62f511300867002afc858860463bdd9fe0c44" +dependencies = [ + "proc-macro2", + "quote", + "syn 3.0.9", +] + +[[package]] +name = "futures-sink" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1944426bf7d03f1d14f708785e4b33efd750b36d48a157b836b3efc15ede8e1d" + +[[package]] +name = "futures-task" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cd417de3d1d015fc3bfd2b1ea46dfc7bab72ef86f1cc7cc9c78e728b34a6d1fd" + +[[package]] +name = "futures-util" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0d50a92467f8ba5dd6e3ee5d4bd04d73ab2e4e1c44474a0674821dfce14b79bc" +dependencies = [ + "futures-channel", + "futures-core", + "futures-io", + "futures-macro", + "futures-sink", + "futures-task", + "memchr", + "pin-project-lite", + "slab", +] + +[[package]] +name = "generic-array" +version = "0.14.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "85649ca51fd72272d7821adaf274ad91c288277713d9c18820d8499a7ff69e9a" +dependencies = [ + "typenum", + "version_check", +] + +[[package]] +name = "glob" +version = "0.3.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e4eba85ea1d0a966a983acd07deee566e67395d2d96b6fb39e62b5a833f1eb0b" + +[[package]] +name = "half" +version = "2.7.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" +dependencies = [ + "cfg-if", + "crunchy", + "zerocopy", +] + +[[package]] +name = "hashbrown" +version = "0.14.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e5274423e17b7c9fc20b6e7e208532f9b19825d82dfd615708b70edd83df41f1" + +[[package]] +name = "hashbrown" +version = "0.15.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9229cfe53dfd69f0609a49f65461bd93001ea1ef889cd5529dd176593f5338a1" +dependencies = [ + "foldhash", +] + +[[package]] +name = "hashbrown" +version = "0.17.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ed5909b6e89a2db4456e54cd5f673791d7eca6732202bbf2a9cc504fe2f9b84a" + +[[package]] +name = "home" +version = "0.5.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cc627f471c528ff0c4a49e1d5e60450c8f6461dd6d10ba9dcd3a61d3dff7728d" +dependencies = [ + "windows-sys 0.61.2", +] + +[[package]] +name = "indexmap" +version = "2.14.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cc4e190f5d26ca7051642629da2c52fc03bde85a03197c99408dcd291734c855" +dependencies = [ + "equivalent", + "hashbrown 0.17.1", +] + +[[package]] +name = "itertools" +version = "0.12.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ba291022dbbd398a455acf126c1e341954079855bc60dfdda641363bd6922569" +dependencies = [ + "either", +] + +[[package]] +name = "itoa" +version = "1.0.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8f42a60cbdf9a97f5d2305f08a87dc4e09308d1276d28c869c684d7777685682" + +[[package]] +name = "lazy_static" +version = "1.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "20870f649af7073d53e38067b2a84312175d56ea15217e1b15bc83506ec50afb" + +[[package]] +name = "lazycell" +version = "1.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "830d08ce1d1d941e6b30645f1a0eb5643013d835ce3779a5fc208261dbe10f55" + +[[package]] +name = "libc" +version = "0.2.191" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "65cd87c9a6d60536580b316ad985726104730a5151c99845f2490428ef91240a" + +[[package]] +name = "libloading" +version = "0.8.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d7c4b02199fee7c5d21a5ae7d8cfa79a6ef5bb2fc834d6e9058e89c825efdc55" +dependencies = [ + "cfg-if", + "windows-link", +] + +[[package]] +name = "libnvvm-sys" +version = "0.2.1" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "cuda-target-spec", + "libloading", + "thiserror 2.0.21", +] + +[[package]] +name = "linux-raw-sys" +version = "0.4.15" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d26c52dbd32dccf2d10cac7725f8eae5296885fb5703b261f7d0a0739ec807ab" + +[[package]] +name = "lock_api" +version = "0.4.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "224399e74b87b5f3557511d98dff8b14089b3dadafcab6bb93eab67d3aace965" +dependencies = [ + "scopeguard", +] + +[[package]] +name = "log" +version = "0.4.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9f8bd3e56ce4dfc153cf470fffbfa98c7620958b312ca5c3a4b8d5181fd13c6" + +[[package]] +name = "memchr" +version = "2.8.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cf8baf1c55e62ffcace7a9f06f4bd9cd3f0c4beb022d3b367256b91b87513d98" + +[[package]] +name = "minimal-lexical" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "68354c5c6bd36d73ff3feceb05efa59b6acb7626617f4962be322a825e61f79a" + +[[package]] +name = "nom" +version = "7.1.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d273983c5a657a70a3e8f2a01329822f3b8c8172b73826411a55751e404a0a4a" +dependencies = [ + "memchr", + "minimal-lexical", +] + +[[package]] +name = "nvjitlink-sys" +version = "0.2.1" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "libloading", + "thiserror 2.0.21", +] + +[[package]] +name = "object" +version = "0.36.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "62948e14d923ea95ea2c7c86c71013138b66525b86bdc08d2dcc262bdb497b87" +dependencies = [ + "crc32fast", + "hashbrown 0.15.5", + "indexmap", + "memchr", +] + +[[package]] +name = "once_cell" +version = "1.21.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9f7c3e4beb33f85d45ae3e3a1792185706c8e16d043238c593331cc7cd313b50" + +[[package]] +name = "oxide-artifacts" +version = "0.2.1" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "object", +] + +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "serde", + "serde_json", + "thiserror 2.0.21", +] + +[[package]] +name = "parking_lot_core" +version = "0.9.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2621685985a2ebf1c516881c026032ac7deafcda1a2c9b7850dc81e3dfcb64c1" +dependencies = [ + "cfg-if", + "libc", + "redox_syscall", + "smallvec", + "windows-link", +] + +[[package]] +name = "pin-project-lite" +version = "0.2.17" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a89322df9ebe1c1578d689c92318e070967d1042b512afbe49518723f4e6d5cd" + +[[package]] +name = "prettyplease" +version = "0.2.37" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "479ca8adacdd7ce8f1fb39ce9ecccbfe93a3f1344b3d0d97f20bc0196208f62b" +dependencies = [ + "proc-macro2", + "syn 2.0.119", +] + +[[package]] +name = "proc-macro2" +version = "1.0.107" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "985e7ec9bb745e6ce6535b544d84d6cd6f7ad8bd711c398938ae983b91a766d9" +dependencies = [ + "unicode-ident", +] + +[[package]] +name = "ptx-parse" +version = "0.1.0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" + +[[package]] +name = "quote" +version = "1.0.47" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1fbf4db142a473a8d80c26bbf18454ed458bf8d26c8219c331daecfdbd079001" +dependencies = [ + "proc-macro2", +] + +[[package]] +name = "redox_syscall" +version = "0.5.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ed2bf2547551a7053d6fdfafda3f938979645c44812fbfcda098faae3f1a362d" +dependencies = [ + "bitflags", +] + +[[package]] +name = "regex" +version = "1.13.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f020237b6c8eed93db2e2cb53c00c60a8e1bc73da7d073199a1180401450218d" +dependencies = [ + "aho-corasick", + "memchr", + "regex-automata", + "regex-syntax", +] + +[[package]] +name = "regex-automata" +version = "0.4.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ad8553b9b26413251cbf30e620595c7a41b3887f03da04579c0e6b0d6a06b4b2" +dependencies = [ + "aho-corasick", + "memchr", + "regex-syntax", +] + +[[package]] +name = "regex-syntax" +version = "0.8.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d6f6ff9a378485b298a5286656da665ba74413d36db0979633275d2e708145d4" + +[[package]] +name = "reserved-oxide-symbols" +version = "0.2.1" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" + +[[package]] +name = "rustc-hash" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "08d43f7aa6b08d49f382cde6a7982047c3426db949b1424bc4b7ec9ae12c6ce2" + +[[package]] +name = "rustc-hash" +version = "2.1.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6b1e7f9a428571be2dc5bc0505c13fb6bf936822b894ec87abf8a08a4e51742d" + +[[package]] +name = "rustix" +version = "0.38.44" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fdb5bc1ae2baa591800df16c9ca78619bf65c0488b41b96ccec5d11220d8c154" +dependencies = [ + "bitflags", + "errno", + "libc", + "linux-raw-sys", + "windows-sys 0.59.0", +] + +[[package]] +name = "scopeguard" +version = "1.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "94143f37725109f92c262ed2cf5e59bce7498c01bcc1502d7b9afe439a4e9f49" + +[[package]] +name = "serde" +version = "1.0.229" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4148590afebada386688f18773da617792bf2ef03ffc1e4cbd2b1d45b023e0ba" +dependencies = [ + "serde_core", + "serde_derive", +] + +[[package]] +name = "serde_core" +version = "1.0.229" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "67dca2c9c51e58a4791a4b1ed58308b39c64224d349a935ab5039aa360942a48" +dependencies = [ + "serde_derive", +] + +[[package]] +name = "serde_derive" +version = "1.0.229" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e7a5d71263a5a7d47b41f6b3f06ba276f10cc18b0931f1799f710578e2309348" +dependencies = [ + "proc-macro2", + "quote", + "syn 3.0.9", +] + +[[package]] +name = "serde_json" +version = "1.0.154" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e7e9cc8b1b85264074fbcc02a88680c4096b1e47df8f739dceb03bf482f04bd6" +dependencies = [ + "itoa", + "memchr", + "serde", + "serde_core", + "zmij", +] + +[[package]] +name = "sha2" +version = "0.10.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a7507d819769d01a365ab707794a4084392c824f54a7a6a7862f8c3d0892b283" +dependencies = [ + "cfg-if", + "cpufeatures", + "digest", +] + +[[package]] +name = "shlex" +version = "1.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0fda2ff0d084019ba4d7c6f371c95d8fd75ce3524c3cb8fb653a3023f6323e64" + +[[package]] +name = "slab" +version = "0.4.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0c790de23124f9ab44544d7ac05d60440adc586479ce501c1d6d7da3cd8c9cf5" + +[[package]] +name = "smallvec" +version = "1.16.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5b3dc8af474f516a851ff4bd12db780f948b9250ad37211e4eec0bccea54e01b" + +[[package]] +name = "syn" +version = "2.0.119" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "872831b642d1a07999a962a351ed35b955ea2cfc8f3862091e2a240a84f17297" +dependencies = [ + "proc-macro2", + "quote", + "unicode-ident", +] + +[[package]] +name = "syn" +version = "3.0.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d78c8dee4c7bf0e14673097256fed6142ce9d3b85a408189d07482442145823b" +dependencies = [ + "proc-macro2", + "quote", + "unicode-ident", +] + +[[package]] +name = "thiserror" +version = "1.0.69" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b6aaf5339b578ea85b50e080feb250a3e8ae8cfcdff9a461c9ec2904bc923f52" +dependencies = [ + "thiserror-impl 1.0.69", +] + +[[package]] +name = "thiserror" +version = "2.0.21" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "09e52cb86a36cede5cb101bf8908837b3e4c6e5e59fe7fd85c23fb56200d189e" +dependencies = [ + "thiserror-impl 2.0.21", +] + +[[package]] +name = "thiserror-impl" +version = "1.0.69" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4fee6c4efc90059e10f81e6d42c60a18f76588c3d74cb83a0b242a2b6c7504c1" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", +] + +[[package]] +name = "thiserror-impl" +version = "2.0.21" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fe5197923287db20a58125f0bc85c062f7f2c892de97b18c356f9efb14b28524" +dependencies = [ + "proc-macro2", + "quote", + "syn 3.0.9", +] + +[[package]] +name = "typenum" +version = "1.20.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b6f5e870be6c3b371b77fe0ee0bafb859fa4964b4404c27de1d380043c4dda20" + +[[package]] +name = "unicode-ident" +version = "1.0.27" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a2c754d6c33795a1c324727428e5a7dedb5b06195f9890bdbcba760d3e246563" + +[[package]] +name = "version_check" +version = "0.9.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0b928f33d975fc6ad9f86c8f283853ad26bdd5b10b7f1542aa2fa15e2289105a" + +[[package]] +name = "which" +version = "4.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "87ba24419a2078cd2b0f2ede2691b6c66d8e47836da3b6db8265ebad47afbfc7" +dependencies = [ + "either", + "home", + "once_cell", + "rustix", +] + +[[package]] +name = "windows-link" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f0805222e57f7521d6a62e36fa9163bc891acd422f971defe97d64e70d0a4fe5" + +[[package]] +name = "windows-sys" +version = "0.59.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1e38bc4d79ed67fd075bcc251a1c39b32a1776bbe92e5bef1f0bf1f8c531853b" +dependencies = [ + "windows-targets", +] + +[[package]] +name = "windows-sys" +version = "0.61.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ae137229bcbd6cdf0f7b80a31df61766145077ddf49416a728b02cb3921ff3fc" +dependencies = [ + "windows-link", +] + +[[package]] +name = "windows-targets" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9b724f72796e036ab90c1021d4780d4d3d648aca59e491e6b98e725b84e99973" +dependencies = [ + "windows_aarch64_gnullvm", + "windows_aarch64_msvc", + "windows_i686_gnu", + "windows_i686_gnullvm", + "windows_i686_msvc", + "windows_x86_64_gnu", + "windows_x86_64_gnullvm", + "windows_x86_64_msvc", +] + +[[package]] +name = "windows_aarch64_gnullvm" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "32a4622180e7a0ec044bb555404c800bc9fd9ec262ec147edd5989ccd0c02cd3" + +[[package]] +name = "windows_aarch64_msvc" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "09ec2a7bb152e2252b53fa7803150007879548bc709c039df7627cabbd05d469" + +[[package]] +name = "windows_i686_gnu" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8e9b5ad5ab802e97eb8e295ac6720e509ee4c243f69d781394014ebfe8bbfa0b" + +[[package]] +name = "windows_i686_gnullvm" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0eee52d38c090b3caa76c563b86c3a4bd71ef1a819287c19d586d7334ae8ed66" + +[[package]] +name = "windows_i686_msvc" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "240948bc05c5e7c6dabba28bf89d89ffce3e303022809e73deaefe4f6ec56c66" + +[[package]] +name = "windows_x86_64_gnu" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "147a5c80aabfbf0c7d901cb5895d1de30ef2907eb21fbbab29ca94c5b08b1a78" + +[[package]] +name = "windows_x86_64_gnullvm" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "24d5b23dc417412679681396f2b49f3de8c1473deb516bd34410872eff51ed0d" + +[[package]] +name = "windows_x86_64_msvc" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "589f6da84c646204747d1270a2a5661ea66ed1cced2631d546fdfb155959f9ec" + +[[package]] +name = "zerocopy" +version = "0.8.63" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e5fe1f8f1b06191a00962174c61aa5005e0bb391a6d80d07e24d115c01a92ed8" +dependencies = [ + "zerocopy-derive", +] + +[[package]] +name = "zerocopy-derive" +version = "0.8.63" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "863ad3ac83293fb4d740aedbfdc9240dd8d1a50c1099acd76ce80ce7c7230c7f" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", +] + +[[package]] +name = "zmij" +version = "1.0.23" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "29666d0abbfad1e3dc4dcf6144730dd3a3ab225bbbdac83319345b1b44ccfc1b" diff --git a/rust/cuda-coop/examples/reduce/Cargo.toml b/rust/cuda-coop/examples/reduce/Cargo.toml new file mode 100644 index 000000000000..1e809cebd868 --- /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 = "96caed46c776eb49226137eac849e795b83b7ce0" } +cuda-host = { git = "https://github.com/tpn/cuda-rust.git", rev = "96caed46c776eb49226137eac849e795b83b7ce0", features = ["async"] } +cuda-core = { git = "https://github.com/tpn/cuda-rust.git", rev = "96caed46c776eb49226137eac849e795b83b7ce0" } +cuda-async = { git = "https://github.com/tpn/cuda-rust.git", rev = "96caed46c776eb49226137eac849e795b83b7ce0" } diff --git a/rust/cuda-coop/examples/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..1245575cc9da --- /dev/null +++ b/rust/cuda-coop/examples/reduce/helper/src/lib.rs @@ -0,0 +1,31 @@ +// 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()) +} 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..82f216ad302d --- /dev/null +++ b/rust/cuda-coop/examples/reduce/src/main.rs @@ -0,0 +1,83 @@ +// 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::*; + + #[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 + ); + } + }}; + } + 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..2c0b8f54e656 --- /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=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "libnvvm-sys", + "nvjitlink-sys", + "serde", + "sha2", + "thiserror 2.0.21", +] + +[[package]] +name = "cuda-async" +version = "0.5.0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "anyhow", + "cuda-bindings", + "cuda-core", + "dashmap", + "futures", + "half", + "once_cell", + "rustc-hash 2.1.3", + "thiserror 1.0.69", +] + +[[package]] +name = "cuda-bindings" +version = "0.5.0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "bindgen", + "libloading", + "prettyplease", + "proc-macro2", + "quote", + "syn 2.0.119", +] + +[[package]] +name = "cuda-coop" +version = "0.0.0" +dependencies = [ + "cuda-coop-codegen", + "cuda-device", +] + +[[package]] +name = "cuda-coop-codegen" +version = "0.0.0" +dependencies = [ + "libloading", + "serde", + "serde_json", + "sha2", +] + +[[package]] +name = "cuda-core" +version = "0.5.0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "anyhow", + "cuda-bindings", + "cuda-core-derive", + "half", + "oxide-artifacts", +] + +[[package]] +name = "cuda-core-derive" +version = "0.5.0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", +] + +[[package]] +name = "cuda-device" +version = "0.2.1" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "cuda-macros", +] + +[[package]] +name = "cuda-host" +version = "0.2.1" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "cuda-artifact-finalizer", + "cuda-async", + "cuda-core", + "cuda-device", + "cuda-macros", + "half", + "libc", + "oxide-artifacts", + "ptx-parse", + "sha2", + "thiserror 2.0.21", +] + +[[package]] +name = "cuda-macros" +version = "0.2.1" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "oxide-tile-export", + "proc-macro2", + "quote", + "reserved-oxide-symbols", + "serde", + "serde_json", + "syn 2.0.119", +] + +[[package]] +name = "cuda-target-spec" +version = "0.1.0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" + +[[package]] +name = "dashmap" +version = "6.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e6361d5c062261c78a176addb82d4c821ae42bed6089de0e12603cd25de2059c" +dependencies = [ + "cfg-if", + "crossbeam-utils", + "hashbrown 0.14.5", + "lock_api", + "once_cell", + "parking_lot_core", +] + +[[package]] +name = "digest" +version = "0.10.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9ed9a281f7bc9b7576e61468ba615a66a5c8cfdff42420a70aa82701a3b1e292" +dependencies = [ + "block-buffer", + "crypto-common", +] + +[[package]] +name = "either" +version = "1.19.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0e9c71c2167ca323c882b99918929403426e2373ea17242ff5653e0d5e1058be" + +[[package]] +name = "equivalent" +version = "1.0.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "00d174d5400e5e8fd687ad1049e2f578285fa914201b1af7e8b112a4546bd826" + +[[package]] +name = "errno" +version = "0.3.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "39cab71617ae0d63f51a36d69f866391735b51691dbda63cf6f96d042b63efeb" +dependencies = [ + "libc", + "windows-sys 0.61.2", +] + +[[package]] +name = "foldhash" +version = "0.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d9c4f5dac5e15c24eb999c26181a6ca40b39fe946cbe4c263c7209467bc83af2" + +[[package]] +name = "futures" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9a31d2a3fbaaeb2af2368bbdd904aa8e812d3c04a1ee10d3171f52d556e5d0a3" +dependencies = [ + "futures-channel", + "futures-core", + "futures-executor", + "futures-io", + "futures-sink", + "futures-task", + "futures-util", +] + +[[package]] +name = "futures-channel" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b1f9e3d69d39e4862ffed03ed071a76f9a13ba1d9109d355b0f0aa6b15e393c4" +dependencies = [ + "futures-core", + "futures-sink", +] + +[[package]] +name = "futures-core" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "92d699e522242e69e3003b94ecc1f960f3a5e015aa7c5d7486e65ad01dd94f5e" + +[[package]] +name = "futures-executor" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "031b47cf1a3c6cc8bc2fc76cd437f521619387907d469316e7c0bc278f1f5432" +dependencies = [ + "futures-core", + "futures-task", + "futures-util", +] + +[[package]] +name = "futures-io" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "53c0fa8157de1303bfffdaa1cc2a673bfffb60102f76b0ef4441659124373fed" + +[[package]] +name = "futures-macro" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9fb9654ba8355388abeb8dcb4fc62f511300867002afc858860463bdd9fe0c44" +dependencies = [ + "proc-macro2", + "quote", + "syn 3.0.9", +] + +[[package]] +name = "futures-sink" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1944426bf7d03f1d14f708785e4b33efd750b36d48a157b836b3efc15ede8e1d" + +[[package]] +name = "futures-task" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cd417de3d1d015fc3bfd2b1ea46dfc7bab72ef86f1cc7cc9c78e728b34a6d1fd" + +[[package]] +name = "futures-util" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0d50a92467f8ba5dd6e3ee5d4bd04d73ab2e4e1c44474a0674821dfce14b79bc" +dependencies = [ + "futures-channel", + "futures-core", + "futures-io", + "futures-macro", + "futures-sink", + "futures-task", + "memchr", + "pin-project-lite", + "slab", +] + +[[package]] +name = "generic-array" +version = "0.14.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "85649ca51fd72272d7821adaf274ad91c288277713d9c18820d8499a7ff69e9a" +dependencies = [ + "typenum", + "version_check", +] + +[[package]] +name = "glob" +version = "0.3.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e4eba85ea1d0a966a983acd07deee566e67395d2d96b6fb39e62b5a833f1eb0b" + +[[package]] +name = "half" +version = "2.7.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" +dependencies = [ + "cfg-if", + "crunchy", + "zerocopy", +] + +[[package]] +name = "hashbrown" +version = "0.14.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e5274423e17b7c9fc20b6e7e208532f9b19825d82dfd615708b70edd83df41f1" + +[[package]] +name = "hashbrown" +version = "0.15.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9229cfe53dfd69f0609a49f65461bd93001ea1ef889cd5529dd176593f5338a1" +dependencies = [ + "foldhash", +] + +[[package]] +name = "hashbrown" +version = "0.17.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ed5909b6e89a2db4456e54cd5f673791d7eca6732202bbf2a9cc504fe2f9b84a" + +[[package]] +name = "home" +version = "0.5.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cc627f471c528ff0c4a49e1d5e60450c8f6461dd6d10ba9dcd3a61d3dff7728d" +dependencies = [ + "windows-sys 0.61.2", +] + +[[package]] +name = "indexmap" +version = "2.14.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cc4e190f5d26ca7051642629da2c52fc03bde85a03197c99408dcd291734c855" +dependencies = [ + "equivalent", + "hashbrown 0.17.1", +] + +[[package]] +name = "itertools" +version = "0.12.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ba291022dbbd398a455acf126c1e341954079855bc60dfdda641363bd6922569" +dependencies = [ + "either", +] + +[[package]] +name = "itoa" +version = "1.0.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8f42a60cbdf9a97f5d2305f08a87dc4e09308d1276d28c869c684d7777685682" + +[[package]] +name = "lazy_static" +version = "1.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "20870f649af7073d53e38067b2a84312175d56ea15217e1b15bc83506ec50afb" + +[[package]] +name = "lazycell" +version = "1.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "830d08ce1d1d941e6b30645f1a0eb5643013d835ce3779a5fc208261dbe10f55" + +[[package]] +name = "libc" +version = "0.2.191" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "65cd87c9a6d60536580b316ad985726104730a5151c99845f2490428ef91240a" + +[[package]] +name = "libloading" +version = "0.8.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d7c4b02199fee7c5d21a5ae7d8cfa79a6ef5bb2fc834d6e9058e89c825efdc55" +dependencies = [ + "cfg-if", + "windows-link", +] + +[[package]] +name = "libnvvm-sys" +version = "0.2.1" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "cuda-target-spec", + "libloading", + "thiserror 2.0.21", +] + +[[package]] +name = "linux-raw-sys" +version = "0.4.15" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d26c52dbd32dccf2d10cac7725f8eae5296885fb5703b261f7d0a0739ec807ab" + +[[package]] +name = "lock_api" +version = "0.4.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "224399e74b87b5f3557511d98dff8b14089b3dadafcab6bb93eab67d3aace965" +dependencies = [ + "scopeguard", +] + +[[package]] +name = "log" +version = "0.4.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9f8bd3e56ce4dfc153cf470fffbfa98c7620958b312ca5c3a4b8d5181fd13c6" + +[[package]] +name = "memchr" +version = "2.8.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cf8baf1c55e62ffcace7a9f06f4bd9cd3f0c4beb022d3b367256b91b87513d98" + +[[package]] +name = "minimal-lexical" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "68354c5c6bd36d73ff3feceb05efa59b6acb7626617f4962be322a825e61f79a" + +[[package]] +name = "nom" +version = "7.1.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d273983c5a657a70a3e8f2a01329822f3b8c8172b73826411a55751e404a0a4a" +dependencies = [ + "memchr", + "minimal-lexical", +] + +[[package]] +name = "nvjitlink-sys" +version = "0.2.1" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "libloading", + "thiserror 2.0.21", +] + +[[package]] +name = "object" +version = "0.36.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "62948e14d923ea95ea2c7c86c71013138b66525b86bdc08d2dcc262bdb497b87" +dependencies = [ + "crc32fast", + "hashbrown 0.15.5", + "indexmap", + "memchr", +] + +[[package]] +name = "once_cell" +version = "1.21.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9f7c3e4beb33f85d45ae3e3a1792185706c8e16d043238c593331cc7cd313b50" + +[[package]] +name = "oxide-artifacts" +version = "0.2.1" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "object", +] + +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "serde", + "serde_json", + "thiserror 2.0.21", +] + +[[package]] +name = "parking_lot_core" +version = "0.9.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2621685985a2ebf1c516881c026032ac7deafcda1a2c9b7850dc81e3dfcb64c1" +dependencies = [ + "cfg-if", + "libc", + "redox_syscall", + "smallvec", + "windows-link", +] + +[[package]] +name = "pin-project-lite" +version = "0.2.17" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a89322df9ebe1c1578d689c92318e070967d1042b512afbe49518723f4e6d5cd" + +[[package]] +name = "prettyplease" +version = "0.2.37" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "479ca8adacdd7ce8f1fb39ce9ecccbfe93a3f1344b3d0d97f20bc0196208f62b" +dependencies = [ + "proc-macro2", + "syn 2.0.119", +] + +[[package]] +name = "proc-macro2" +version = "1.0.107" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "985e7ec9bb745e6ce6535b544d84d6cd6f7ad8bd711c398938ae983b91a766d9" +dependencies = [ + "unicode-ident", +] + +[[package]] +name = "ptx-parse" +version = "0.1.0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" + +[[package]] +name = "quote" +version = "1.0.47" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1fbf4db142a473a8d80c26bbf18454ed458bf8d26c8219c331daecfdbd079001" +dependencies = [ + "proc-macro2", +] + +[[package]] +name = "redox_syscall" +version = "0.5.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ed2bf2547551a7053d6fdfafda3f938979645c44812fbfcda098faae3f1a362d" +dependencies = [ + "bitflags", +] + +[[package]] +name = "regex" +version = "1.13.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f020237b6c8eed93db2e2cb53c00c60a8e1bc73da7d073199a1180401450218d" +dependencies = [ + "aho-corasick", + "memchr", + "regex-automata", + "regex-syntax", +] + +[[package]] +name = "regex-automata" +version = "0.4.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ad8553b9b26413251cbf30e620595c7a41b3887f03da04579c0e6b0d6a06b4b2" +dependencies = [ + "aho-corasick", + "memchr", + "regex-syntax", +] + +[[package]] +name = "regex-syntax" +version = "0.8.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d6f6ff9a378485b298a5286656da665ba74413d36db0979633275d2e708145d4" + +[[package]] +name = "reserved-oxide-symbols" +version = "0.2.1" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" + +[[package]] +name = "rustc-hash" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "08d43f7aa6b08d49f382cde6a7982047c3426db949b1424bc4b7ec9ae12c6ce2" + +[[package]] +name = "rustc-hash" +version = "2.1.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6b1e7f9a428571be2dc5bc0505c13fb6bf936822b894ec87abf8a08a4e51742d" + +[[package]] +name = "rustix" +version = "0.38.44" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fdb5bc1ae2baa591800df16c9ca78619bf65c0488b41b96ccec5d11220d8c154" +dependencies = [ + "bitflags", + "errno", + "libc", + "linux-raw-sys", + "windows-sys 0.59.0", +] + +[[package]] +name = "scopeguard" +version = "1.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "94143f37725109f92c262ed2cf5e59bce7498c01bcc1502d7b9afe439a4e9f49" + +[[package]] +name = "serde" +version = "1.0.229" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4148590afebada386688f18773da617792bf2ef03ffc1e4cbd2b1d45b023e0ba" +dependencies = [ + "serde_core", + "serde_derive", +] + +[[package]] +name = "serde_core" +version = "1.0.229" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "67dca2c9c51e58a4791a4b1ed58308b39c64224d349a935ab5039aa360942a48" +dependencies = [ + "serde_derive", +] + +[[package]] +name = "serde_derive" +version = "1.0.229" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e7a5d71263a5a7d47b41f6b3f06ba276f10cc18b0931f1799f710578e2309348" +dependencies = [ + "proc-macro2", + "quote", + "syn 3.0.9", +] + +[[package]] +name = "serde_json" +version = "1.0.154" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e7e9cc8b1b85264074fbcc02a88680c4096b1e47df8f739dceb03bf482f04bd6" +dependencies = [ + "itoa", + "memchr", + "serde", + "serde_core", + "zmij", +] + +[[package]] +name = "sha2" +version = "0.10.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a7507d819769d01a365ab707794a4084392c824f54a7a6a7862f8c3d0892b283" +dependencies = [ + "cfg-if", + "cpufeatures", + "digest", +] + +[[package]] +name = "shlex" +version = "1.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0fda2ff0d084019ba4d7c6f371c95d8fd75ce3524c3cb8fb653a3023f6323e64" + +[[package]] +name = "slab" +version = "0.4.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0c790de23124f9ab44544d7ac05d60440adc586479ce501c1d6d7da3cd8c9cf5" + +[[package]] +name = "smallvec" +version = "1.16.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5b3dc8af474f516a851ff4bd12db780f948b9250ad37211e4eec0bccea54e01b" + +[[package]] +name = "syn" +version = "2.0.119" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "872831b642d1a07999a962a351ed35b955ea2cfc8f3862091e2a240a84f17297" +dependencies = [ + "proc-macro2", + "quote", + "unicode-ident", +] + +[[package]] +name = "syn" +version = "3.0.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d78c8dee4c7bf0e14673097256fed6142ce9d3b85a408189d07482442145823b" +dependencies = [ + "proc-macro2", + "quote", + "unicode-ident", +] + +[[package]] +name = "thiserror" +version = "1.0.69" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b6aaf5339b578ea85b50e080feb250a3e8ae8cfcdff9a461c9ec2904bc923f52" +dependencies = [ + "thiserror-impl 1.0.69", +] + +[[package]] +name = "thiserror" +version = "2.0.21" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "09e52cb86a36cede5cb101bf8908837b3e4c6e5e59fe7fd85c23fb56200d189e" +dependencies = [ + "thiserror-impl 2.0.21", +] + +[[package]] +name = "thiserror-impl" +version = "1.0.69" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4fee6c4efc90059e10f81e6d42c60a18f76588c3d74cb83a0b242a2b6c7504c1" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", +] + +[[package]] +name = "thiserror-impl" +version = "2.0.21" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fe5197923287db20a58125f0bc85c062f7f2c892de97b18c356f9efb14b28524" +dependencies = [ + "proc-macro2", + "quote", + "syn 3.0.9", +] + +[[package]] +name = "typenum" +version = "1.20.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b6f5e870be6c3b371b77fe0ee0bafb859fa4964b4404c27de1d380043c4dda20" + +[[package]] +name = "unicode-ident" +version = "1.0.27" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a2c754d6c33795a1c324727428e5a7dedb5b06195f9890bdbcba760d3e246563" + +[[package]] +name = "version_check" +version = "0.9.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0b928f33d975fc6ad9f86c8f283853ad26bdd5b10b7f1542aa2fa15e2289105a" + +[[package]] +name = "which" +version = "4.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "87ba24419a2078cd2b0f2ede2691b6c66d8e47836da3b6db8265ebad47afbfc7" +dependencies = [ + "either", + "home", + "once_cell", + "rustix", +] + +[[package]] +name = "windows-link" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f0805222e57f7521d6a62e36fa9163bc891acd422f971defe97d64e70d0a4fe5" + +[[package]] +name = "windows-sys" +version = "0.59.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1e38bc4d79ed67fd075bcc251a1c39b32a1776bbe92e5bef1f0bf1f8c531853b" +dependencies = [ + "windows-targets", +] + +[[package]] +name = "windows-sys" +version = "0.61.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ae137229bcbd6cdf0f7b80a31df61766145077ddf49416a728b02cb3921ff3fc" +dependencies = [ + "windows-link", +] + +[[package]] +name = "windows-targets" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9b724f72796e036ab90c1021d4780d4d3d648aca59e491e6b98e725b84e99973" +dependencies = [ + "windows_aarch64_gnullvm", + "windows_aarch64_msvc", + "windows_i686_gnu", + "windows_i686_gnullvm", + "windows_i686_msvc", + "windows_x86_64_gnu", + "windows_x86_64_gnullvm", + "windows_x86_64_msvc", +] + +[[package]] +name = "windows_aarch64_gnullvm" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "32a4622180e7a0ec044bb555404c800bc9fd9ec262ec147edd5989ccd0c02cd3" + +[[package]] +name = "windows_aarch64_msvc" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "09ec2a7bb152e2252b53fa7803150007879548bc709c039df7627cabbd05d469" + +[[package]] +name = "windows_i686_gnu" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8e9b5ad5ab802e97eb8e295ac6720e509ee4c243f69d781394014ebfe8bbfa0b" + +[[package]] +name = "windows_i686_gnullvm" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0eee52d38c090b3caa76c563b86c3a4bd71ef1a819287c19d586d7334ae8ed66" + +[[package]] +name = "windows_i686_msvc" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "240948bc05c5e7c6dabba28bf89d89ffce3e303022809e73deaefe4f6ec56c66" + +[[package]] +name = "windows_x86_64_gnu" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "147a5c80aabfbf0c7d901cb5895d1de30ef2907eb21fbbab29ca94c5b08b1a78" + +[[package]] +name = "windows_x86_64_gnullvm" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "24d5b23dc417412679681396f2b49f3de8c1473deb516bd34410872eff51ed0d" + +[[package]] +name = "windows_x86_64_msvc" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "589f6da84c646204747d1270a2a5661ea66ed1cced2631d546fdfb155959f9ec" + +[[package]] +name = "zerocopy" +version = "0.8.63" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e5fe1f8f1b06191a00962174c61aa5005e0bb391a6d80d07e24d115c01a92ed8" +dependencies = [ + "zerocopy-derive", +] + +[[package]] +name = "zerocopy-derive" +version = "0.8.63" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "863ad3ac83293fb4d740aedbfdc9240dd8d1a50c1099acd76ce80ce7c7230c7f" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", +] + +[[package]] +name = "zmij" +version = "1.0.23" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "29666d0abbfad1e3dc4dcf6144730dd3a3ab225bbbdac83319345b1b44ccfc1b" diff --git a/rust/cuda-coop/examples/reduce_batched/Cargo.toml b/rust/cuda-coop/examples/reduce_batched/Cargo.toml new file mode 100644 index 000000000000..58091c530866 --- /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 = "96caed46c776eb49226137eac849e795b83b7ce0" } +cuda-host = { git = "https://github.com/tpn/cuda-rust.git", rev = "96caed46c776eb49226137eac849e795b83b7ce0", features = ["async"] } +cuda-core = { git = "https://github.com/tpn/cuda-rust.git", rev = "96caed46c776eb49226137eac849e795b83b7ce0" } +cuda-async = { git = "https://github.com/tpn/cuda-rust.git", rev = "96caed46c776eb49226137eac849e795b83b7ce0" } diff --git a/rust/cuda-coop/examples/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..565676cddbda --- /dev/null +++ b/rust/cuda-coop/examples/scan/Cargo.lock @@ -0,0 +1,1019 @@ +# This file is automatically @generated by Cargo. +# It is not intended for manual editing. +version = 4 + +[[package]] +name = "aho-corasick" +version = "1.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c982642fa9e8606056828ee9a8505737230110bb1099153c79efe865c59d12ba" +dependencies = [ + "memchr", +] + +[[package]] +name = "anyhow" +version = "1.0.104" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "330a5ed07fa54e4702c9d6c4174f74427fc0ef6e214bbd677ae50a5099946470" + +[[package]] +name = "bindgen" +version = "0.69.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "271383c67ccabffb7381723dea0672a673f292304fcb45c01cc648c7a8d58088" +dependencies = [ + "bitflags", + "cexpr", + "clang-sys", + "itertools", + "lazy_static", + "lazycell", + "log", + "prettyplease", + "proc-macro2", + "quote", + "regex", + "rustc-hash 1.1.0", + "shlex", + "syn 2.0.119", + "which", +] + +[[package]] +name = "bitflags" +version = "2.13.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3ded4057c258ba199e2d26386d3af3780957ecaee6c4ef4041c6b4b8b97c0b06" + +[[package]] +name = "block-buffer" +version = "0.10.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3078c7629b62d3f0439517fa394996acacc5cbc91c5a20d8c658e77abd503a71" +dependencies = [ + "generic-array", +] + +[[package]] +name = "cexpr" +version = "0.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6fac387a98bb7c37292057cffc56d62ecb629900026402633ae9160df93a8766" +dependencies = [ + "nom", +] + +[[package]] +name = "cfg-if" +version = "1.0.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4e7648175b45a9a48536d676f68d918270699102aa8dab5496df06904c914600" + +[[package]] +name = "clang-sys" +version = "1.9.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "157a8ba7b480713b56f4c09fd13fc3e0a22a5dfab8097ba61cbc5feef950788a" +dependencies = [ + "glob", + "libc", + "libloading", +] + +[[package]] +name = "coop-scan" +version = "0.0.0" +dependencies = [ + "coop-scan-helper", + "cuda-async", + "cuda-coop", + "cuda-core", + "cuda-device", + "cuda-host", +] + +[[package]] +name = "coop-scan-helper" +version = "0.0.0" +dependencies = [ + "cuda-coop", +] + +[[package]] +name = "cpufeatures" +version = "0.2.17" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "59ed5838eebb26a2bb2e58f6d5b5316989ae9d08bab10e0e6d103e656d1b0280" +dependencies = [ + "libc", +] + +[[package]] +name = "crc32fast" +version = "1.5.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0ba106e844045ddc02c20768ac4b0b7cb45f2cc2ecbd833ac9f5dacb183c2055" +dependencies = [ + "cfg-if", +] + +[[package]] +name = "crossbeam-utils" +version = "0.8.23" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a31eee39dddec8330830986fcd7625edb5a24ec90ea038215273bbc3adb08ac6" + +[[package]] +name = "crunchy" +version = "0.2.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "460fbee9c2c2f33933d720630a6a0bac33ba7053db5344fac858d4b8952d77d5" + +[[package]] +name = "crypto-common" +version = "0.1.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "78c8292055d1c1df0cce5d180393dc8cce0abec0a7102adb6c7b1eef6016d60a" +dependencies = [ + "generic-array", + "typenum", +] + +[[package]] +name = "cuda-artifact-finalizer" +version = "0.2.1" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "libnvvm-sys", + "nvjitlink-sys", + "serde", + "sha2", + "thiserror 2.0.21", +] + +[[package]] +name = "cuda-async" +version = "0.5.0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "anyhow", + "cuda-bindings", + "cuda-core", + "dashmap", + "futures", + "half", + "once_cell", + "rustc-hash 2.1.3", + "thiserror 1.0.69", +] + +[[package]] +name = "cuda-bindings" +version = "0.5.0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "bindgen", + "libloading", + "prettyplease", + "proc-macro2", + "quote", + "syn 2.0.119", +] + +[[package]] +name = "cuda-coop" +version = "0.0.0" +dependencies = [ + "cuda-coop-codegen", + "cuda-device", +] + +[[package]] +name = "cuda-coop-codegen" +version = "0.0.0" +dependencies = [ + "libloading", + "serde", + "serde_json", + "sha2", +] + +[[package]] +name = "cuda-core" +version = "0.5.0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "anyhow", + "cuda-bindings", + "cuda-core-derive", + "half", + "oxide-artifacts", +] + +[[package]] +name = "cuda-core-derive" +version = "0.5.0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", +] + +[[package]] +name = "cuda-device" +version = "0.2.1" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "cuda-macros", +] + +[[package]] +name = "cuda-host" +version = "0.2.1" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "cuda-artifact-finalizer", + "cuda-async", + "cuda-core", + "cuda-device", + "cuda-macros", + "half", + "libc", + "oxide-artifacts", + "ptx-parse", + "sha2", + "thiserror 2.0.21", +] + +[[package]] +name = "cuda-macros" +version = "0.2.1" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "oxide-tile-export", + "proc-macro2", + "quote", + "reserved-oxide-symbols", + "serde", + "serde_json", + "syn 2.0.119", +] + +[[package]] +name = "cuda-target-spec" +version = "0.1.0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" + +[[package]] +name = "dashmap" +version = "6.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e6361d5c062261c78a176addb82d4c821ae42bed6089de0e12603cd25de2059c" +dependencies = [ + "cfg-if", + "crossbeam-utils", + "hashbrown 0.14.5", + "lock_api", + "once_cell", + "parking_lot_core", +] + +[[package]] +name = "digest" +version = "0.10.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9ed9a281f7bc9b7576e61468ba615a66a5c8cfdff42420a70aa82701a3b1e292" +dependencies = [ + "block-buffer", + "crypto-common", +] + +[[package]] +name = "either" +version = "1.19.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0e9c71c2167ca323c882b99918929403426e2373ea17242ff5653e0d5e1058be" + +[[package]] +name = "equivalent" +version = "1.0.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "00d174d5400e5e8fd687ad1049e2f578285fa914201b1af7e8b112a4546bd826" + +[[package]] +name = "errno" +version = "0.3.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "39cab71617ae0d63f51a36d69f866391735b51691dbda63cf6f96d042b63efeb" +dependencies = [ + "libc", + "windows-sys 0.61.2", +] + +[[package]] +name = "foldhash" +version = "0.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d9c4f5dac5e15c24eb999c26181a6ca40b39fe946cbe4c263c7209467bc83af2" + +[[package]] +name = "futures" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9a31d2a3fbaaeb2af2368bbdd904aa8e812d3c04a1ee10d3171f52d556e5d0a3" +dependencies = [ + "futures-channel", + "futures-core", + "futures-executor", + "futures-io", + "futures-sink", + "futures-task", + "futures-util", +] + +[[package]] +name = "futures-channel" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b1f9e3d69d39e4862ffed03ed071a76f9a13ba1d9109d355b0f0aa6b15e393c4" +dependencies = [ + "futures-core", + "futures-sink", +] + +[[package]] +name = "futures-core" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "92d699e522242e69e3003b94ecc1f960f3a5e015aa7c5d7486e65ad01dd94f5e" + +[[package]] +name = "futures-executor" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "031b47cf1a3c6cc8bc2fc76cd437f521619387907d469316e7c0bc278f1f5432" +dependencies = [ + "futures-core", + "futures-task", + "futures-util", +] + +[[package]] +name = "futures-io" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "53c0fa8157de1303bfffdaa1cc2a673bfffb60102f76b0ef4441659124373fed" + +[[package]] +name = "futures-macro" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9fb9654ba8355388abeb8dcb4fc62f511300867002afc858860463bdd9fe0c44" +dependencies = [ + "proc-macro2", + "quote", + "syn 3.0.9", +] + +[[package]] +name = "futures-sink" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1944426bf7d03f1d14f708785e4b33efd750b36d48a157b836b3efc15ede8e1d" + +[[package]] +name = "futures-task" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cd417de3d1d015fc3bfd2b1ea46dfc7bab72ef86f1cc7cc9c78e728b34a6d1fd" + +[[package]] +name = "futures-util" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0d50a92467f8ba5dd6e3ee5d4bd04d73ab2e4e1c44474a0674821dfce14b79bc" +dependencies = [ + "futures-channel", + "futures-core", + "futures-io", + "futures-macro", + "futures-sink", + "futures-task", + "memchr", + "pin-project-lite", + "slab", +] + +[[package]] +name = "generic-array" +version = "0.14.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "85649ca51fd72272d7821adaf274ad91c288277713d9c18820d8499a7ff69e9a" +dependencies = [ + "typenum", + "version_check", +] + +[[package]] +name = "glob" +version = "0.3.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e4eba85ea1d0a966a983acd07deee566e67395d2d96b6fb39e62b5a833f1eb0b" + +[[package]] +name = "half" +version = "2.7.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" +dependencies = [ + "cfg-if", + "crunchy", + "zerocopy", +] + +[[package]] +name = "hashbrown" +version = "0.14.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e5274423e17b7c9fc20b6e7e208532f9b19825d82dfd615708b70edd83df41f1" + +[[package]] +name = "hashbrown" +version = "0.15.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9229cfe53dfd69f0609a49f65461bd93001ea1ef889cd5529dd176593f5338a1" +dependencies = [ + "foldhash", +] + +[[package]] +name = "hashbrown" +version = "0.17.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ed5909b6e89a2db4456e54cd5f673791d7eca6732202bbf2a9cc504fe2f9b84a" + +[[package]] +name = "home" +version = "0.5.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cc627f471c528ff0c4a49e1d5e60450c8f6461dd6d10ba9dcd3a61d3dff7728d" +dependencies = [ + "windows-sys 0.61.2", +] + +[[package]] +name = "indexmap" +version = "2.14.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cc4e190f5d26ca7051642629da2c52fc03bde85a03197c99408dcd291734c855" +dependencies = [ + "equivalent", + "hashbrown 0.17.1", +] + +[[package]] +name = "itertools" +version = "0.12.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ba291022dbbd398a455acf126c1e341954079855bc60dfdda641363bd6922569" +dependencies = [ + "either", +] + +[[package]] +name = "itoa" +version = "1.0.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8f42a60cbdf9a97f5d2305f08a87dc4e09308d1276d28c869c684d7777685682" + +[[package]] +name = "lazy_static" +version = "1.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "20870f649af7073d53e38067b2a84312175d56ea15217e1b15bc83506ec50afb" + +[[package]] +name = "lazycell" +version = "1.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "830d08ce1d1d941e6b30645f1a0eb5643013d835ce3779a5fc208261dbe10f55" + +[[package]] +name = "libc" +version = "0.2.191" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "65cd87c9a6d60536580b316ad985726104730a5151c99845f2490428ef91240a" + +[[package]] +name = "libloading" +version = "0.8.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d7c4b02199fee7c5d21a5ae7d8cfa79a6ef5bb2fc834d6e9058e89c825efdc55" +dependencies = [ + "cfg-if", + "windows-link", +] + +[[package]] +name = "libnvvm-sys" +version = "0.2.1" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "cuda-target-spec", + "libloading", + "thiserror 2.0.21", +] + +[[package]] +name = "linux-raw-sys" +version = "0.4.15" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d26c52dbd32dccf2d10cac7725f8eae5296885fb5703b261f7d0a0739ec807ab" + +[[package]] +name = "lock_api" +version = "0.4.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "224399e74b87b5f3557511d98dff8b14089b3dadafcab6bb93eab67d3aace965" +dependencies = [ + "scopeguard", +] + +[[package]] +name = "log" +version = "0.4.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9f8bd3e56ce4dfc153cf470fffbfa98c7620958b312ca5c3a4b8d5181fd13c6" + +[[package]] +name = "memchr" +version = "2.8.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cf8baf1c55e62ffcace7a9f06f4bd9cd3f0c4beb022d3b367256b91b87513d98" + +[[package]] +name = "minimal-lexical" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "68354c5c6bd36d73ff3feceb05efa59b6acb7626617f4962be322a825e61f79a" + +[[package]] +name = "nom" +version = "7.1.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d273983c5a657a70a3e8f2a01329822f3b8c8172b73826411a55751e404a0a4a" +dependencies = [ + "memchr", + "minimal-lexical", +] + +[[package]] +name = "nvjitlink-sys" +version = "0.2.1" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "libloading", + "thiserror 2.0.21", +] + +[[package]] +name = "object" +version = "0.36.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "62948e14d923ea95ea2c7c86c71013138b66525b86bdc08d2dcc262bdb497b87" +dependencies = [ + "crc32fast", + "hashbrown 0.15.5", + "indexmap", + "memchr", +] + +[[package]] +name = "once_cell" +version = "1.21.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9f7c3e4beb33f85d45ae3e3a1792185706c8e16d043238c593331cc7cd313b50" + +[[package]] +name = "oxide-artifacts" +version = "0.2.1" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "object", +] + +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "serde", + "serde_json", + "thiserror 2.0.21", +] + +[[package]] +name = "parking_lot_core" +version = "0.9.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2621685985a2ebf1c516881c026032ac7deafcda1a2c9b7850dc81e3dfcb64c1" +dependencies = [ + "cfg-if", + "libc", + "redox_syscall", + "smallvec", + "windows-link", +] + +[[package]] +name = "pin-project-lite" +version = "0.2.17" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a89322df9ebe1c1578d689c92318e070967d1042b512afbe49518723f4e6d5cd" + +[[package]] +name = "prettyplease" +version = "0.2.37" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "479ca8adacdd7ce8f1fb39ce9ecccbfe93a3f1344b3d0d97f20bc0196208f62b" +dependencies = [ + "proc-macro2", + "syn 2.0.119", +] + +[[package]] +name = "proc-macro2" +version = "1.0.107" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "985e7ec9bb745e6ce6535b544d84d6cd6f7ad8bd711c398938ae983b91a766d9" +dependencies = [ + "unicode-ident", +] + +[[package]] +name = "ptx-parse" +version = "0.1.0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" + +[[package]] +name = "quote" +version = "1.0.47" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1fbf4db142a473a8d80c26bbf18454ed458bf8d26c8219c331daecfdbd079001" +dependencies = [ + "proc-macro2", +] + +[[package]] +name = "redox_syscall" +version = "0.5.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ed2bf2547551a7053d6fdfafda3f938979645c44812fbfcda098faae3f1a362d" +dependencies = [ + "bitflags", +] + +[[package]] +name = "regex" +version = "1.13.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f020237b6c8eed93db2e2cb53c00c60a8e1bc73da7d073199a1180401450218d" +dependencies = [ + "aho-corasick", + "memchr", + "regex-automata", + "regex-syntax", +] + +[[package]] +name = "regex-automata" +version = "0.4.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ad8553b9b26413251cbf30e620595c7a41b3887f03da04579c0e6b0d6a06b4b2" +dependencies = [ + "aho-corasick", + "memchr", + "regex-syntax", +] + +[[package]] +name = "regex-syntax" +version = "0.8.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d6f6ff9a378485b298a5286656da665ba74413d36db0979633275d2e708145d4" + +[[package]] +name = "reserved-oxide-symbols" +version = "0.2.1" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" + +[[package]] +name = "rustc-hash" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "08d43f7aa6b08d49f382cde6a7982047c3426db949b1424bc4b7ec9ae12c6ce2" + +[[package]] +name = "rustc-hash" +version = "2.1.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6b1e7f9a428571be2dc5bc0505c13fb6bf936822b894ec87abf8a08a4e51742d" + +[[package]] +name = "rustix" +version = "0.38.44" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fdb5bc1ae2baa591800df16c9ca78619bf65c0488b41b96ccec5d11220d8c154" +dependencies = [ + "bitflags", + "errno", + "libc", + "linux-raw-sys", + "windows-sys 0.59.0", +] + +[[package]] +name = "scopeguard" +version = "1.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "94143f37725109f92c262ed2cf5e59bce7498c01bcc1502d7b9afe439a4e9f49" + +[[package]] +name = "serde" +version = "1.0.229" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4148590afebada386688f18773da617792bf2ef03ffc1e4cbd2b1d45b023e0ba" +dependencies = [ + "serde_core", + "serde_derive", +] + +[[package]] +name = "serde_core" +version = "1.0.229" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "67dca2c9c51e58a4791a4b1ed58308b39c64224d349a935ab5039aa360942a48" +dependencies = [ + "serde_derive", +] + +[[package]] +name = "serde_derive" +version = "1.0.229" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e7a5d71263a5a7d47b41f6b3f06ba276f10cc18b0931f1799f710578e2309348" +dependencies = [ + "proc-macro2", + "quote", + "syn 3.0.9", +] + +[[package]] +name = "serde_json" +version = "1.0.154" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e7e9cc8b1b85264074fbcc02a88680c4096b1e47df8f739dceb03bf482f04bd6" +dependencies = [ + "itoa", + "memchr", + "serde", + "serde_core", + "zmij", +] + +[[package]] +name = "sha2" +version = "0.10.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a7507d819769d01a365ab707794a4084392c824f54a7a6a7862f8c3d0892b283" +dependencies = [ + "cfg-if", + "cpufeatures", + "digest", +] + +[[package]] +name = "shlex" +version = "1.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0fda2ff0d084019ba4d7c6f371c95d8fd75ce3524c3cb8fb653a3023f6323e64" + +[[package]] +name = "slab" +version = "0.4.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0c790de23124f9ab44544d7ac05d60440adc586479ce501c1d6d7da3cd8c9cf5" + +[[package]] +name = "smallvec" +version = "1.16.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5b3dc8af474f516a851ff4bd12db780f948b9250ad37211e4eec0bccea54e01b" + +[[package]] +name = "syn" +version = "2.0.119" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "872831b642d1a07999a962a351ed35b955ea2cfc8f3862091e2a240a84f17297" +dependencies = [ + "proc-macro2", + "quote", + "unicode-ident", +] + +[[package]] +name = "syn" +version = "3.0.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d78c8dee4c7bf0e14673097256fed6142ce9d3b85a408189d07482442145823b" +dependencies = [ + "proc-macro2", + "quote", + "unicode-ident", +] + +[[package]] +name = "thiserror" +version = "1.0.69" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b6aaf5339b578ea85b50e080feb250a3e8ae8cfcdff9a461c9ec2904bc923f52" +dependencies = [ + "thiserror-impl 1.0.69", +] + +[[package]] +name = "thiserror" +version = "2.0.21" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "09e52cb86a36cede5cb101bf8908837b3e4c6e5e59fe7fd85c23fb56200d189e" +dependencies = [ + "thiserror-impl 2.0.21", +] + +[[package]] +name = "thiserror-impl" +version = "1.0.69" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4fee6c4efc90059e10f81e6d42c60a18f76588c3d74cb83a0b242a2b6c7504c1" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", +] + +[[package]] +name = "thiserror-impl" +version = "2.0.21" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fe5197923287db20a58125f0bc85c062f7f2c892de97b18c356f9efb14b28524" +dependencies = [ + "proc-macro2", + "quote", + "syn 3.0.9", +] + +[[package]] +name = "typenum" +version = "1.20.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b6f5e870be6c3b371b77fe0ee0bafb859fa4964b4404c27de1d380043c4dda20" + +[[package]] +name = "unicode-ident" +version = "1.0.27" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a2c754d6c33795a1c324727428e5a7dedb5b06195f9890bdbcba760d3e246563" + +[[package]] +name = "version_check" +version = "0.9.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0b928f33d975fc6ad9f86c8f283853ad26bdd5b10b7f1542aa2fa15e2289105a" + +[[package]] +name = "which" +version = "4.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "87ba24419a2078cd2b0f2ede2691b6c66d8e47836da3b6db8265ebad47afbfc7" +dependencies = [ + "either", + "home", + "once_cell", + "rustix", +] + +[[package]] +name = "windows-link" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f0805222e57f7521d6a62e36fa9163bc891acd422f971defe97d64e70d0a4fe5" + +[[package]] +name = "windows-sys" +version = "0.59.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1e38bc4d79ed67fd075bcc251a1c39b32a1776bbe92e5bef1f0bf1f8c531853b" +dependencies = [ + "windows-targets", +] + +[[package]] +name = "windows-sys" +version = "0.61.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ae137229bcbd6cdf0f7b80a31df61766145077ddf49416a728b02cb3921ff3fc" +dependencies = [ + "windows-link", +] + +[[package]] +name = "windows-targets" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9b724f72796e036ab90c1021d4780d4d3d648aca59e491e6b98e725b84e99973" +dependencies = [ + "windows_aarch64_gnullvm", + "windows_aarch64_msvc", + "windows_i686_gnu", + "windows_i686_gnullvm", + "windows_i686_msvc", + "windows_x86_64_gnu", + "windows_x86_64_gnullvm", + "windows_x86_64_msvc", +] + +[[package]] +name = "windows_aarch64_gnullvm" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "32a4622180e7a0ec044bb555404c800bc9fd9ec262ec147edd5989ccd0c02cd3" + +[[package]] +name = "windows_aarch64_msvc" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "09ec2a7bb152e2252b53fa7803150007879548bc709c039df7627cabbd05d469" + +[[package]] +name = "windows_i686_gnu" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8e9b5ad5ab802e97eb8e295ac6720e509ee4c243f69d781394014ebfe8bbfa0b" + +[[package]] +name = "windows_i686_gnullvm" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0eee52d38c090b3caa76c563b86c3a4bd71ef1a819287c19d586d7334ae8ed66" + +[[package]] +name = "windows_i686_msvc" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "240948bc05c5e7c6dabba28bf89d89ffce3e303022809e73deaefe4f6ec56c66" + +[[package]] +name = "windows_x86_64_gnu" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "147a5c80aabfbf0c7d901cb5895d1de30ef2907eb21fbbab29ca94c5b08b1a78" + +[[package]] +name = "windows_x86_64_gnullvm" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "24d5b23dc417412679681396f2b49f3de8c1473deb516bd34410872eff51ed0d" + +[[package]] +name = "windows_x86_64_msvc" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "589f6da84c646204747d1270a2a5661ea66ed1cced2631d546fdfb155959f9ec" + +[[package]] +name = "zerocopy" +version = "0.8.63" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e5fe1f8f1b06191a00962174c61aa5005e0bb391a6d80d07e24d115c01a92ed8" +dependencies = [ + "zerocopy-derive", +] + +[[package]] +name = "zerocopy-derive" +version = "0.8.63" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "863ad3ac83293fb4d740aedbfdc9240dd8d1a50c1099acd76ce80ce7c7230c7f" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", +] + +[[package]] +name = "zmij" +version = "1.0.23" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "29666d0abbfad1e3dc4dcf6144730dd3a3ab225bbbdac83319345b1b44ccfc1b" diff --git a/rust/cuda-coop/examples/scan/Cargo.toml b/rust/cuda-coop/examples/scan/Cargo.toml new file mode 100644 index 000000000000..fe413c7daac5 --- /dev/null +++ b/rust/cuda-coop/examples/scan/Cargo.toml @@ -0,0 +1,19 @@ +# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. +# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +[package] +name = "coop-scan" +version = "0.0.0" +edition = "2024" +publish = false + +[workspace] +members = ["helper"] + +[dependencies] +cuda-coop = { path = "../.." } +coop-scan-helper = { path = "helper" } +cuda-device = { git = "https://github.com/tpn/cuda-rust.git", rev = "96caed46c776eb49226137eac849e795b83b7ce0" } +cuda-host = { git = "https://github.com/tpn/cuda-rust.git", rev = "96caed46c776eb49226137eac849e795b83b7ce0", features = ["async"] } +cuda-core = { git = "https://github.com/tpn/cuda-rust.git", rev = "96caed46c776eb49226137eac849e795b83b7ce0" } +cuda-async = { git = "https://github.com/tpn/cuda-rust.git", rev = "96caed46c776eb49226137eac849e795b83b7ce0" } diff --git a/rust/cuda-coop/examples/scan/helper/Cargo.toml b/rust/cuda-coop/examples/scan/helper/Cargo.toml new file mode 100644 index 000000000000..9300b682f9fe --- /dev/null +++ b/rust/cuda-coop/examples/scan/helper/Cargo.toml @@ -0,0 +1,11 @@ +# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. +# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +[package] +name = "coop-scan-helper" +version = "0.0.0" +edition = "2024" +publish = false + +[dependencies] +cuda-coop = { path = "../../.." } diff --git a/rust/cuda-coop/examples/scan/helper/src/lib.rs b/rust/cuda-coop/examples/scan/helper/src/lib.rs new file mode 100644 index 000000000000..6b36174a4290 --- /dev/null +++ b/rust/cuda-coop/examples/scan/helper/src/lib.rs @@ -0,0 +1,39 @@ +// Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +//! A reusable, element-type-preserving block prefix helper. +#![no_std] +use cuda_coop as coop; + +/// Write exclusive sums for the valid prefix of one block tile. +/// +/// `zero` fills the missing tail before Scan, which requires every input slot +/// initialized. The caller supplies the additive identity for `T`. +pub fn prefix( + block: coop::ThreadGroup, + input: &[T], + output: &mut coop::DisjointSlice<'_, T, coop::CollectiveTiles>, + offset: usize, + valid: usize, + zero: T, +) { + let mut values = coop::ThreadData::<_, ITEMS_PER_THREAD>::new(); + coop::load( + block, + input, + &mut values, + coop::LoadOptions::default() + .offset(offset) + .valid_items(valid) + .oob_default(zero), + ); + let prefixes = coop::scan(block, &values, Default::default()); + coop::store( + block, + output, + &prefixes, + coop::StoreOptions::default() + .offset(offset) + .valid_items(valid), + ); +} diff --git a/rust/cuda-coop/examples/scan/src/main.rs b/rust/cuda-coop/examples/scan/src/main.rs new file mode 100644 index 000000000000..5fa3b659352d --- /dev/null +++ b/rust/cuda-coop/examples/scan/src/main.rs @@ -0,0 +1,75 @@ +// Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +use cuda_coop as coop; +use cuda_device::{kernel, launch_contract}; +use cuda_host::cuda_module; + +#[cuda_module] +mod kernels { + use super::*; + + #[kernel] + #[launch_contract(domain = 1, block = (128, 1, 1), dynamic_shared = inferred)] + pub fn prefix( + input: &[i32], + mut output: coop::DisjointSlice>, + ) { + let block = coop::this_block(); + let tile = block.count() as usize * ITEMS_PER_THREAD; + let offset = block.rank_at::() as usize * tile; + let remaining = if offset <= input.len() { + input.len() - offset + } else { + 0 + }; + coop_scan_helper::prefix::<_, ITEMS_PER_THREAD>( + block, + input, + &mut output, + offset, + remaining.min(tile), + 0_i32, + ); + } +} + +fn main() -> Result<(), Box> { + use cuda_async::simt::device_context::{init_device_contexts, with_cuda_context}; + use cuda_async::simt::device_operation::DeviceOperation; + use cuda_core::{DeviceBuffer, LaunchConfig1D}; + + init_device_contexts(0, 1)?; + let context = with_cuda_context(0, |context| context.clone())?; + let stream = context.default_stream(); + // SAFETY: the typed declarations and embedded kernels compile together. + let module = unsafe { kernels::load_async(0)? }; + macro_rules! run { + ($items:literal) => {{ + let tile = 128 * $items; + let input: Vec = (0..3 * tile - 7).map(|i| (i % 7) as i32).collect(); + let mut expected = Vec::with_capacity(input.len()); + for chunk in input.chunks(tile) { + let mut prefix = 0; + for &value in chunk { + expected.push(prefix); + prefix += value; + } + } + let source = DeviceBuffer::from_host(&stream, &input)?; + let target = DeviceBuffer::from_host(&stream, &vec![-1_i32; input.len()])?; + let plan = module.prepare_prefix::<$items>(LaunchConfig1D::auto_shared(3, 128))?; + let (_, target) = module + .prefix_async_owned::<$items, _, _>(&plan, source, target) + .sync()?; + let mut output = vec![0_i32; input.len()]; + target.copy_to_host(&stream, &mut output)?; + stream.synchronize()?; + assert_eq!(output, expected); + println!("items/thread={}: {} prefixes checked", $items, output.len()); + }}; + } + run!(1); + run!(4); + Ok(()) +} diff --git a/rust/cuda-coop/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..8b2898a222d5 --- /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=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "libnvvm-sys", + "nvjitlink-sys", + "serde", + "sha2", + "thiserror 2.0.21", +] + +[[package]] +name = "cuda-async" +version = "0.5.0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "anyhow", + "cuda-bindings", + "cuda-core", + "dashmap", + "futures", + "half", + "once_cell", + "rustc-hash 2.1.3", + "thiserror 1.0.69", +] + +[[package]] +name = "cuda-bindings" +version = "0.5.0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "bindgen", + "libloading", + "prettyplease", + "proc-macro2", + "quote", + "syn 2.0.119", +] + +[[package]] +name = "cuda-coop" +version = "0.0.0" +dependencies = [ + "cuda-coop-codegen", + "cuda-device", +] + +[[package]] +name = "cuda-coop-codegen" +version = "0.0.0" +dependencies = [ + "libloading", + "serde", + "serde_json", + "sha2", +] + +[[package]] +name = "cuda-core" +version = "0.5.0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "anyhow", + "cuda-bindings", + "cuda-core-derive", + "half", + "oxide-artifacts", +] + +[[package]] +name = "cuda-core-derive" +version = "0.5.0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", +] + +[[package]] +name = "cuda-device" +version = "0.2.1" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "cuda-macros", +] + +[[package]] +name = "cuda-host" +version = "0.2.1" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "cuda-artifact-finalizer", + "cuda-async", + "cuda-core", + "cuda-device", + "cuda-macros", + "half", + "libc", + "oxide-artifacts", + "ptx-parse", + "sha2", + "thiserror 2.0.21", +] + +[[package]] +name = "cuda-macros" +version = "0.2.1" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "oxide-tile-export", + "proc-macro2", + "quote", + "reserved-oxide-symbols", + "serde", + "serde_json", + "syn 2.0.119", +] + +[[package]] +name = "cuda-target-spec" +version = "0.1.0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" + +[[package]] +name = "dashmap" +version = "6.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e6361d5c062261c78a176addb82d4c821ae42bed6089de0e12603cd25de2059c" +dependencies = [ + "cfg-if", + "crossbeam-utils", + "hashbrown 0.14.5", + "lock_api", + "once_cell", + "parking_lot_core", +] + +[[package]] +name = "digest" +version = "0.10.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9ed9a281f7bc9b7576e61468ba615a66a5c8cfdff42420a70aa82701a3b1e292" +dependencies = [ + "block-buffer", + "crypto-common", +] + +[[package]] +name = "either" +version = "1.19.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0e9c71c2167ca323c882b99918929403426e2373ea17242ff5653e0d5e1058be" + +[[package]] +name = "equivalent" +version = "1.0.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "00d174d5400e5e8fd687ad1049e2f578285fa914201b1af7e8b112a4546bd826" + +[[package]] +name = "errno" +version = "0.3.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "39cab71617ae0d63f51a36d69f866391735b51691dbda63cf6f96d042b63efeb" +dependencies = [ + "libc", + "windows-sys 0.61.2", +] + +[[package]] +name = "foldhash" +version = "0.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d9c4f5dac5e15c24eb999c26181a6ca40b39fe946cbe4c263c7209467bc83af2" + +[[package]] +name = "futures" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9a31d2a3fbaaeb2af2368bbdd904aa8e812d3c04a1ee10d3171f52d556e5d0a3" +dependencies = [ + "futures-channel", + "futures-core", + "futures-executor", + "futures-io", + "futures-sink", + "futures-task", + "futures-util", +] + +[[package]] +name = "futures-channel" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b1f9e3d69d39e4862ffed03ed071a76f9a13ba1d9109d355b0f0aa6b15e393c4" +dependencies = [ + "futures-core", + "futures-sink", +] + +[[package]] +name = "futures-core" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "92d699e522242e69e3003b94ecc1f960f3a5e015aa7c5d7486e65ad01dd94f5e" + +[[package]] +name = "futures-executor" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "031b47cf1a3c6cc8bc2fc76cd437f521619387907d469316e7c0bc278f1f5432" +dependencies = [ + "futures-core", + "futures-task", + "futures-util", +] + +[[package]] +name = "futures-io" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "53c0fa8157de1303bfffdaa1cc2a673bfffb60102f76b0ef4441659124373fed" + +[[package]] +name = "futures-macro" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9fb9654ba8355388abeb8dcb4fc62f511300867002afc858860463bdd9fe0c44" +dependencies = [ + "proc-macro2", + "quote", + "syn 3.0.9", +] + +[[package]] +name = "futures-sink" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1944426bf7d03f1d14f708785e4b33efd750b36d48a157b836b3efc15ede8e1d" + +[[package]] +name = "futures-task" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cd417de3d1d015fc3bfd2b1ea46dfc7bab72ef86f1cc7cc9c78e728b34a6d1fd" + +[[package]] +name = "futures-util" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0d50a92467f8ba5dd6e3ee5d4bd04d73ab2e4e1c44474a0674821dfce14b79bc" +dependencies = [ + "futures-channel", + "futures-core", + "futures-io", + "futures-macro", + "futures-sink", + "futures-task", + "memchr", + "pin-project-lite", + "slab", +] + +[[package]] +name = "generic-array" +version = "0.14.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "85649ca51fd72272d7821adaf274ad91c288277713d9c18820d8499a7ff69e9a" +dependencies = [ + "typenum", + "version_check", +] + +[[package]] +name = "glob" +version = "0.3.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e4eba85ea1d0a966a983acd07deee566e67395d2d96b6fb39e62b5a833f1eb0b" + +[[package]] +name = "half" +version = "2.7.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" +dependencies = [ + "cfg-if", + "crunchy", + "zerocopy", +] + +[[package]] +name = "hashbrown" +version = "0.14.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e5274423e17b7c9fc20b6e7e208532f9b19825d82dfd615708b70edd83df41f1" + +[[package]] +name = "hashbrown" +version = "0.15.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9229cfe53dfd69f0609a49f65461bd93001ea1ef889cd5529dd176593f5338a1" +dependencies = [ + "foldhash", +] + +[[package]] +name = "hashbrown" +version = "0.17.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ed5909b6e89a2db4456e54cd5f673791d7eca6732202bbf2a9cc504fe2f9b84a" + +[[package]] +name = "home" +version = "0.5.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cc627f471c528ff0c4a49e1d5e60450c8f6461dd6d10ba9dcd3a61d3dff7728d" +dependencies = [ + "windows-sys 0.61.2", +] + +[[package]] +name = "indexmap" +version = "2.14.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cc4e190f5d26ca7051642629da2c52fc03bde85a03197c99408dcd291734c855" +dependencies = [ + "equivalent", + "hashbrown 0.17.1", +] + +[[package]] +name = "itertools" +version = "0.12.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ba291022dbbd398a455acf126c1e341954079855bc60dfdda641363bd6922569" +dependencies = [ + "either", +] + +[[package]] +name = "itoa" +version = "1.0.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8f42a60cbdf9a97f5d2305f08a87dc4e09308d1276d28c869c684d7777685682" + +[[package]] +name = "lazy_static" +version = "1.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "20870f649af7073d53e38067b2a84312175d56ea15217e1b15bc83506ec50afb" + +[[package]] +name = "lazycell" +version = "1.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "830d08ce1d1d941e6b30645f1a0eb5643013d835ce3779a5fc208261dbe10f55" + +[[package]] +name = "libc" +version = "0.2.191" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "65cd87c9a6d60536580b316ad985726104730a5151c99845f2490428ef91240a" + +[[package]] +name = "libloading" +version = "0.8.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d7c4b02199fee7c5d21a5ae7d8cfa79a6ef5bb2fc834d6e9058e89c825efdc55" +dependencies = [ + "cfg-if", + "windows-link", +] + +[[package]] +name = "libnvvm-sys" +version = "0.2.1" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "cuda-target-spec", + "libloading", + "thiserror 2.0.21", +] + +[[package]] +name = "linux-raw-sys" +version = "0.4.15" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d26c52dbd32dccf2d10cac7725f8eae5296885fb5703b261f7d0a0739ec807ab" + +[[package]] +name = "lock_api" +version = "0.4.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "224399e74b87b5f3557511d98dff8b14089b3dadafcab6bb93eab67d3aace965" +dependencies = [ + "scopeguard", +] + +[[package]] +name = "log" +version = "0.4.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9f8bd3e56ce4dfc153cf470fffbfa98c7620958b312ca5c3a4b8d5181fd13c6" + +[[package]] +name = "memchr" +version = "2.8.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cf8baf1c55e62ffcace7a9f06f4bd9cd3f0c4beb022d3b367256b91b87513d98" + +[[package]] +name = "minimal-lexical" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "68354c5c6bd36d73ff3feceb05efa59b6acb7626617f4962be322a825e61f79a" + +[[package]] +name = "nom" +version = "7.1.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d273983c5a657a70a3e8f2a01329822f3b8c8172b73826411a55751e404a0a4a" +dependencies = [ + "memchr", + "minimal-lexical", +] + +[[package]] +name = "nvjitlink-sys" +version = "0.2.1" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "libloading", + "thiserror 2.0.21", +] + +[[package]] +name = "object" +version = "0.36.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "62948e14d923ea95ea2c7c86c71013138b66525b86bdc08d2dcc262bdb497b87" +dependencies = [ + "crc32fast", + "hashbrown 0.15.5", + "indexmap", + "memchr", +] + +[[package]] +name = "once_cell" +version = "1.21.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9f7c3e4beb33f85d45ae3e3a1792185706c8e16d043238c593331cc7cd313b50" + +[[package]] +name = "oxide-artifacts" +version = "0.2.1" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "object", +] + +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "serde", + "serde_json", + "thiserror 2.0.21", +] + +[[package]] +name = "parking_lot_core" +version = "0.9.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2621685985a2ebf1c516881c026032ac7deafcda1a2c9b7850dc81e3dfcb64c1" +dependencies = [ + "cfg-if", + "libc", + "redox_syscall", + "smallvec", + "windows-link", +] + +[[package]] +name = "pin-project-lite" +version = "0.2.17" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a89322df9ebe1c1578d689c92318e070967d1042b512afbe49518723f4e6d5cd" + +[[package]] +name = "prettyplease" +version = "0.2.37" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "479ca8adacdd7ce8f1fb39ce9ecccbfe93a3f1344b3d0d97f20bc0196208f62b" +dependencies = [ + "proc-macro2", + "syn 2.0.119", +] + +[[package]] +name = "proc-macro2" +version = "1.0.107" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "985e7ec9bb745e6ce6535b544d84d6cd6f7ad8bd711c398938ae983b91a766d9" +dependencies = [ + "unicode-ident", +] + +[[package]] +name = "ptx-parse" +version = "0.1.0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" + +[[package]] +name = "quote" +version = "1.0.47" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1fbf4db142a473a8d80c26bbf18454ed458bf8d26c8219c331daecfdbd079001" +dependencies = [ + "proc-macro2", +] + +[[package]] +name = "redox_syscall" +version = "0.5.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ed2bf2547551a7053d6fdfafda3f938979645c44812fbfcda098faae3f1a362d" +dependencies = [ + "bitflags", +] + +[[package]] +name = "regex" +version = "1.13.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f020237b6c8eed93db2e2cb53c00c60a8e1bc73da7d073199a1180401450218d" +dependencies = [ + "aho-corasick", + "memchr", + "regex-automata", + "regex-syntax", +] + +[[package]] +name = "regex-automata" +version = "0.4.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ad8553b9b26413251cbf30e620595c7a41b3887f03da04579c0e6b0d6a06b4b2" +dependencies = [ + "aho-corasick", + "memchr", + "regex-syntax", +] + +[[package]] +name = "regex-syntax" +version = "0.8.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d6f6ff9a378485b298a5286656da665ba74413d36db0979633275d2e708145d4" + +[[package]] +name = "reserved-oxide-symbols" +version = "0.2.1" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" + +[[package]] +name = "rustc-hash" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "08d43f7aa6b08d49f382cde6a7982047c3426db949b1424bc4b7ec9ae12c6ce2" + +[[package]] +name = "rustc-hash" +version = "2.1.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6b1e7f9a428571be2dc5bc0505c13fb6bf936822b894ec87abf8a08a4e51742d" + +[[package]] +name = "rustix" +version = "0.38.44" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fdb5bc1ae2baa591800df16c9ca78619bf65c0488b41b96ccec5d11220d8c154" +dependencies = [ + "bitflags", + "errno", + "libc", + "linux-raw-sys", + "windows-sys 0.59.0", +] + +[[package]] +name = "scopeguard" +version = "1.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "94143f37725109f92c262ed2cf5e59bce7498c01bcc1502d7b9afe439a4e9f49" + +[[package]] +name = "serde" +version = "1.0.229" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4148590afebada386688f18773da617792bf2ef03ffc1e4cbd2b1d45b023e0ba" +dependencies = [ + "serde_core", + "serde_derive", +] + +[[package]] +name = "serde_core" +version = "1.0.229" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "67dca2c9c51e58a4791a4b1ed58308b39c64224d349a935ab5039aa360942a48" +dependencies = [ + "serde_derive", +] + +[[package]] +name = "serde_derive" +version = "1.0.229" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e7a5d71263a5a7d47b41f6b3f06ba276f10cc18b0931f1799f710578e2309348" +dependencies = [ + "proc-macro2", + "quote", + "syn 3.0.9", +] + +[[package]] +name = "serde_json" +version = "1.0.154" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e7e9cc8b1b85264074fbcc02a88680c4096b1e47df8f739dceb03bf482f04bd6" +dependencies = [ + "itoa", + "memchr", + "serde", + "serde_core", + "zmij", +] + +[[package]] +name = "sha2" +version = "0.10.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a7507d819769d01a365ab707794a4084392c824f54a7a6a7862f8c3d0892b283" +dependencies = [ + "cfg-if", + "cpufeatures", + "digest", +] + +[[package]] +name = "shlex" +version = "1.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0fda2ff0d084019ba4d7c6f371c95d8fd75ce3524c3cb8fb653a3023f6323e64" + +[[package]] +name = "slab" +version = "0.4.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0c790de23124f9ab44544d7ac05d60440adc586479ce501c1d6d7da3cd8c9cf5" + +[[package]] +name = "smallvec" +version = "1.16.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5b3dc8af474f516a851ff4bd12db780f948b9250ad37211e4eec0bccea54e01b" + +[[package]] +name = "syn" +version = "2.0.119" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "872831b642d1a07999a962a351ed35b955ea2cfc8f3862091e2a240a84f17297" +dependencies = [ + "proc-macro2", + "quote", + "unicode-ident", +] + +[[package]] +name = "syn" +version = "3.0.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d78c8dee4c7bf0e14673097256fed6142ce9d3b85a408189d07482442145823b" +dependencies = [ + "proc-macro2", + "quote", + "unicode-ident", +] + +[[package]] +name = "thiserror" +version = "1.0.69" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b6aaf5339b578ea85b50e080feb250a3e8ae8cfcdff9a461c9ec2904bc923f52" +dependencies = [ + "thiserror-impl 1.0.69", +] + +[[package]] +name = "thiserror" +version = "2.0.21" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "09e52cb86a36cede5cb101bf8908837b3e4c6e5e59fe7fd85c23fb56200d189e" +dependencies = [ + "thiserror-impl 2.0.21", +] + +[[package]] +name = "thiserror-impl" +version = "1.0.69" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4fee6c4efc90059e10f81e6d42c60a18f76588c3d74cb83a0b242a2b6c7504c1" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", +] + +[[package]] +name = "thiserror-impl" +version = "2.0.21" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fe5197923287db20a58125f0bc85c062f7f2c892de97b18c356f9efb14b28524" +dependencies = [ + "proc-macro2", + "quote", + "syn 3.0.9", +] + +[[package]] +name = "typenum" +version = "1.20.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b6f5e870be6c3b371b77fe0ee0bafb859fa4964b4404c27de1d380043c4dda20" + +[[package]] +name = "unicode-ident" +version = "1.0.27" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a2c754d6c33795a1c324727428e5a7dedb5b06195f9890bdbcba760d3e246563" + +[[package]] +name = "version_check" +version = "0.9.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0b928f33d975fc6ad9f86c8f283853ad26bdd5b10b7f1542aa2fa15e2289105a" + +[[package]] +name = "which" +version = "4.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "87ba24419a2078cd2b0f2ede2691b6c66d8e47836da3b6db8265ebad47afbfc7" +dependencies = [ + "either", + "home", + "once_cell", + "rustix", +] + +[[package]] +name = "windows-link" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f0805222e57f7521d6a62e36fa9163bc891acd422f971defe97d64e70d0a4fe5" + +[[package]] +name = "windows-sys" +version = "0.59.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1e38bc4d79ed67fd075bcc251a1c39b32a1776bbe92e5bef1f0bf1f8c531853b" +dependencies = [ + "windows-targets", +] + +[[package]] +name = "windows-sys" +version = "0.61.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ae137229bcbd6cdf0f7b80a31df61766145077ddf49416a728b02cb3921ff3fc" +dependencies = [ + "windows-link", +] + +[[package]] +name = "windows-targets" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9b724f72796e036ab90c1021d4780d4d3d648aca59e491e6b98e725b84e99973" +dependencies = [ + "windows_aarch64_gnullvm", + "windows_aarch64_msvc", + "windows_i686_gnu", + "windows_i686_gnullvm", + "windows_i686_msvc", + "windows_x86_64_gnu", + "windows_x86_64_gnullvm", + "windows_x86_64_msvc", +] + +[[package]] +name = "windows_aarch64_gnullvm" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "32a4622180e7a0ec044bb555404c800bc9fd9ec262ec147edd5989ccd0c02cd3" + +[[package]] +name = "windows_aarch64_msvc" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "09ec2a7bb152e2252b53fa7803150007879548bc709c039df7627cabbd05d469" + +[[package]] +name = "windows_i686_gnu" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8e9b5ad5ab802e97eb8e295ac6720e509ee4c243f69d781394014ebfe8bbfa0b" + +[[package]] +name = "windows_i686_gnullvm" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0eee52d38c090b3caa76c563b86c3a4bd71ef1a819287c19d586d7334ae8ed66" + +[[package]] +name = "windows_i686_msvc" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "240948bc05c5e7c6dabba28bf89d89ffce3e303022809e73deaefe4f6ec56c66" + +[[package]] +name = "windows_x86_64_gnu" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "147a5c80aabfbf0c7d901cb5895d1de30ef2907eb21fbbab29ca94c5b08b1a78" + +[[package]] +name = "windows_x86_64_gnullvm" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "24d5b23dc417412679681396f2b49f3de8c1473deb516bd34410872eff51ed0d" + +[[package]] +name = "windows_x86_64_msvc" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "589f6da84c646204747d1270a2a5661ea66ed1cced2631d546fdfb155959f9ec" + +[[package]] +name = "zerocopy" +version = "0.8.63" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e5fe1f8f1b06191a00962174c61aa5005e0bb391a6d80d07e24d115c01a92ed8" +dependencies = [ + "zerocopy-derive", +] + +[[package]] +name = "zerocopy-derive" +version = "0.8.63" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "863ad3ac83293fb4d740aedbfdc9240dd8d1a50c1099acd76ce80ce7c7230c7f" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", +] + +[[package]] +name = "zmij" +version = "1.0.23" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "29666d0abbfad1e3dc4dcf6144730dd3a3ab225bbbdac83319345b1b44ccfc1b" diff --git a/rust/cuda-coop/examples/shuffle/Cargo.toml b/rust/cuda-coop/examples/shuffle/Cargo.toml new file mode 100644 index 000000000000..cea6058c2dd3 --- /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 = "96caed46c776eb49226137eac849e795b83b7ce0" } +cuda-host = { git = "https://github.com/tpn/cuda-rust.git", rev = "96caed46c776eb49226137eac849e795b83b7ce0", features = ["async"] } +cuda-core = { git = "https://github.com/tpn/cuda-rust.git", rev = "96caed46c776eb49226137eac849e795b83b7ce0" } +cuda-async = { git = "https://github.com/tpn/cuda-rust.git", rev = "96caed46c776eb49226137eac849e795b83b7ce0" } diff --git a/rust/cuda-coop/examples/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/lib.rs b/rust/cuda-coop/src/lib.rs new file mode 100644 index 000000000000..e55f911b6b76 --- /dev/null +++ b/rust/cuda-coop/src/lib.rs @@ -0,0 +1,125 @@ +// 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. +//! 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 load_store; +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 load_store::{ + Algorithm, Direct, LoadOptions, LoadSource, LoadStoreGroup, StoreOptions, StoreValue, Striped, + Transpose, Vectorize, WarpTranspose, WarpTransposeTimesliced, load, store, +}; +pub use reduce::{ + RakingCommutativeOnly, ReduceAlgorithm, ReduceGroup, ReduceOperator, ReduceOptions, + ReduceValue, SumOptions, WarpReductions, reduce, sum, +}; +pub use 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/reduce.rs b/rust/cuda-coop/src/reduce.rs new file mode 100644 index 000000000000..6ed392645bb8 --- /dev/null +++ b/rust/cuda-coop/src/reduce.rs @@ -0,0 +1,457 @@ +// 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; } + )+}; +} +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. Custom operators are +/// not yet supported; all current operators permit the commutative algorithm. +pub trait ReduceOperator: sealed::Operator {} +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); + +/// 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 a built-in operator, retaining its actual generic value. + 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); + boundary::reduce::(&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, + { + boundary::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 { + 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..8cbc4b45377d --- /dev/null +++ b/rust/cuda-coop/src/reduce/boundary.rs @@ -0,0 +1,427 @@ +// Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +use super::{ReduceAlgorithm, ReduceGroup, ReduceOperator, 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: ReduceOperator, + 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) +} 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..e567d14abd08 --- /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=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "libnvvm-sys", + "nvjitlink-sys", + "serde", + "sha2", + "thiserror 2.0.21", +] + +[[package]] +name = "cuda-async" +version = "0.5.0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "anyhow", + "cuda-bindings", + "cuda-core", + "dashmap", + "futures", + "half", + "once_cell", + "rustc-hash 2.1.3", + "thiserror 1.0.69", +] + +[[package]] +name = "cuda-bindings" +version = "0.5.0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "bindgen", + "libloading", + "prettyplease", + "proc-macro2", + "quote", + "syn 2.0.119", +] + +[[package]] +name = "cuda-coop" +version = "0.0.0" +dependencies = [ + "cuda-coop-codegen", + "cuda-device", +] + +[[package]] +name = "cuda-coop-codegen" +version = "0.0.0" +dependencies = [ + "libloading", + "serde", + "serde_json", + "sha2", +] + +[[package]] +name = "cuda-core" +version = "0.5.0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "anyhow", + "cuda-bindings", + "cuda-core-derive", + "half", + "oxide-artifacts", +] + +[[package]] +name = "cuda-core-derive" +version = "0.5.0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", +] + +[[package]] +name = "cuda-device" +version = "0.2.1" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "cuda-macros", +] + +[[package]] +name = "cuda-host" +version = "0.2.1" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "cuda-artifact-finalizer", + "cuda-async", + "cuda-core", + "cuda-device", + "cuda-macros", + "half", + "libc", + "oxide-artifacts", + "ptx-parse", + "sha2", + "thiserror 2.0.21", +] + +[[package]] +name = "cuda-macros" +version = "0.2.1" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "oxide-tile-export", + "proc-macro2", + "quote", + "reserved-oxide-symbols", + "serde", + "serde_json", + "syn 2.0.119", +] + +[[package]] +name = "cuda-target-spec" +version = "0.1.0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" + +[[package]] +name = "dashmap" +version = "6.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e6361d5c062261c78a176addb82d4c821ae42bed6089de0e12603cd25de2059c" +dependencies = [ + "cfg-if", + "crossbeam-utils", + "hashbrown 0.14.5", + "lock_api", + "once_cell", + "parking_lot_core", +] + +[[package]] +name = "digest" +version = "0.10.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9ed9a281f7bc9b7576e61468ba615a66a5c8cfdff42420a70aa82701a3b1e292" +dependencies = [ + "block-buffer", + "crypto-common", +] + +[[package]] +name = "either" +version = "1.19.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0e9c71c2167ca323c882b99918929403426e2373ea17242ff5653e0d5e1058be" + +[[package]] +name = "equivalent" +version = "1.0.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "00d174d5400e5e8fd687ad1049e2f578285fa914201b1af7e8b112a4546bd826" + +[[package]] +name = "errno" +version = "0.3.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "39cab71617ae0d63f51a36d69f866391735b51691dbda63cf6f96d042b63efeb" +dependencies = [ + "libc", + "windows-sys 0.61.2", +] + +[[package]] +name = "foldhash" +version = "0.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d9c4f5dac5e15c24eb999c26181a6ca40b39fe946cbe4c263c7209467bc83af2" + +[[package]] +name = "futures" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9a31d2a3fbaaeb2af2368bbdd904aa8e812d3c04a1ee10d3171f52d556e5d0a3" +dependencies = [ + "futures-channel", + "futures-core", + "futures-executor", + "futures-io", + "futures-sink", + "futures-task", + "futures-util", +] + +[[package]] +name = "futures-channel" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b1f9e3d69d39e4862ffed03ed071a76f9a13ba1d9109d355b0f0aa6b15e393c4" +dependencies = [ + "futures-core", + "futures-sink", +] + +[[package]] +name = "futures-core" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "92d699e522242e69e3003b94ecc1f960f3a5e015aa7c5d7486e65ad01dd94f5e" + +[[package]] +name = "futures-executor" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "031b47cf1a3c6cc8bc2fc76cd437f521619387907d469316e7c0bc278f1f5432" +dependencies = [ + "futures-core", + "futures-task", + "futures-util", +] + +[[package]] +name = "futures-io" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "53c0fa8157de1303bfffdaa1cc2a673bfffb60102f76b0ef4441659124373fed" + +[[package]] +name = "futures-macro" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9fb9654ba8355388abeb8dcb4fc62f511300867002afc858860463bdd9fe0c44" +dependencies = [ + "proc-macro2", + "quote", + "syn 3.0.9", +] + +[[package]] +name = "futures-sink" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1944426bf7d03f1d14f708785e4b33efd750b36d48a157b836b3efc15ede8e1d" + +[[package]] +name = "futures-task" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cd417de3d1d015fc3bfd2b1ea46dfc7bab72ef86f1cc7cc9c78e728b34a6d1fd" + +[[package]] +name = "futures-util" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0d50a92467f8ba5dd6e3ee5d4bd04d73ab2e4e1c44474a0674821dfce14b79bc" +dependencies = [ + "futures-channel", + "futures-core", + "futures-io", + "futures-macro", + "futures-sink", + "futures-task", + "memchr", + "pin-project-lite", + "slab", +] + +[[package]] +name = "generic-array" +version = "0.14.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "85649ca51fd72272d7821adaf274ad91c288277713d9c18820d8499a7ff69e9a" +dependencies = [ + "typenum", + "version_check", +] + +[[package]] +name = "glob" +version = "0.3.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e4eba85ea1d0a966a983acd07deee566e67395d2d96b6fb39e62b5a833f1eb0b" + +[[package]] +name = "half" +version = "2.7.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" +dependencies = [ + "cfg-if", + "crunchy", + "zerocopy", +] + +[[package]] +name = "hashbrown" +version = "0.14.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e5274423e17b7c9fc20b6e7e208532f9b19825d82dfd615708b70edd83df41f1" + +[[package]] +name = "hashbrown" +version = "0.15.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9229cfe53dfd69f0609a49f65461bd93001ea1ef889cd5529dd176593f5338a1" +dependencies = [ + "foldhash", +] + +[[package]] +name = "hashbrown" +version = "0.17.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ed5909b6e89a2db4456e54cd5f673791d7eca6732202bbf2a9cc504fe2f9b84a" + +[[package]] +name = "home" +version = "0.5.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cc627f471c528ff0c4a49e1d5e60450c8f6461dd6d10ba9dcd3a61d3dff7728d" +dependencies = [ + "windows-sys 0.61.2", +] + +[[package]] +name = "indexmap" +version = "2.14.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cc4e190f5d26ca7051642629da2c52fc03bde85a03197c99408dcd291734c855" +dependencies = [ + "equivalent", + "hashbrown 0.17.1", +] + +[[package]] +name = "itertools" +version = "0.12.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ba291022dbbd398a455acf126c1e341954079855bc60dfdda641363bd6922569" +dependencies = [ + "either", +] + +[[package]] +name = "itoa" +version = "1.0.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8f42a60cbdf9a97f5d2305f08a87dc4e09308d1276d28c869c684d7777685682" + +[[package]] +name = "lazy_static" +version = "1.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "20870f649af7073d53e38067b2a84312175d56ea15217e1b15bc83506ec50afb" + +[[package]] +name = "lazycell" +version = "1.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "830d08ce1d1d941e6b30645f1a0eb5643013d835ce3779a5fc208261dbe10f55" + +[[package]] +name = "libc" +version = "0.2.191" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "65cd87c9a6d60536580b316ad985726104730a5151c99845f2490428ef91240a" + +[[package]] +name = "libloading" +version = "0.8.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d7c4b02199fee7c5d21a5ae7d8cfa79a6ef5bb2fc834d6e9058e89c825efdc55" +dependencies = [ + "cfg-if", + "windows-link", +] + +[[package]] +name = "libnvvm-sys" +version = "0.2.1" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "cuda-target-spec", + "libloading", + "thiserror 2.0.21", +] + +[[package]] +name = "linux-raw-sys" +version = "0.4.15" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d26c52dbd32dccf2d10cac7725f8eae5296885fb5703b261f7d0a0739ec807ab" + +[[package]] +name = "lock_api" +version = "0.4.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "224399e74b87b5f3557511d98dff8b14089b3dadafcab6bb93eab67d3aace965" +dependencies = [ + "scopeguard", +] + +[[package]] +name = "log" +version = "0.4.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9f8bd3e56ce4dfc153cf470fffbfa98c7620958b312ca5c3a4b8d5181fd13c6" + +[[package]] +name = "memchr" +version = "2.8.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cf8baf1c55e62ffcace7a9f06f4bd9cd3f0c4beb022d3b367256b91b87513d98" + +[[package]] +name = "minimal-lexical" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "68354c5c6bd36d73ff3feceb05efa59b6acb7626617f4962be322a825e61f79a" + +[[package]] +name = "nom" +version = "7.1.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d273983c5a657a70a3e8f2a01329822f3b8c8172b73826411a55751e404a0a4a" +dependencies = [ + "memchr", + "minimal-lexical", +] + +[[package]] +name = "nvjitlink-sys" +version = "0.2.1" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "libloading", + "thiserror 2.0.21", +] + +[[package]] +name = "object" +version = "0.36.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "62948e14d923ea95ea2c7c86c71013138b66525b86bdc08d2dcc262bdb497b87" +dependencies = [ + "crc32fast", + "hashbrown 0.15.5", + "indexmap", + "memchr", +] + +[[package]] +name = "once_cell" +version = "1.21.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9f7c3e4beb33f85d45ae3e3a1792185706c8e16d043238c593331cc7cd313b50" + +[[package]] +name = "oxide-artifacts" +version = "0.2.1" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "object", +] + +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "serde", + "serde_json", + "thiserror 2.0.21", +] + +[[package]] +name = "parking_lot_core" +version = "0.9.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2621685985a2ebf1c516881c026032ac7deafcda1a2c9b7850dc81e3dfcb64c1" +dependencies = [ + "cfg-if", + "libc", + "redox_syscall", + "smallvec", + "windows-link", +] + +[[package]] +name = "pin-project-lite" +version = "0.2.17" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a89322df9ebe1c1578d689c92318e070967d1042b512afbe49518723f4e6d5cd" + +[[package]] +name = "prettyplease" +version = "0.2.37" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "479ca8adacdd7ce8f1fb39ce9ecccbfe93a3f1344b3d0d97f20bc0196208f62b" +dependencies = [ + "proc-macro2", + "syn 2.0.119", +] + +[[package]] +name = "proc-macro2" +version = "1.0.107" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "985e7ec9bb745e6ce6535b544d84d6cd6f7ad8bd711c398938ae983b91a766d9" +dependencies = [ + "unicode-ident", +] + +[[package]] +name = "ptx-parse" +version = "0.1.0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" + +[[package]] +name = "quote" +version = "1.0.47" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1fbf4db142a473a8d80c26bbf18454ed458bf8d26c8219c331daecfdbd079001" +dependencies = [ + "proc-macro2", +] + +[[package]] +name = "redox_syscall" +version = "0.5.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ed2bf2547551a7053d6fdfafda3f938979645c44812fbfcda098faae3f1a362d" +dependencies = [ + "bitflags", +] + +[[package]] +name = "regex" +version = "1.13.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f020237b6c8eed93db2e2cb53c00c60a8e1bc73da7d073199a1180401450218d" +dependencies = [ + "aho-corasick", + "memchr", + "regex-automata", + "regex-syntax", +] + +[[package]] +name = "regex-automata" +version = "0.4.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ad8553b9b26413251cbf30e620595c7a41b3887f03da04579c0e6b0d6a06b4b2" +dependencies = [ + "aho-corasick", + "memchr", + "regex-syntax", +] + +[[package]] +name = "regex-syntax" +version = "0.8.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d6f6ff9a378485b298a5286656da665ba74413d36db0979633275d2e708145d4" + +[[package]] +name = "reserved-oxide-symbols" +version = "0.2.1" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" + +[[package]] +name = "rustc-hash" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "08d43f7aa6b08d49f382cde6a7982047c3426db949b1424bc4b7ec9ae12c6ce2" + +[[package]] +name = "rustc-hash" +version = "2.1.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6b1e7f9a428571be2dc5bc0505c13fb6bf936822b894ec87abf8a08a4e51742d" + +[[package]] +name = "rustix" +version = "0.38.44" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fdb5bc1ae2baa591800df16c9ca78619bf65c0488b41b96ccec5d11220d8c154" +dependencies = [ + "bitflags", + "errno", + "libc", + "linux-raw-sys", + "windows-sys 0.59.0", +] + +[[package]] +name = "scopeguard" +version = "1.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "94143f37725109f92c262ed2cf5e59bce7498c01bcc1502d7b9afe439a4e9f49" + +[[package]] +name = "serde" +version = "1.0.229" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4148590afebada386688f18773da617792bf2ef03ffc1e4cbd2b1d45b023e0ba" +dependencies = [ + "serde_core", + "serde_derive", +] + +[[package]] +name = "serde_core" +version = "1.0.229" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "67dca2c9c51e58a4791a4b1ed58308b39c64224d349a935ab5039aa360942a48" +dependencies = [ + "serde_derive", +] + +[[package]] +name = "serde_derive" +version = "1.0.229" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e7a5d71263a5a7d47b41f6b3f06ba276f10cc18b0931f1799f710578e2309348" +dependencies = [ + "proc-macro2", + "quote", + "syn 3.0.9", +] + +[[package]] +name = "serde_json" +version = "1.0.154" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e7e9cc8b1b85264074fbcc02a88680c4096b1e47df8f739dceb03bf482f04bd6" +dependencies = [ + "itoa", + "memchr", + "serde", + "serde_core", + "zmij", +] + +[[package]] +name = "sha2" +version = "0.10.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a7507d819769d01a365ab707794a4084392c824f54a7a6a7862f8c3d0892b283" +dependencies = [ + "cfg-if", + "cpufeatures", + "digest", +] + +[[package]] +name = "shlex" +version = "1.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0fda2ff0d084019ba4d7c6f371c95d8fd75ce3524c3cb8fb653a3023f6323e64" + +[[package]] +name = "slab" +version = "0.4.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0c790de23124f9ab44544d7ac05d60440adc586479ce501c1d6d7da3cd8c9cf5" + +[[package]] +name = "smallvec" +version = "1.16.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5b3dc8af474f516a851ff4bd12db780f948b9250ad37211e4eec0bccea54e01b" + +[[package]] +name = "syn" +version = "2.0.119" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "872831b642d1a07999a962a351ed35b955ea2cfc8f3862091e2a240a84f17297" +dependencies = [ + "proc-macro2", + "quote", + "unicode-ident", +] + +[[package]] +name = "syn" +version = "3.0.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d78c8dee4c7bf0e14673097256fed6142ce9d3b85a408189d07482442145823b" +dependencies = [ + "proc-macro2", + "quote", + "unicode-ident", +] + +[[package]] +name = "thiserror" +version = "1.0.69" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b6aaf5339b578ea85b50e080feb250a3e8ae8cfcdff9a461c9ec2904bc923f52" +dependencies = [ + "thiserror-impl 1.0.69", +] + +[[package]] +name = "thiserror" +version = "2.0.21" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "09e52cb86a36cede5cb101bf8908837b3e4c6e5e59fe7fd85c23fb56200d189e" +dependencies = [ + "thiserror-impl 2.0.21", +] + +[[package]] +name = "thiserror-impl" +version = "1.0.69" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4fee6c4efc90059e10f81e6d42c60a18f76588c3d74cb83a0b242a2b6c7504c1" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", +] + +[[package]] +name = "thiserror-impl" +version = "2.0.21" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fe5197923287db20a58125f0bc85c062f7f2c892de97b18c356f9efb14b28524" +dependencies = [ + "proc-macro2", + "quote", + "syn 3.0.9", +] + +[[package]] +name = "typenum" +version = "1.20.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b6f5e870be6c3b371b77fe0ee0bafb859fa4964b4404c27de1d380043c4dda20" + +[[package]] +name = "unicode-ident" +version = "1.0.27" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a2c754d6c33795a1c324727428e5a7dedb5b06195f9890bdbcba760d3e246563" + +[[package]] +name = "version_check" +version = "0.9.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0b928f33d975fc6ad9f86c8f283853ad26bdd5b10b7f1542aa2fa15e2289105a" + +[[package]] +name = "which" +version = "4.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "87ba24419a2078cd2b0f2ede2691b6c66d8e47836da3b6db8265ebad47afbfc7" +dependencies = [ + "either", + "home", + "once_cell", + "rustix", +] + +[[package]] +name = "windows-link" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f0805222e57f7521d6a62e36fa9163bc891acd422f971defe97d64e70d0a4fe5" + +[[package]] +name = "windows-sys" +version = "0.59.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1e38bc4d79ed67fd075bcc251a1c39b32a1776bbe92e5bef1f0bf1f8c531853b" +dependencies = [ + "windows-targets", +] + +[[package]] +name = "windows-sys" +version = "0.61.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ae137229bcbd6cdf0f7b80a31df61766145077ddf49416a728b02cb3921ff3fc" +dependencies = [ + "windows-link", +] + +[[package]] +name = "windows-targets" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9b724f72796e036ab90c1021d4780d4d3d648aca59e491e6b98e725b84e99973" +dependencies = [ + "windows_aarch64_gnullvm", + "windows_aarch64_msvc", + "windows_i686_gnu", + "windows_i686_gnullvm", + "windows_i686_msvc", + "windows_x86_64_gnu", + "windows_x86_64_gnullvm", + "windows_x86_64_msvc", +] + +[[package]] +name = "windows_aarch64_gnullvm" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "32a4622180e7a0ec044bb555404c800bc9fd9ec262ec147edd5989ccd0c02cd3" + +[[package]] +name = "windows_aarch64_msvc" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "09ec2a7bb152e2252b53fa7803150007879548bc709c039df7627cabbd05d469" + +[[package]] +name = "windows_i686_gnu" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8e9b5ad5ab802e97eb8e295ac6720e509ee4c243f69d781394014ebfe8bbfa0b" + +[[package]] +name = "windows_i686_gnullvm" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0eee52d38c090b3caa76c563b86c3a4bd71ef1a819287c19d586d7334ae8ed66" + +[[package]] +name = "windows_i686_msvc" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "240948bc05c5e7c6dabba28bf89d89ffce3e303022809e73deaefe4f6ec56c66" + +[[package]] +name = "windows_x86_64_gnu" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "147a5c80aabfbf0c7d901cb5895d1de30ef2907eb21fbbab29ca94c5b08b1a78" + +[[package]] +name = "windows_x86_64_gnullvm" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "24d5b23dc417412679681396f2b49f3de8c1473deb516bd34410872eff51ed0d" + +[[package]] +name = "windows_x86_64_msvc" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "589f6da84c646204747d1270a2a5661ea66ed1cced2631d546fdfb155959f9ec" + +[[package]] +name = "zerocopy" +version = "0.8.63" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e5fe1f8f1b06191a00962174c61aa5005e0bb391a6d80d07e24d115c01a92ed8" +dependencies = [ + "zerocopy-derive", +] + +[[package]] +name = "zerocopy-derive" +version = "0.8.63" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "863ad3ac83293fb4d740aedbfdc9240dd8d1a50c1099acd76ce80ce7c7230c7f" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", +] + +[[package]] +name = "zmij" +version = "1.0.23" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "29666d0abbfad1e3dc4dcf6144730dd3a3ab225bbbdac83319345b1b44ccfc1b" diff --git a/rust/cuda-coop/tests/adjacent_difference/Cargo.toml b/rust/cuda-coop/tests/adjacent_difference/Cargo.toml new file mode 100644 index 000000000000..74bffb62bf05 --- /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 = "96caed46c776eb49226137eac849e795b83b7ce0" } +cuda-host = { git = "https://github.com/tpn/cuda-rust.git", rev = "96caed46c776eb49226137eac849e795b83b7ce0", features = ["async"] } +cuda-core = { git = "https://github.com/tpn/cuda-rust.git", rev = "96caed46c776eb49226137eac849e795b83b7ce0" } +cuda-async = { git = "https://github.com/tpn/cuda-rust.git", rev = "96caed46c776eb49226137eac849e795b83b7ce0" } diff --git a/rust/cuda-coop/tests/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..69b7fa259711 --- /dev/null +++ b/rust/cuda-coop/tests/check_load_store.py @@ -0,0 +1,1646 @@ +# 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, + ), +) + +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 == "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..a6275a7a9ac0 --- /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=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "libnvvm-sys", + "nvjitlink-sys", + "serde", + "sha2", + "thiserror 2.0.21", +] + +[[package]] +name = "cuda-async" +version = "0.5.0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "anyhow", + "cuda-bindings", + "cuda-core", + "dashmap", + "futures", + "half", + "once_cell", + "rustc-hash 2.1.3", + "thiserror 1.0.69", +] + +[[package]] +name = "cuda-bindings" +version = "0.5.0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "bindgen", + "libloading", + "prettyplease", + "proc-macro2", + "quote", + "syn 2.0.119", +] + +[[package]] +name = "cuda-coop" +version = "0.0.0" +dependencies = [ + "cuda-coop-codegen", + "cuda-device", +] + +[[package]] +name = "cuda-coop-codegen" +version = "0.0.0" +dependencies = [ + "libloading", + "serde", + "serde_json", + "sha2", +] + +[[package]] +name = "cuda-core" +version = "0.5.0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "anyhow", + "cuda-bindings", + "cuda-core-derive", + "half", + "oxide-artifacts", +] + +[[package]] +name = "cuda-core-derive" +version = "0.5.0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", +] + +[[package]] +name = "cuda-device" +version = "0.2.1" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "cuda-macros", +] + +[[package]] +name = "cuda-host" +version = "0.2.1" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "cuda-artifact-finalizer", + "cuda-async", + "cuda-core", + "cuda-device", + "cuda-macros", + "half", + "libc", + "oxide-artifacts", + "ptx-parse", + "sha2", + "thiserror 2.0.21", +] + +[[package]] +name = "cuda-macros" +version = "0.2.1" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "oxide-tile-export", + "proc-macro2", + "quote", + "reserved-oxide-symbols", + "serde", + "serde_json", + "syn 2.0.119", +] + +[[package]] +name = "cuda-target-spec" +version = "0.1.0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" + +[[package]] +name = "dashmap" +version = "6.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e6361d5c062261c78a176addb82d4c821ae42bed6089de0e12603cd25de2059c" +dependencies = [ + "cfg-if", + "crossbeam-utils", + "hashbrown 0.14.5", + "lock_api", + "once_cell", + "parking_lot_core", +] + +[[package]] +name = "digest" +version = "0.10.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9ed9a281f7bc9b7576e61468ba615a66a5c8cfdff42420a70aa82701a3b1e292" +dependencies = [ + "block-buffer", + "crypto-common", +] + +[[package]] +name = "either" +version = "1.19.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0e9c71c2167ca323c882b99918929403426e2373ea17242ff5653e0d5e1058be" + +[[package]] +name = "equivalent" +version = "1.0.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "00d174d5400e5e8fd687ad1049e2f578285fa914201b1af7e8b112a4546bd826" + +[[package]] +name = "errno" +version = "0.3.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "39cab71617ae0d63f51a36d69f866391735b51691dbda63cf6f96d042b63efeb" +dependencies = [ + "libc", + "windows-sys 0.61.2", +] + +[[package]] +name = "foldhash" +version = "0.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d9c4f5dac5e15c24eb999c26181a6ca40b39fe946cbe4c263c7209467bc83af2" + +[[package]] +name = "futures" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9a31d2a3fbaaeb2af2368bbdd904aa8e812d3c04a1ee10d3171f52d556e5d0a3" +dependencies = [ + "futures-channel", + "futures-core", + "futures-executor", + "futures-io", + "futures-sink", + "futures-task", + "futures-util", +] + +[[package]] +name = "futures-channel" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b1f9e3d69d39e4862ffed03ed071a76f9a13ba1d9109d355b0f0aa6b15e393c4" +dependencies = [ + "futures-core", + "futures-sink", +] + +[[package]] +name = "futures-core" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "92d699e522242e69e3003b94ecc1f960f3a5e015aa7c5d7486e65ad01dd94f5e" + +[[package]] +name = "futures-executor" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "031b47cf1a3c6cc8bc2fc76cd437f521619387907d469316e7c0bc278f1f5432" +dependencies = [ + "futures-core", + "futures-task", + "futures-util", +] + +[[package]] +name = "futures-io" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "53c0fa8157de1303bfffdaa1cc2a673bfffb60102f76b0ef4441659124373fed" + +[[package]] +name = "futures-macro" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9fb9654ba8355388abeb8dcb4fc62f511300867002afc858860463bdd9fe0c44" +dependencies = [ + "proc-macro2", + "quote", + "syn 3.0.9", +] + +[[package]] +name = "futures-sink" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1944426bf7d03f1d14f708785e4b33efd750b36d48a157b836b3efc15ede8e1d" + +[[package]] +name = "futures-task" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cd417de3d1d015fc3bfd2b1ea46dfc7bab72ef86f1cc7cc9c78e728b34a6d1fd" + +[[package]] +name = "futures-util" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0d50a92467f8ba5dd6e3ee5d4bd04d73ab2e4e1c44474a0674821dfce14b79bc" +dependencies = [ + "futures-channel", + "futures-core", + "futures-io", + "futures-macro", + "futures-sink", + "futures-task", + "memchr", + "pin-project-lite", + "slab", +] + +[[package]] +name = "generic-array" +version = "0.14.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "85649ca51fd72272d7821adaf274ad91c288277713d9c18820d8499a7ff69e9a" +dependencies = [ + "typenum", + "version_check", +] + +[[package]] +name = "glob" +version = "0.3.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e4eba85ea1d0a966a983acd07deee566e67395d2d96b6fb39e62b5a833f1eb0b" + +[[package]] +name = "half" +version = "2.7.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" +dependencies = [ + "cfg-if", + "crunchy", + "zerocopy", +] + +[[package]] +name = "hashbrown" +version = "0.14.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e5274423e17b7c9fc20b6e7e208532f9b19825d82dfd615708b70edd83df41f1" + +[[package]] +name = "hashbrown" +version = "0.15.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9229cfe53dfd69f0609a49f65461bd93001ea1ef889cd5529dd176593f5338a1" +dependencies = [ + "foldhash", +] + +[[package]] +name = "hashbrown" +version = "0.17.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ed5909b6e89a2db4456e54cd5f673791d7eca6732202bbf2a9cc504fe2f9b84a" + +[[package]] +name = "home" +version = "0.5.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cc627f471c528ff0c4a49e1d5e60450c8f6461dd6d10ba9dcd3a61d3dff7728d" +dependencies = [ + "windows-sys 0.61.2", +] + +[[package]] +name = "indexmap" +version = "2.14.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cc4e190f5d26ca7051642629da2c52fc03bde85a03197c99408dcd291734c855" +dependencies = [ + "equivalent", + "hashbrown 0.17.1", +] + +[[package]] +name = "itertools" +version = "0.12.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ba291022dbbd398a455acf126c1e341954079855bc60dfdda641363bd6922569" +dependencies = [ + "either", +] + +[[package]] +name = "itoa" +version = "1.0.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8f42a60cbdf9a97f5d2305f08a87dc4e09308d1276d28c869c684d7777685682" + +[[package]] +name = "lazy_static" +version = "1.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "20870f649af7073d53e38067b2a84312175d56ea15217e1b15bc83506ec50afb" + +[[package]] +name = "lazycell" +version = "1.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "830d08ce1d1d941e6b30645f1a0eb5643013d835ce3779a5fc208261dbe10f55" + +[[package]] +name = "libc" +version = "0.2.191" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "65cd87c9a6d60536580b316ad985726104730a5151c99845f2490428ef91240a" + +[[package]] +name = "libloading" +version = "0.8.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d7c4b02199fee7c5d21a5ae7d8cfa79a6ef5bb2fc834d6e9058e89c825efdc55" +dependencies = [ + "cfg-if", + "windows-link", +] + +[[package]] +name = "libnvvm-sys" +version = "0.2.1" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "cuda-target-spec", + "libloading", + "thiserror 2.0.21", +] + +[[package]] +name = "linux-raw-sys" +version = "0.4.15" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d26c52dbd32dccf2d10cac7725f8eae5296885fb5703b261f7d0a0739ec807ab" + +[[package]] +name = "lock_api" +version = "0.4.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "224399e74b87b5f3557511d98dff8b14089b3dadafcab6bb93eab67d3aace965" +dependencies = [ + "scopeguard", +] + +[[package]] +name = "log" +version = "0.4.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9f8bd3e56ce4dfc153cf470fffbfa98c7620958b312ca5c3a4b8d5181fd13c6" + +[[package]] +name = "memchr" +version = "2.8.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cf8baf1c55e62ffcace7a9f06f4bd9cd3f0c4beb022d3b367256b91b87513d98" + +[[package]] +name = "minimal-lexical" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "68354c5c6bd36d73ff3feceb05efa59b6acb7626617f4962be322a825e61f79a" + +[[package]] +name = "nom" +version = "7.1.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d273983c5a657a70a3e8f2a01329822f3b8c8172b73826411a55751e404a0a4a" +dependencies = [ + "memchr", + "minimal-lexical", +] + +[[package]] +name = "nvjitlink-sys" +version = "0.2.1" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "libloading", + "thiserror 2.0.21", +] + +[[package]] +name = "object" +version = "0.36.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "62948e14d923ea95ea2c7c86c71013138b66525b86bdc08d2dcc262bdb497b87" +dependencies = [ + "crc32fast", + "hashbrown 0.15.5", + "indexmap", + "memchr", +] + +[[package]] +name = "once_cell" +version = "1.21.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9f7c3e4beb33f85d45ae3e3a1792185706c8e16d043238c593331cc7cd313b50" + +[[package]] +name = "oxide-artifacts" +version = "0.2.1" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "object", +] + +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "serde", + "serde_json", + "thiserror 2.0.21", +] + +[[package]] +name = "parking_lot_core" +version = "0.9.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2621685985a2ebf1c516881c026032ac7deafcda1a2c9b7850dc81e3dfcb64c1" +dependencies = [ + "cfg-if", + "libc", + "redox_syscall", + "smallvec", + "windows-link", +] + +[[package]] +name = "pin-project-lite" +version = "0.2.17" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a89322df9ebe1c1578d689c92318e070967d1042b512afbe49518723f4e6d5cd" + +[[package]] +name = "prettyplease" +version = "0.2.37" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "479ca8adacdd7ce8f1fb39ce9ecccbfe93a3f1344b3d0d97f20bc0196208f62b" +dependencies = [ + "proc-macro2", + "syn 2.0.119", +] + +[[package]] +name = "proc-macro2" +version = "1.0.107" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "985e7ec9bb745e6ce6535b544d84d6cd6f7ad8bd711c398938ae983b91a766d9" +dependencies = [ + "unicode-ident", +] + +[[package]] +name = "ptx-parse" +version = "0.1.0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" + +[[package]] +name = "quote" +version = "1.0.47" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1fbf4db142a473a8d80c26bbf18454ed458bf8d26c8219c331daecfdbd079001" +dependencies = [ + "proc-macro2", +] + +[[package]] +name = "redox_syscall" +version = "0.5.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ed2bf2547551a7053d6fdfafda3f938979645c44812fbfcda098faae3f1a362d" +dependencies = [ + "bitflags", +] + +[[package]] +name = "regex" +version = "1.13.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f020237b6c8eed93db2e2cb53c00c60a8e1bc73da7d073199a1180401450218d" +dependencies = [ + "aho-corasick", + "memchr", + "regex-automata", + "regex-syntax", +] + +[[package]] +name = "regex-automata" +version = "0.4.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ad8553b9b26413251cbf30e620595c7a41b3887f03da04579c0e6b0d6a06b4b2" +dependencies = [ + "aho-corasick", + "memchr", + "regex-syntax", +] + +[[package]] +name = "regex-syntax" +version = "0.8.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d6f6ff9a378485b298a5286656da665ba74413d36db0979633275d2e708145d4" + +[[package]] +name = "reserved-oxide-symbols" +version = "0.2.1" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" + +[[package]] +name = "rustc-hash" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "08d43f7aa6b08d49f382cde6a7982047c3426db949b1424bc4b7ec9ae12c6ce2" + +[[package]] +name = "rustc-hash" +version = "2.1.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6b1e7f9a428571be2dc5bc0505c13fb6bf936822b894ec87abf8a08a4e51742d" + +[[package]] +name = "rustix" +version = "0.38.44" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fdb5bc1ae2baa591800df16c9ca78619bf65c0488b41b96ccec5d11220d8c154" +dependencies = [ + "bitflags", + "errno", + "libc", + "linux-raw-sys", + "windows-sys 0.59.0", +] + +[[package]] +name = "scopeguard" +version = "1.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "94143f37725109f92c262ed2cf5e59bce7498c01bcc1502d7b9afe439a4e9f49" + +[[package]] +name = "serde" +version = "1.0.229" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4148590afebada386688f18773da617792bf2ef03ffc1e4cbd2b1d45b023e0ba" +dependencies = [ + "serde_core", + "serde_derive", +] + +[[package]] +name = "serde_core" +version = "1.0.229" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "67dca2c9c51e58a4791a4b1ed58308b39c64224d349a935ab5039aa360942a48" +dependencies = [ + "serde_derive", +] + +[[package]] +name = "serde_derive" +version = "1.0.229" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e7a5d71263a5a7d47b41f6b3f06ba276f10cc18b0931f1799f710578e2309348" +dependencies = [ + "proc-macro2", + "quote", + "syn 3.0.9", +] + +[[package]] +name = "serde_json" +version = "1.0.154" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e7e9cc8b1b85264074fbcc02a88680c4096b1e47df8f739dceb03bf482f04bd6" +dependencies = [ + "itoa", + "memchr", + "serde", + "serde_core", + "zmij", +] + +[[package]] +name = "sha2" +version = "0.10.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a7507d819769d01a365ab707794a4084392c824f54a7a6a7862f8c3d0892b283" +dependencies = [ + "cfg-if", + "cpufeatures", + "digest", +] + +[[package]] +name = "shlex" +version = "1.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0fda2ff0d084019ba4d7c6f371c95d8fd75ce3524c3cb8fb653a3023f6323e64" + +[[package]] +name = "slab" +version = "0.4.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0c790de23124f9ab44544d7ac05d60440adc586479ce501c1d6d7da3cd8c9cf5" + +[[package]] +name = "smallvec" +version = "1.16.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5b3dc8af474f516a851ff4bd12db780f948b9250ad37211e4eec0bccea54e01b" + +[[package]] +name = "syn" +version = "2.0.119" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "872831b642d1a07999a962a351ed35b955ea2cfc8f3862091e2a240a84f17297" +dependencies = [ + "proc-macro2", + "quote", + "unicode-ident", +] + +[[package]] +name = "syn" +version = "3.0.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d78c8dee4c7bf0e14673097256fed6142ce9d3b85a408189d07482442145823b" +dependencies = [ + "proc-macro2", + "quote", + "unicode-ident", +] + +[[package]] +name = "thiserror" +version = "1.0.69" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b6aaf5339b578ea85b50e080feb250a3e8ae8cfcdff9a461c9ec2904bc923f52" +dependencies = [ + "thiserror-impl 1.0.69", +] + +[[package]] +name = "thiserror" +version = "2.0.21" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "09e52cb86a36cede5cb101bf8908837b3e4c6e5e59fe7fd85c23fb56200d189e" +dependencies = [ + "thiserror-impl 2.0.21", +] + +[[package]] +name = "thiserror-impl" +version = "1.0.69" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4fee6c4efc90059e10f81e6d42c60a18f76588c3d74cb83a0b242a2b6c7504c1" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", +] + +[[package]] +name = "thiserror-impl" +version = "2.0.21" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fe5197923287db20a58125f0bc85c062f7f2c892de97b18c356f9efb14b28524" +dependencies = [ + "proc-macro2", + "quote", + "syn 3.0.9", +] + +[[package]] +name = "typenum" +version = "1.20.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b6f5e870be6c3b371b77fe0ee0bafb859fa4964b4404c27de1d380043c4dda20" + +[[package]] +name = "unicode-ident" +version = "1.0.27" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a2c754d6c33795a1c324727428e5a7dedb5b06195f9890bdbcba760d3e246563" + +[[package]] +name = "version_check" +version = "0.9.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0b928f33d975fc6ad9f86c8f283853ad26bdd5b10b7f1542aa2fa15e2289105a" + +[[package]] +name = "which" +version = "4.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "87ba24419a2078cd2b0f2ede2691b6c66d8e47836da3b6db8265ebad47afbfc7" +dependencies = [ + "either", + "home", + "once_cell", + "rustix", +] + +[[package]] +name = "windows-link" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f0805222e57f7521d6a62e36fa9163bc891acd422f971defe97d64e70d0a4fe5" + +[[package]] +name = "windows-sys" +version = "0.59.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1e38bc4d79ed67fd075bcc251a1c39b32a1776bbe92e5bef1f0bf1f8c531853b" +dependencies = [ + "windows-targets", +] + +[[package]] +name = "windows-sys" +version = "0.61.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ae137229bcbd6cdf0f7b80a31df61766145077ddf49416a728b02cb3921ff3fc" +dependencies = [ + "windows-link", +] + +[[package]] +name = "windows-targets" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9b724f72796e036ab90c1021d4780d4d3d648aca59e491e6b98e725b84e99973" +dependencies = [ + "windows_aarch64_gnullvm", + "windows_aarch64_msvc", + "windows_i686_gnu", + "windows_i686_gnullvm", + "windows_i686_msvc", + "windows_x86_64_gnu", + "windows_x86_64_gnullvm", + "windows_x86_64_msvc", +] + +[[package]] +name = "windows_aarch64_gnullvm" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "32a4622180e7a0ec044bb555404c800bc9fd9ec262ec147edd5989ccd0c02cd3" + +[[package]] +name = "windows_aarch64_msvc" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "09ec2a7bb152e2252b53fa7803150007879548bc709c039df7627cabbd05d469" + +[[package]] +name = "windows_i686_gnu" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8e9b5ad5ab802e97eb8e295ac6720e509ee4c243f69d781394014ebfe8bbfa0b" + +[[package]] +name = "windows_i686_gnullvm" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0eee52d38c090b3caa76c563b86c3a4bd71ef1a819287c19d586d7334ae8ed66" + +[[package]] +name = "windows_i686_msvc" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "240948bc05c5e7c6dabba28bf89d89ffce3e303022809e73deaefe4f6ec56c66" + +[[package]] +name = "windows_x86_64_gnu" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "147a5c80aabfbf0c7d901cb5895d1de30ef2907eb21fbbab29ca94c5b08b1a78" + +[[package]] +name = "windows_x86_64_gnullvm" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "24d5b23dc417412679681396f2b49f3de8c1473deb516bd34410872eff51ed0d" + +[[package]] +name = "windows_x86_64_msvc" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "589f6da84c646204747d1270a2a5661ea66ed1cced2631d546fdfb155959f9ec" + +[[package]] +name = "zerocopy" +version = "0.8.63" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e5fe1f8f1b06191a00962174c61aa5005e0bb391a6d80d07e24d115c01a92ed8" +dependencies = [ + "zerocopy-derive", +] + +[[package]] +name = "zerocopy-derive" +version = "0.8.63" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "863ad3ac83293fb4d740aedbfdc9240dd8d1a50c1099acd76ce80ce7c7230c7f" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", +] + +[[package]] +name = "zmij" +version = "1.0.23" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "29666d0abbfad1e3dc4dcf6144730dd3a3ab225bbbdac83319345b1b44ccfc1b" diff --git a/rust/cuda-coop/tests/discontinuity/Cargo.toml b/rust/cuda-coop/tests/discontinuity/Cargo.toml new file mode 100644 index 000000000000..1e3d2960d9a8 --- /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 = "96caed46c776eb49226137eac849e795b83b7ce0" } +cuda-host = { git = "https://github.com/tpn/cuda-rust.git", rev = "96caed46c776eb49226137eac849e795b83b7ce0", features = ["async"] } +cuda-core = { git = "https://github.com/tpn/cuda-rust.git", rev = "96caed46c776eb49226137eac849e795b83b7ce0" } +cuda-async = { git = "https://github.com/tpn/cuda-rust.git", rev = "96caed46c776eb49226137eac849e795b83b7ce0" } diff --git a/rust/cuda-coop/tests/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..038fdbb07c68 --- /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=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "libnvvm-sys", + "nvjitlink-sys", + "serde", + "sha2", + "thiserror 2.0.21", +] + +[[package]] +name = "cuda-async" +version = "0.5.0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "anyhow", + "cuda-bindings", + "cuda-core", + "dashmap", + "futures", + "half", + "once_cell", + "rustc-hash 2.1.3", + "thiserror 1.0.69", +] + +[[package]] +name = "cuda-bindings" +version = "0.5.0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "bindgen", + "libloading", + "prettyplease", + "proc-macro2", + "quote", + "syn 2.0.119", +] + +[[package]] +name = "cuda-coop" +version = "0.0.0" +dependencies = [ + "cuda-coop-codegen", + "cuda-device", +] + +[[package]] +name = "cuda-coop-codegen" +version = "0.0.0" +dependencies = [ + "libloading", + "serde", + "serde_json", + "sha2", +] + +[[package]] +name = "cuda-core" +version = "0.5.0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "anyhow", + "cuda-bindings", + "cuda-core-derive", + "half", + "oxide-artifacts", +] + +[[package]] +name = "cuda-core-derive" +version = "0.5.0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", +] + +[[package]] +name = "cuda-device" +version = "0.2.1" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "cuda-macros", +] + +[[package]] +name = "cuda-host" +version = "0.2.1" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "cuda-artifact-finalizer", + "cuda-async", + "cuda-core", + "cuda-device", + "cuda-macros", + "half", + "libc", + "oxide-artifacts", + "ptx-parse", + "sha2", + "thiserror 2.0.21", +] + +[[package]] +name = "cuda-macros" +version = "0.2.1" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "oxide-tile-export", + "proc-macro2", + "quote", + "reserved-oxide-symbols", + "serde", + "serde_json", + "syn 2.0.119", +] + +[[package]] +name = "cuda-target-spec" +version = "0.1.0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" + +[[package]] +name = "dashmap" +version = "6.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e6361d5c062261c78a176addb82d4c821ae42bed6089de0e12603cd25de2059c" +dependencies = [ + "cfg-if", + "crossbeam-utils", + "hashbrown 0.14.5", + "lock_api", + "once_cell", + "parking_lot_core", +] + +[[package]] +name = "digest" +version = "0.10.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9ed9a281f7bc9b7576e61468ba615a66a5c8cfdff42420a70aa82701a3b1e292" +dependencies = [ + "block-buffer", + "crypto-common", +] + +[[package]] +name = "either" +version = "1.19.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0e9c71c2167ca323c882b99918929403426e2373ea17242ff5653e0d5e1058be" + +[[package]] +name = "equivalent" +version = "1.0.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "00d174d5400e5e8fd687ad1049e2f578285fa914201b1af7e8b112a4546bd826" + +[[package]] +name = "errno" +version = "0.3.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "39cab71617ae0d63f51a36d69f866391735b51691dbda63cf6f96d042b63efeb" +dependencies = [ + "libc", + "windows-sys 0.61.2", +] + +[[package]] +name = "foldhash" +version = "0.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d9c4f5dac5e15c24eb999c26181a6ca40b39fe946cbe4c263c7209467bc83af2" + +[[package]] +name = "futures" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9a31d2a3fbaaeb2af2368bbdd904aa8e812d3c04a1ee10d3171f52d556e5d0a3" +dependencies = [ + "futures-channel", + "futures-core", + "futures-executor", + "futures-io", + "futures-sink", + "futures-task", + "futures-util", +] + +[[package]] +name = "futures-channel" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b1f9e3d69d39e4862ffed03ed071a76f9a13ba1d9109d355b0f0aa6b15e393c4" +dependencies = [ + "futures-core", + "futures-sink", +] + +[[package]] +name = "futures-core" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "92d699e522242e69e3003b94ecc1f960f3a5e015aa7c5d7486e65ad01dd94f5e" + +[[package]] +name = "futures-executor" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "031b47cf1a3c6cc8bc2fc76cd437f521619387907d469316e7c0bc278f1f5432" +dependencies = [ + "futures-core", + "futures-task", + "futures-util", +] + +[[package]] +name = "futures-io" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "53c0fa8157de1303bfffdaa1cc2a673bfffb60102f76b0ef4441659124373fed" + +[[package]] +name = "futures-macro" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9fb9654ba8355388abeb8dcb4fc62f511300867002afc858860463bdd9fe0c44" +dependencies = [ + "proc-macro2", + "quote", + "syn 3.0.9", +] + +[[package]] +name = "futures-sink" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1944426bf7d03f1d14f708785e4b33efd750b36d48a157b836b3efc15ede8e1d" + +[[package]] +name = "futures-task" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cd417de3d1d015fc3bfd2b1ea46dfc7bab72ef86f1cc7cc9c78e728b34a6d1fd" + +[[package]] +name = "futures-util" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0d50a92467f8ba5dd6e3ee5d4bd04d73ab2e4e1c44474a0674821dfce14b79bc" +dependencies = [ + "futures-channel", + "futures-core", + "futures-io", + "futures-macro", + "futures-sink", + "futures-task", + "memchr", + "pin-project-lite", + "slab", +] + +[[package]] +name = "generic-array" +version = "0.14.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "85649ca51fd72272d7821adaf274ad91c288277713d9c18820d8499a7ff69e9a" +dependencies = [ + "typenum", + "version_check", +] + +[[package]] +name = "glob" +version = "0.3.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e4eba85ea1d0a966a983acd07deee566e67395d2d96b6fb39e62b5a833f1eb0b" + +[[package]] +name = "half" +version = "2.7.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" +dependencies = [ + "cfg-if", + "crunchy", + "zerocopy", +] + +[[package]] +name = "hashbrown" +version = "0.14.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e5274423e17b7c9fc20b6e7e208532f9b19825d82dfd615708b70edd83df41f1" + +[[package]] +name = "hashbrown" +version = "0.15.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9229cfe53dfd69f0609a49f65461bd93001ea1ef889cd5529dd176593f5338a1" +dependencies = [ + "foldhash", +] + +[[package]] +name = "hashbrown" +version = "0.17.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ed5909b6e89a2db4456e54cd5f673791d7eca6732202bbf2a9cc504fe2f9b84a" + +[[package]] +name = "home" +version = "0.5.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cc627f471c528ff0c4a49e1d5e60450c8f6461dd6d10ba9dcd3a61d3dff7728d" +dependencies = [ + "windows-sys 0.61.2", +] + +[[package]] +name = "indexmap" +version = "2.14.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cc4e190f5d26ca7051642629da2c52fc03bde85a03197c99408dcd291734c855" +dependencies = [ + "equivalent", + "hashbrown 0.17.1", +] + +[[package]] +name = "itertools" +version = "0.12.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ba291022dbbd398a455acf126c1e341954079855bc60dfdda641363bd6922569" +dependencies = [ + "either", +] + +[[package]] +name = "itoa" +version = "1.0.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8f42a60cbdf9a97f5d2305f08a87dc4e09308d1276d28c869c684d7777685682" + +[[package]] +name = "lazy_static" +version = "1.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "20870f649af7073d53e38067b2a84312175d56ea15217e1b15bc83506ec50afb" + +[[package]] +name = "lazycell" +version = "1.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "830d08ce1d1d941e6b30645f1a0eb5643013d835ce3779a5fc208261dbe10f55" + +[[package]] +name = "libc" +version = "0.2.191" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "65cd87c9a6d60536580b316ad985726104730a5151c99845f2490428ef91240a" + +[[package]] +name = "libloading" +version = "0.8.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d7c4b02199fee7c5d21a5ae7d8cfa79a6ef5bb2fc834d6e9058e89c825efdc55" +dependencies = [ + "cfg-if", + "windows-link", +] + +[[package]] +name = "libnvvm-sys" +version = "0.2.1" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "cuda-target-spec", + "libloading", + "thiserror 2.0.21", +] + +[[package]] +name = "linux-raw-sys" +version = "0.4.15" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d26c52dbd32dccf2d10cac7725f8eae5296885fb5703b261f7d0a0739ec807ab" + +[[package]] +name = "lock_api" +version = "0.4.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "224399e74b87b5f3557511d98dff8b14089b3dadafcab6bb93eab67d3aace965" +dependencies = [ + "scopeguard", +] + +[[package]] +name = "log" +version = "0.4.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9f8bd3e56ce4dfc153cf470fffbfa98c7620958b312ca5c3a4b8d5181fd13c6" + +[[package]] +name = "memchr" +version = "2.8.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cf8baf1c55e62ffcace7a9f06f4bd9cd3f0c4beb022d3b367256b91b87513d98" + +[[package]] +name = "minimal-lexical" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "68354c5c6bd36d73ff3feceb05efa59b6acb7626617f4962be322a825e61f79a" + +[[package]] +name = "nom" +version = "7.1.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d273983c5a657a70a3e8f2a01329822f3b8c8172b73826411a55751e404a0a4a" +dependencies = [ + "memchr", + "minimal-lexical", +] + +[[package]] +name = "nvjitlink-sys" +version = "0.2.1" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "libloading", + "thiserror 2.0.21", +] + +[[package]] +name = "object" +version = "0.36.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "62948e14d923ea95ea2c7c86c71013138b66525b86bdc08d2dcc262bdb497b87" +dependencies = [ + "crc32fast", + "hashbrown 0.15.5", + "indexmap", + "memchr", +] + +[[package]] +name = "once_cell" +version = "1.21.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9f7c3e4beb33f85d45ae3e3a1792185706c8e16d043238c593331cc7cd313b50" + +[[package]] +name = "oxide-artifacts" +version = "0.2.1" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "object", +] + +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "serde", + "serde_json", + "thiserror 2.0.21", +] + +[[package]] +name = "parking_lot_core" +version = "0.9.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2621685985a2ebf1c516881c026032ac7deafcda1a2c9b7850dc81e3dfcb64c1" +dependencies = [ + "cfg-if", + "libc", + "redox_syscall", + "smallvec", + "windows-link", +] + +[[package]] +name = "pin-project-lite" +version = "0.2.17" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a89322df9ebe1c1578d689c92318e070967d1042b512afbe49518723f4e6d5cd" + +[[package]] +name = "prettyplease" +version = "0.2.37" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "479ca8adacdd7ce8f1fb39ce9ecccbfe93a3f1344b3d0d97f20bc0196208f62b" +dependencies = [ + "proc-macro2", + "syn 2.0.119", +] + +[[package]] +name = "proc-macro2" +version = "1.0.107" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "985e7ec9bb745e6ce6535b544d84d6cd6f7ad8bd711c398938ae983b91a766d9" +dependencies = [ + "unicode-ident", +] + +[[package]] +name = "ptx-parse" +version = "0.1.0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" + +[[package]] +name = "quote" +version = "1.0.47" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1fbf4db142a473a8d80c26bbf18454ed458bf8d26c8219c331daecfdbd079001" +dependencies = [ + "proc-macro2", +] + +[[package]] +name = "redox_syscall" +version = "0.5.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ed2bf2547551a7053d6fdfafda3f938979645c44812fbfcda098faae3f1a362d" +dependencies = [ + "bitflags", +] + +[[package]] +name = "regex" +version = "1.13.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f020237b6c8eed93db2e2cb53c00c60a8e1bc73da7d073199a1180401450218d" +dependencies = [ + "aho-corasick", + "memchr", + "regex-automata", + "regex-syntax", +] + +[[package]] +name = "regex-automata" +version = "0.4.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ad8553b9b26413251cbf30e620595c7a41b3887f03da04579c0e6b0d6a06b4b2" +dependencies = [ + "aho-corasick", + "memchr", + "regex-syntax", +] + +[[package]] +name = "regex-syntax" +version = "0.8.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d6f6ff9a378485b298a5286656da665ba74413d36db0979633275d2e708145d4" + +[[package]] +name = "reserved-oxide-symbols" +version = "0.2.1" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" + +[[package]] +name = "rustc-hash" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "08d43f7aa6b08d49f382cde6a7982047c3426db949b1424bc4b7ec9ae12c6ce2" + +[[package]] +name = "rustc-hash" +version = "2.1.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6b1e7f9a428571be2dc5bc0505c13fb6bf936822b894ec87abf8a08a4e51742d" + +[[package]] +name = "rustix" +version = "0.38.44" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fdb5bc1ae2baa591800df16c9ca78619bf65c0488b41b96ccec5d11220d8c154" +dependencies = [ + "bitflags", + "errno", + "libc", + "linux-raw-sys", + "windows-sys 0.59.0", +] + +[[package]] +name = "scopeguard" +version = "1.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "94143f37725109f92c262ed2cf5e59bce7498c01bcc1502d7b9afe439a4e9f49" + +[[package]] +name = "serde" +version = "1.0.229" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4148590afebada386688f18773da617792bf2ef03ffc1e4cbd2b1d45b023e0ba" +dependencies = [ + "serde_core", + "serde_derive", +] + +[[package]] +name = "serde_core" +version = "1.0.229" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "67dca2c9c51e58a4791a4b1ed58308b39c64224d349a935ab5039aa360942a48" +dependencies = [ + "serde_derive", +] + +[[package]] +name = "serde_derive" +version = "1.0.229" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e7a5d71263a5a7d47b41f6b3f06ba276f10cc18b0931f1799f710578e2309348" +dependencies = [ + "proc-macro2", + "quote", + "syn 3.0.9", +] + +[[package]] +name = "serde_json" +version = "1.0.154" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e7e9cc8b1b85264074fbcc02a88680c4096b1e47df8f739dceb03bf482f04bd6" +dependencies = [ + "itoa", + "memchr", + "serde", + "serde_core", + "zmij", +] + +[[package]] +name = "sha2" +version = "0.10.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a7507d819769d01a365ab707794a4084392c824f54a7a6a7862f8c3d0892b283" +dependencies = [ + "cfg-if", + "cpufeatures", + "digest", +] + +[[package]] +name = "shlex" +version = "1.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0fda2ff0d084019ba4d7c6f371c95d8fd75ce3524c3cb8fb653a3023f6323e64" + +[[package]] +name = "slab" +version = "0.4.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0c790de23124f9ab44544d7ac05d60440adc586479ce501c1d6d7da3cd8c9cf5" + +[[package]] +name = "smallvec" +version = "1.16.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5b3dc8af474f516a851ff4bd12db780f948b9250ad37211e4eec0bccea54e01b" + +[[package]] +name = "syn" +version = "2.0.119" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "872831b642d1a07999a962a351ed35b955ea2cfc8f3862091e2a240a84f17297" +dependencies = [ + "proc-macro2", + "quote", + "unicode-ident", +] + +[[package]] +name = "syn" +version = "3.0.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d78c8dee4c7bf0e14673097256fed6142ce9d3b85a408189d07482442145823b" +dependencies = [ + "proc-macro2", + "quote", + "unicode-ident", +] + +[[package]] +name = "thiserror" +version = "1.0.69" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b6aaf5339b578ea85b50e080feb250a3e8ae8cfcdff9a461c9ec2904bc923f52" +dependencies = [ + "thiserror-impl 1.0.69", +] + +[[package]] +name = "thiserror" +version = "2.0.21" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "09e52cb86a36cede5cb101bf8908837b3e4c6e5e59fe7fd85c23fb56200d189e" +dependencies = [ + "thiserror-impl 2.0.21", +] + +[[package]] +name = "thiserror-impl" +version = "1.0.69" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4fee6c4efc90059e10f81e6d42c60a18f76588c3d74cb83a0b242a2b6c7504c1" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", +] + +[[package]] +name = "thiserror-impl" +version = "2.0.21" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fe5197923287db20a58125f0bc85c062f7f2c892de97b18c356f9efb14b28524" +dependencies = [ + "proc-macro2", + "quote", + "syn 3.0.9", +] + +[[package]] +name = "typenum" +version = "1.20.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b6f5e870be6c3b371b77fe0ee0bafb859fa4964b4404c27de1d380043c4dda20" + +[[package]] +name = "unicode-ident" +version = "1.0.27" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a2c754d6c33795a1c324727428e5a7dedb5b06195f9890bdbcba760d3e246563" + +[[package]] +name = "version_check" +version = "0.9.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0b928f33d975fc6ad9f86c8f283853ad26bdd5b10b7f1542aa2fa15e2289105a" + +[[package]] +name = "which" +version = "4.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "87ba24419a2078cd2b0f2ede2691b6c66d8e47836da3b6db8265ebad47afbfc7" +dependencies = [ + "either", + "home", + "once_cell", + "rustix", +] + +[[package]] +name = "windows-link" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f0805222e57f7521d6a62e36fa9163bc891acd422f971defe97d64e70d0a4fe5" + +[[package]] +name = "windows-sys" +version = "0.59.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1e38bc4d79ed67fd075bcc251a1c39b32a1776bbe92e5bef1f0bf1f8c531853b" +dependencies = [ + "windows-targets", +] + +[[package]] +name = "windows-sys" +version = "0.61.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ae137229bcbd6cdf0f7b80a31df61766145077ddf49416a728b02cb3921ff3fc" +dependencies = [ + "windows-link", +] + +[[package]] +name = "windows-targets" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9b724f72796e036ab90c1021d4780d4d3d648aca59e491e6b98e725b84e99973" +dependencies = [ + "windows_aarch64_gnullvm", + "windows_aarch64_msvc", + "windows_i686_gnu", + "windows_i686_gnullvm", + "windows_i686_msvc", + "windows_x86_64_gnu", + "windows_x86_64_gnullvm", + "windows_x86_64_msvc", +] + +[[package]] +name = "windows_aarch64_gnullvm" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "32a4622180e7a0ec044bb555404c800bc9fd9ec262ec147edd5989ccd0c02cd3" + +[[package]] +name = "windows_aarch64_msvc" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "09ec2a7bb152e2252b53fa7803150007879548bc709c039df7627cabbd05d469" + +[[package]] +name = "windows_i686_gnu" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8e9b5ad5ab802e97eb8e295ac6720e509ee4c243f69d781394014ebfe8bbfa0b" + +[[package]] +name = "windows_i686_gnullvm" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0eee52d38c090b3caa76c563b86c3a4bd71ef1a819287c19d586d7334ae8ed66" + +[[package]] +name = "windows_i686_msvc" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "240948bc05c5e7c6dabba28bf89d89ffce3e303022809e73deaefe4f6ec56c66" + +[[package]] +name = "windows_x86_64_gnu" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "147a5c80aabfbf0c7d901cb5895d1de30ef2907eb21fbbab29ca94c5b08b1a78" + +[[package]] +name = "windows_x86_64_gnullvm" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "24d5b23dc417412679681396f2b49f3de8c1473deb516bd34410872eff51ed0d" + +[[package]] +name = "windows_x86_64_msvc" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "589f6da84c646204747d1270a2a5661ea66ed1cced2631d546fdfb155959f9ec" + +[[package]] +name = "zerocopy" +version = "0.8.63" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e5fe1f8f1b06191a00962174c61aa5005e0bb391a6d80d07e24d115c01a92ed8" +dependencies = [ + "zerocopy-derive", +] + +[[package]] +name = "zerocopy-derive" +version = "0.8.63" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "863ad3ac83293fb4d740aedbfdc9240dd8d1a50c1099acd76ce80ce7c7230c7f" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", +] + +[[package]] +name = "zmij" +version = "1.0.23" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "29666d0abbfad1e3dc4dcf6144730dd3a3ab225bbbdac83319345b1b44ccfc1b" diff --git a/rust/cuda-coop/tests/exchange/Cargo.toml b/rust/cuda-coop/tests/exchange/Cargo.toml new file mode 100644 index 000000000000..881a367c4c17 --- /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 = "96caed46c776eb49226137eac849e795b83b7ce0" } +cuda-host = { git = "https://github.com/tpn/cuda-rust.git", rev = "96caed46c776eb49226137eac849e795b83b7ce0", features = ["async"] } +cuda-core = { git = "https://github.com/tpn/cuda-rust.git", rev = "96caed46c776eb49226137eac849e795b83b7ce0" } +cuda-async = { git = "https://github.com/tpn/cuda-rust.git", rev = "96caed46c776eb49226137eac849e795b83b7ce0" } diff --git a/rust/cuda-coop/tests/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/load_store/.gitignore b/rust/cuda-coop/tests/load_store/.gitignore new file mode 100644 index 000000000000..74bffa8af4a4 --- /dev/null +++ b/rust/cuda-coop/tests/load_store/.gitignore @@ -0,0 +1,7 @@ +# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. +# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception +/target/ +/.oxide-artifacts/ +/*.ll +/*.options +/*.target diff --git a/rust/cuda-coop/tests/load_store/Cargo.lock b/rust/cuda-coop/tests/load_store/Cargo.lock new file mode 100644 index 000000000000..514e255e8636 --- /dev/null +++ b/rust/cuda-coop/tests/load_store/Cargo.lock @@ -0,0 +1,1011 @@ +# This file is automatically @generated by Cargo. +# It is not intended for manual editing. +version = 4 + +[[package]] +name = "aho-corasick" +version = "1.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c982642fa9e8606056828ee9a8505737230110bb1099153c79efe865c59d12ba" +dependencies = [ + "memchr", +] + +[[package]] +name = "anyhow" +version = "1.0.104" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "330a5ed07fa54e4702c9d6c4174f74427fc0ef6e214bbd677ae50a5099946470" + +[[package]] +name = "bindgen" +version = "0.69.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "271383c67ccabffb7381723dea0672a673f292304fcb45c01cc648c7a8d58088" +dependencies = [ + "bitflags", + "cexpr", + "clang-sys", + "itertools", + "lazy_static", + "lazycell", + "log", + "prettyplease", + "proc-macro2", + "quote", + "regex", + "rustc-hash 1.1.0", + "shlex", + "syn 2.0.119", + "which", +] + +[[package]] +name = "bitflags" +version = "2.13.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3ded4057c258ba199e2d26386d3af3780957ecaee6c4ef4041c6b4b8b97c0b06" + +[[package]] +name = "block-buffer" +version = "0.10.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3078c7629b62d3f0439517fa394996acacc5cbc91c5a20d8c658e77abd503a71" +dependencies = [ + "generic-array", +] + +[[package]] +name = "cexpr" +version = "0.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6fac387a98bb7c37292057cffc56d62ecb629900026402633ae9160df93a8766" +dependencies = [ + "nom", +] + +[[package]] +name = "cfg-if" +version = "1.0.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4e7648175b45a9a48536d676f68d918270699102aa8dab5496df06904c914600" + +[[package]] +name = "clang-sys" +version = "1.9.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "157a8ba7b480713b56f4c09fd13fc3e0a22a5dfab8097ba61cbc5feef950788a" +dependencies = [ + "glob", + "libc", + "libloading", +] + +[[package]] +name = "coop-load-store-tests" +version = "0.0.0" +dependencies = [ + "cuda-async", + "cuda-coop", + "cuda-core", + "cuda-device", + "cuda-host", +] + +[[package]] +name = "cpufeatures" +version = "0.2.17" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "59ed5838eebb26a2bb2e58f6d5b5316989ae9d08bab10e0e6d103e656d1b0280" +dependencies = [ + "libc", +] + +[[package]] +name = "crc32fast" +version = "1.5.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0ba106e844045ddc02c20768ac4b0b7cb45f2cc2ecbd833ac9f5dacb183c2055" +dependencies = [ + "cfg-if", +] + +[[package]] +name = "crossbeam-utils" +version = "0.8.23" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a31eee39dddec8330830986fcd7625edb5a24ec90ea038215273bbc3adb08ac6" + +[[package]] +name = "crunchy" +version = "0.2.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "460fbee9c2c2f33933d720630a6a0bac33ba7053db5344fac858d4b8952d77d5" + +[[package]] +name = "crypto-common" +version = "0.1.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "78c8292055d1c1df0cce5d180393dc8cce0abec0a7102adb6c7b1eef6016d60a" +dependencies = [ + "generic-array", + "typenum", +] + +[[package]] +name = "cuda-artifact-finalizer" +version = "0.2.1" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "libnvvm-sys", + "nvjitlink-sys", + "serde", + "sha2", + "thiserror 2.0.21", +] + +[[package]] +name = "cuda-async" +version = "0.5.0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "anyhow", + "cuda-bindings", + "cuda-core", + "dashmap", + "futures", + "half", + "once_cell", + "rustc-hash 2.1.3", + "thiserror 1.0.69", +] + +[[package]] +name = "cuda-bindings" +version = "0.5.0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "bindgen", + "libloading", + "prettyplease", + "proc-macro2", + "quote", + "syn 2.0.119", +] + +[[package]] +name = "cuda-coop" +version = "0.0.0" +dependencies = [ + "cuda-coop-codegen", + "cuda-device", +] + +[[package]] +name = "cuda-coop-codegen" +version = "0.0.0" +dependencies = [ + "libloading", + "serde", + "serde_json", + "sha2", +] + +[[package]] +name = "cuda-core" +version = "0.5.0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "anyhow", + "cuda-bindings", + "cuda-core-derive", + "half", + "oxide-artifacts", +] + +[[package]] +name = "cuda-core-derive" +version = "0.5.0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", +] + +[[package]] +name = "cuda-device" +version = "0.2.1" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "cuda-macros", +] + +[[package]] +name = "cuda-host" +version = "0.2.1" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "cuda-artifact-finalizer", + "cuda-async", + "cuda-core", + "cuda-device", + "cuda-macros", + "half", + "libc", + "oxide-artifacts", + "ptx-parse", + "sha2", + "thiserror 2.0.21", +] + +[[package]] +name = "cuda-macros" +version = "0.2.1" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "oxide-tile-export", + "proc-macro2", + "quote", + "reserved-oxide-symbols", + "serde", + "serde_json", + "syn 2.0.119", +] + +[[package]] +name = "cuda-target-spec" +version = "0.1.0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" + +[[package]] +name = "dashmap" +version = "6.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e6361d5c062261c78a176addb82d4c821ae42bed6089de0e12603cd25de2059c" +dependencies = [ + "cfg-if", + "crossbeam-utils", + "hashbrown 0.14.5", + "lock_api", + "once_cell", + "parking_lot_core", +] + +[[package]] +name = "digest" +version = "0.10.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9ed9a281f7bc9b7576e61468ba615a66a5c8cfdff42420a70aa82701a3b1e292" +dependencies = [ + "block-buffer", + "crypto-common", +] + +[[package]] +name = "either" +version = "1.19.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0e9c71c2167ca323c882b99918929403426e2373ea17242ff5653e0d5e1058be" + +[[package]] +name = "equivalent" +version = "1.0.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "00d174d5400e5e8fd687ad1049e2f578285fa914201b1af7e8b112a4546bd826" + +[[package]] +name = "errno" +version = "0.3.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "39cab71617ae0d63f51a36d69f866391735b51691dbda63cf6f96d042b63efeb" +dependencies = [ + "libc", + "windows-sys 0.61.2", +] + +[[package]] +name = "foldhash" +version = "0.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d9c4f5dac5e15c24eb999c26181a6ca40b39fe946cbe4c263c7209467bc83af2" + +[[package]] +name = "futures" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9a31d2a3fbaaeb2af2368bbdd904aa8e812d3c04a1ee10d3171f52d556e5d0a3" +dependencies = [ + "futures-channel", + "futures-core", + "futures-executor", + "futures-io", + "futures-sink", + "futures-task", + "futures-util", +] + +[[package]] +name = "futures-channel" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b1f9e3d69d39e4862ffed03ed071a76f9a13ba1d9109d355b0f0aa6b15e393c4" +dependencies = [ + "futures-core", + "futures-sink", +] + +[[package]] +name = "futures-core" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "92d699e522242e69e3003b94ecc1f960f3a5e015aa7c5d7486e65ad01dd94f5e" + +[[package]] +name = "futures-executor" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "031b47cf1a3c6cc8bc2fc76cd437f521619387907d469316e7c0bc278f1f5432" +dependencies = [ + "futures-core", + "futures-task", + "futures-util", +] + +[[package]] +name = "futures-io" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "53c0fa8157de1303bfffdaa1cc2a673bfffb60102f76b0ef4441659124373fed" + +[[package]] +name = "futures-macro" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9fb9654ba8355388abeb8dcb4fc62f511300867002afc858860463bdd9fe0c44" +dependencies = [ + "proc-macro2", + "quote", + "syn 3.0.9", +] + +[[package]] +name = "futures-sink" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1944426bf7d03f1d14f708785e4b33efd750b36d48a157b836b3efc15ede8e1d" + +[[package]] +name = "futures-task" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cd417de3d1d015fc3bfd2b1ea46dfc7bab72ef86f1cc7cc9c78e728b34a6d1fd" + +[[package]] +name = "futures-util" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0d50a92467f8ba5dd6e3ee5d4bd04d73ab2e4e1c44474a0674821dfce14b79bc" +dependencies = [ + "futures-channel", + "futures-core", + "futures-io", + "futures-macro", + "futures-sink", + "futures-task", + "memchr", + "pin-project-lite", + "slab", +] + +[[package]] +name = "generic-array" +version = "0.14.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "85649ca51fd72272d7821adaf274ad91c288277713d9c18820d8499a7ff69e9a" +dependencies = [ + "typenum", + "version_check", +] + +[[package]] +name = "glob" +version = "0.3.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e4eba85ea1d0a966a983acd07deee566e67395d2d96b6fb39e62b5a833f1eb0b" + +[[package]] +name = "half" +version = "2.7.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" +dependencies = [ + "cfg-if", + "crunchy", + "zerocopy", +] + +[[package]] +name = "hashbrown" +version = "0.14.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e5274423e17b7c9fc20b6e7e208532f9b19825d82dfd615708b70edd83df41f1" + +[[package]] +name = "hashbrown" +version = "0.15.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9229cfe53dfd69f0609a49f65461bd93001ea1ef889cd5529dd176593f5338a1" +dependencies = [ + "foldhash", +] + +[[package]] +name = "hashbrown" +version = "0.17.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ed5909b6e89a2db4456e54cd5f673791d7eca6732202bbf2a9cc504fe2f9b84a" + +[[package]] +name = "home" +version = "0.5.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cc627f471c528ff0c4a49e1d5e60450c8f6461dd6d10ba9dcd3a61d3dff7728d" +dependencies = [ + "windows-sys 0.61.2", +] + +[[package]] +name = "indexmap" +version = "2.14.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cc4e190f5d26ca7051642629da2c52fc03bde85a03197c99408dcd291734c855" +dependencies = [ + "equivalent", + "hashbrown 0.17.1", +] + +[[package]] +name = "itertools" +version = "0.12.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ba291022dbbd398a455acf126c1e341954079855bc60dfdda641363bd6922569" +dependencies = [ + "either", +] + +[[package]] +name = "itoa" +version = "1.0.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8f42a60cbdf9a97f5d2305f08a87dc4e09308d1276d28c869c684d7777685682" + +[[package]] +name = "lazy_static" +version = "1.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "20870f649af7073d53e38067b2a84312175d56ea15217e1b15bc83506ec50afb" + +[[package]] +name = "lazycell" +version = "1.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "830d08ce1d1d941e6b30645f1a0eb5643013d835ce3779a5fc208261dbe10f55" + +[[package]] +name = "libc" +version = "0.2.191" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "65cd87c9a6d60536580b316ad985726104730a5151c99845f2490428ef91240a" + +[[package]] +name = "libloading" +version = "0.8.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d7c4b02199fee7c5d21a5ae7d8cfa79a6ef5bb2fc834d6e9058e89c825efdc55" +dependencies = [ + "cfg-if", + "windows-link", +] + +[[package]] +name = "libnvvm-sys" +version = "0.2.1" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "cuda-target-spec", + "libloading", + "thiserror 2.0.21", +] + +[[package]] +name = "linux-raw-sys" +version = "0.4.15" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d26c52dbd32dccf2d10cac7725f8eae5296885fb5703b261f7d0a0739ec807ab" + +[[package]] +name = "lock_api" +version = "0.4.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "224399e74b87b5f3557511d98dff8b14089b3dadafcab6bb93eab67d3aace965" +dependencies = [ + "scopeguard", +] + +[[package]] +name = "log" +version = "0.4.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9f8bd3e56ce4dfc153cf470fffbfa98c7620958b312ca5c3a4b8d5181fd13c6" + +[[package]] +name = "memchr" +version = "2.8.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cf8baf1c55e62ffcace7a9f06f4bd9cd3f0c4beb022d3b367256b91b87513d98" + +[[package]] +name = "minimal-lexical" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "68354c5c6bd36d73ff3feceb05efa59b6acb7626617f4962be322a825e61f79a" + +[[package]] +name = "nom" +version = "7.1.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d273983c5a657a70a3e8f2a01329822f3b8c8172b73826411a55751e404a0a4a" +dependencies = [ + "memchr", + "minimal-lexical", +] + +[[package]] +name = "nvjitlink-sys" +version = "0.2.1" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "libloading", + "thiserror 2.0.21", +] + +[[package]] +name = "object" +version = "0.36.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "62948e14d923ea95ea2c7c86c71013138b66525b86bdc08d2dcc262bdb497b87" +dependencies = [ + "crc32fast", + "hashbrown 0.15.5", + "indexmap", + "memchr", +] + +[[package]] +name = "once_cell" +version = "1.21.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9f7c3e4beb33f85d45ae3e3a1792185706c8e16d043238c593331cc7cd313b50" + +[[package]] +name = "oxide-artifacts" +version = "0.2.1" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "object", +] + +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "serde", + "serde_json", + "thiserror 2.0.21", +] + +[[package]] +name = "parking_lot_core" +version = "0.9.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2621685985a2ebf1c516881c026032ac7deafcda1a2c9b7850dc81e3dfcb64c1" +dependencies = [ + "cfg-if", + "libc", + "redox_syscall", + "smallvec", + "windows-link", +] + +[[package]] +name = "pin-project-lite" +version = "0.2.17" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a89322df9ebe1c1578d689c92318e070967d1042b512afbe49518723f4e6d5cd" + +[[package]] +name = "prettyplease" +version = "0.2.37" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "479ca8adacdd7ce8f1fb39ce9ecccbfe93a3f1344b3d0d97f20bc0196208f62b" +dependencies = [ + "proc-macro2", + "syn 2.0.119", +] + +[[package]] +name = "proc-macro2" +version = "1.0.107" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "985e7ec9bb745e6ce6535b544d84d6cd6f7ad8bd711c398938ae983b91a766d9" +dependencies = [ + "unicode-ident", +] + +[[package]] +name = "ptx-parse" +version = "0.1.0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" + +[[package]] +name = "quote" +version = "1.0.47" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1fbf4db142a473a8d80c26bbf18454ed458bf8d26c8219c331daecfdbd079001" +dependencies = [ + "proc-macro2", +] + +[[package]] +name = "redox_syscall" +version = "0.5.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ed2bf2547551a7053d6fdfafda3f938979645c44812fbfcda098faae3f1a362d" +dependencies = [ + "bitflags", +] + +[[package]] +name = "regex" +version = "1.13.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f020237b6c8eed93db2e2cb53c00c60a8e1bc73da7d073199a1180401450218d" +dependencies = [ + "aho-corasick", + "memchr", + "regex-automata", + "regex-syntax", +] + +[[package]] +name = "regex-automata" +version = "0.4.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ad8553b9b26413251cbf30e620595c7a41b3887f03da04579c0e6b0d6a06b4b2" +dependencies = [ + "aho-corasick", + "memchr", + "regex-syntax", +] + +[[package]] +name = "regex-syntax" +version = "0.8.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d6f6ff9a378485b298a5286656da665ba74413d36db0979633275d2e708145d4" + +[[package]] +name = "reserved-oxide-symbols" +version = "0.2.1" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" + +[[package]] +name = "rustc-hash" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "08d43f7aa6b08d49f382cde6a7982047c3426db949b1424bc4b7ec9ae12c6ce2" + +[[package]] +name = "rustc-hash" +version = "2.1.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6b1e7f9a428571be2dc5bc0505c13fb6bf936822b894ec87abf8a08a4e51742d" + +[[package]] +name = "rustix" +version = "0.38.44" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fdb5bc1ae2baa591800df16c9ca78619bf65c0488b41b96ccec5d11220d8c154" +dependencies = [ + "bitflags", + "errno", + "libc", + "linux-raw-sys", + "windows-sys 0.59.0", +] + +[[package]] +name = "scopeguard" +version = "1.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "94143f37725109f92c262ed2cf5e59bce7498c01bcc1502d7b9afe439a4e9f49" + +[[package]] +name = "serde" +version = "1.0.229" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4148590afebada386688f18773da617792bf2ef03ffc1e4cbd2b1d45b023e0ba" +dependencies = [ + "serde_core", + "serde_derive", +] + +[[package]] +name = "serde_core" +version = "1.0.229" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "67dca2c9c51e58a4791a4b1ed58308b39c64224d349a935ab5039aa360942a48" +dependencies = [ + "serde_derive", +] + +[[package]] +name = "serde_derive" +version = "1.0.229" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e7a5d71263a5a7d47b41f6b3f06ba276f10cc18b0931f1799f710578e2309348" +dependencies = [ + "proc-macro2", + "quote", + "syn 3.0.9", +] + +[[package]] +name = "serde_json" +version = "1.0.154" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e7e9cc8b1b85264074fbcc02a88680c4096b1e47df8f739dceb03bf482f04bd6" +dependencies = [ + "itoa", + "memchr", + "serde", + "serde_core", + "zmij", +] + +[[package]] +name = "sha2" +version = "0.10.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a7507d819769d01a365ab707794a4084392c824f54a7a6a7862f8c3d0892b283" +dependencies = [ + "cfg-if", + "cpufeatures", + "digest", +] + +[[package]] +name = "shlex" +version = "1.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0fda2ff0d084019ba4d7c6f371c95d8fd75ce3524c3cb8fb653a3023f6323e64" + +[[package]] +name = "slab" +version = "0.4.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0c790de23124f9ab44544d7ac05d60440adc586479ce501c1d6d7da3cd8c9cf5" + +[[package]] +name = "smallvec" +version = "1.16.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5b3dc8af474f516a851ff4bd12db780f948b9250ad37211e4eec0bccea54e01b" + +[[package]] +name = "syn" +version = "2.0.119" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "872831b642d1a07999a962a351ed35b955ea2cfc8f3862091e2a240a84f17297" +dependencies = [ + "proc-macro2", + "quote", + "unicode-ident", +] + +[[package]] +name = "syn" +version = "3.0.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d78c8dee4c7bf0e14673097256fed6142ce9d3b85a408189d07482442145823b" +dependencies = [ + "proc-macro2", + "quote", + "unicode-ident", +] + +[[package]] +name = "thiserror" +version = "1.0.69" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b6aaf5339b578ea85b50e080feb250a3e8ae8cfcdff9a461c9ec2904bc923f52" +dependencies = [ + "thiserror-impl 1.0.69", +] + +[[package]] +name = "thiserror" +version = "2.0.21" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "09e52cb86a36cede5cb101bf8908837b3e4c6e5e59fe7fd85c23fb56200d189e" +dependencies = [ + "thiserror-impl 2.0.21", +] + +[[package]] +name = "thiserror-impl" +version = "1.0.69" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4fee6c4efc90059e10f81e6d42c60a18f76588c3d74cb83a0b242a2b6c7504c1" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", +] + +[[package]] +name = "thiserror-impl" +version = "2.0.21" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fe5197923287db20a58125f0bc85c062f7f2c892de97b18c356f9efb14b28524" +dependencies = [ + "proc-macro2", + "quote", + "syn 3.0.9", +] + +[[package]] +name = "typenum" +version = "1.20.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b6f5e870be6c3b371b77fe0ee0bafb859fa4964b4404c27de1d380043c4dda20" + +[[package]] +name = "unicode-ident" +version = "1.0.27" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a2c754d6c33795a1c324727428e5a7dedb5b06195f9890bdbcba760d3e246563" + +[[package]] +name = "version_check" +version = "0.9.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0b928f33d975fc6ad9f86c8f283853ad26bdd5b10b7f1542aa2fa15e2289105a" + +[[package]] +name = "which" +version = "4.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "87ba24419a2078cd2b0f2ede2691b6c66d8e47836da3b6db8265ebad47afbfc7" +dependencies = [ + "either", + "home", + "once_cell", + "rustix", +] + +[[package]] +name = "windows-link" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f0805222e57f7521d6a62e36fa9163bc891acd422f971defe97d64e70d0a4fe5" + +[[package]] +name = "windows-sys" +version = "0.59.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1e38bc4d79ed67fd075bcc251a1c39b32a1776bbe92e5bef1f0bf1f8c531853b" +dependencies = [ + "windows-targets", +] + +[[package]] +name = "windows-sys" +version = "0.61.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ae137229bcbd6cdf0f7b80a31df61766145077ddf49416a728b02cb3921ff3fc" +dependencies = [ + "windows-link", +] + +[[package]] +name = "windows-targets" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9b724f72796e036ab90c1021d4780d4d3d648aca59e491e6b98e725b84e99973" +dependencies = [ + "windows_aarch64_gnullvm", + "windows_aarch64_msvc", + "windows_i686_gnu", + "windows_i686_gnullvm", + "windows_i686_msvc", + "windows_x86_64_gnu", + "windows_x86_64_gnullvm", + "windows_x86_64_msvc", +] + +[[package]] +name = "windows_aarch64_gnullvm" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "32a4622180e7a0ec044bb555404c800bc9fd9ec262ec147edd5989ccd0c02cd3" + +[[package]] +name = "windows_aarch64_msvc" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "09ec2a7bb152e2252b53fa7803150007879548bc709c039df7627cabbd05d469" + +[[package]] +name = "windows_i686_gnu" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8e9b5ad5ab802e97eb8e295ac6720e509ee4c243f69d781394014ebfe8bbfa0b" + +[[package]] +name = "windows_i686_gnullvm" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0eee52d38c090b3caa76c563b86c3a4bd71ef1a819287c19d586d7334ae8ed66" + +[[package]] +name = "windows_i686_msvc" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "240948bc05c5e7c6dabba28bf89d89ffce3e303022809e73deaefe4f6ec56c66" + +[[package]] +name = "windows_x86_64_gnu" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "147a5c80aabfbf0c7d901cb5895d1de30ef2907eb21fbbab29ca94c5b08b1a78" + +[[package]] +name = "windows_x86_64_gnullvm" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "24d5b23dc417412679681396f2b49f3de8c1473deb516bd34410872eff51ed0d" + +[[package]] +name = "windows_x86_64_msvc" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "589f6da84c646204747d1270a2a5661ea66ed1cced2631d546fdfb155959f9ec" + +[[package]] +name = "zerocopy" +version = "0.8.63" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e5fe1f8f1b06191a00962174c61aa5005e0bb391a6d80d07e24d115c01a92ed8" +dependencies = [ + "zerocopy-derive", +] + +[[package]] +name = "zerocopy-derive" +version = "0.8.63" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "863ad3ac83293fb4d740aedbfdc9240dd8d1a50c1099acd76ce80ce7c7230c7f" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", +] + +[[package]] +name = "zmij" +version = "1.0.23" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "29666d0abbfad1e3dc4dcf6144730dd3a3ab225bbbdac83319345b1b44ccfc1b" diff --git a/rust/cuda-coop/tests/load_store/Cargo.toml b/rust/cuda-coop/tests/load_store/Cargo.toml new file mode 100644 index 000000000000..626172bef415 --- /dev/null +++ b/rust/cuda-coop/tests/load_store/Cargo.toml @@ -0,0 +1,17 @@ +# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. +# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +[package] +name = "coop-load-store-tests" +version = "0.0.0" +edition = "2024" +publish = false + +[workspace] + +[dependencies] +cuda-coop = { path = "../.." } +cuda-device = { git = "https://github.com/tpn/cuda-rust.git", rev = "96caed46c776eb49226137eac849e795b83b7ce0" } +cuda-host = { git = "https://github.com/tpn/cuda-rust.git", rev = "96caed46c776eb49226137eac849e795b83b7ce0", features = ["async"] } +cuda-core = { git = "https://github.com/tpn/cuda-rust.git", rev = "96caed46c776eb49226137eac849e795b83b7ce0" } +cuda-async = { git = "https://github.com/tpn/cuda-rust.git", rev = "96caed46c776eb49226137eac849e795b83b7ce0" } diff --git a/rust/cuda-coop/tests/load_store/src/block_warp.rs b/rust/cuda-coop/tests/load_store/src/block_warp.rs new file mode 100644 index 000000000000..42e058a840b5 --- /dev/null +++ b/rust/cuda-coop/tests/load_store/src/block_warp.rs @@ -0,0 +1,366 @@ +// Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +//! Full-block warp exchange observes every payload slot before and after Store. +use cuda_coop as coop; +use cuda_device::{CollectiveTiles, DisjointSlice, kernel, launch_contract, thread}; +use cuda_host::cuda_module; + +#[cuda_module] +mod kernels { + use super::*; + + fn expected( + input: &[T], + offset: usize, + index: usize, + items: usize, + valid: usize, + default: T, + mode: u32, + ) -> Option { + if mode == 1 || mode == 5 || mode == 6 { + if index < valid { + Some(input[offset + index]) + } else { + None + } + } else if mode == 2 || mode == 4 { + if index < valid { + Some(input[offset + index]) + } else { + Some(default) + } + } else if mode == 3 { + let lane = index / items; + let item = index % items; + if index < valid || (lane & 1 == 0 && item & 1 == 1) || (lane & 1 == 1 && item & 1 == 0) + { + Some(input[offset + index]) + } else { + None + } + } else { + Some(input[offset + index]) + } + } + + #[inline(never)] + fn exercise< + T: coop::Numeric + PartialEq, + const ITEMS: usize, + const CAPACITY: usize, + A: coop::Algorithm, + >( + input: &[T], + output: &mut DisjointSlice>, + valid: usize, + default: T, + mode: u32, + algorithm: A, + ) { + let block = coop::this_block(); + let x = thread::blockDim_x() as usize; + let y = thread::blockDim_y() as usize; + let threads = x * y * thread::blockDim_z() as usize; + let lane = thread::threadIdx_x() as usize + + x * (thread::threadIdx_y() as usize + y * thread::threadIdx_z() as usize); + let group = thread::blockIdx_x() as usize + + thread::gridDim_x() as usize + * (thread::blockIdx_y() as usize + + thread::gridDim_y() as usize * thread::blockIdx_z() as usize); + let offset = group * threads * CAPACITY; + let mut values = coop::ThreadData::<_, ITEMS>::new(); + if mode != 9 { + coop::load( + block, + input, + &mut values, + coop::LoadOptions::default() + .offset(offset) + .algorithm(algorithm), + ); + } + if mode == 1 || mode == 5 || mode == 6 { + coop::load( + block, + input, + &mut values, + coop::LoadOptions::default() + .offset(offset) + .valid_items(valid) + .algorithm(algorithm), + ); + } else if mode == 2 || mode == 4 { + coop::load( + block, + input, + &mut values, + coop::LoadOptions::default() + .offset(offset) + .valid_items(valid) + .oob_default(default) + .algorithm(algorithm), + ); + } else if mode == 3 { + let incoming = values; + values = coop::ThreadData::new(); + let mut item = 0; + while item < ITEMS { + let index = lane * ITEMS + item; + let initialize = index < valid + || (lane & 1 == 0 && item & 1 == 1) + || (lane & 1 == 1 && item & 1 == 0); + if let (true, Some(value)) = (initialize, incoming.get(item)) { + values.write(item, value); + } + item += 1; + } + } + let mut correct = true; + let mut item = 0; + while item < ITEMS { + correct &= values.get(item) + == expected( + input, + offset, + lane * ITEMS + item, + ITEMS, + valid, + default, + mode, + ); + item += 1; + } + if mode == 0 || mode == 2 { + coop::store( + block, + &mut *output, + &mut values, + coop::StoreOptions::default() + .offset(offset) + .algorithm(algorithm), + ); + } else { + let written = if mode == 5 { + valid + 1 + } else if mode == 6 { + valid / 2 + } else { + valid + }; + coop::store( + block, + &mut *output, + &mut values, + coop::StoreOptions::default() + .offset(offset) + .valid_items(written) + .algorithm(algorithm), + ); + } + item = 0; + while item < ITEMS { + let source = (lane / 32) * (32 * ITEMS) + 32 * item + lane % 32; + correct &= + values.get(item) == expected(input, offset, source, ITEMS, valid, default, mode); + item += 1; + } + // Managed scratch reuse inserts completion before this second operation. + // A zero count still permutes all values and bits, with no global writes. + coop::store( + block, + &mut *output, + &mut values, + coop::StoreOptions::default() + .offset(offset) + .valid_items(0) + .algorithm(algorithm), + ); + item = 0; + while item < ITEMS { + let first = (lane / 32) * (32 * ITEMS) + 32 * item + lane % 32; + let source = + ((first / ITEMS) / 32) * (32 * ITEMS) + 32 * (first % ITEMS) + (first / ITEMS) % 32; + correct &= + values.get(item) == expected(input, offset, source, ITEMS, valid, default, mode); + item += 1; + } + if mode == 4 { + thread::sync_threads(); + coop::load( + block, + &mut *output, + &mut values, + coop::LoadOptions::default() + .offset(offset) + .valid_items(valid) + .oob_default(default) + .algorithm(algorithm), + ); + coop::store( + block, + &mut *output, + &mut values, + coop::StoreOptions::default() + .offset(offset) + .valid_items(valid) + .algorithm(algorithm), + ); + item = 0; + while item < ITEMS { + correct &= values.get(item) + == expected( + input, + offset, + (lane / 32) * (32 * ITEMS) + 32 * item + lane % 32, + ITEMS, + valid, + default, + mode, + ); + item += 1; + } + } + assert!(correct); + } + #[kernel] + #[launch_contract(domain=3,block=(8,4,2),dynamic_shared=inferred)] + pub fn block_warp( + input: &[i32], + mut output: DisjointSlice>, + valid: usize, + mode: u32, + ) { + exercise::<_, ITEMS_PER_THREAD, 4, _>( + input, + &mut output, + valid, + -17, + mode, + coop::WarpTranspose, + ); + } + #[kernel] + #[launch_contract(domain=3,block=(8,4,2),dynamic_shared=inferred)] + pub fn block_warp_timesliced( + input: &[i32], + mut output: DisjointSlice>, + valid: usize, + mode: u32, + ) { + exercise::<_, ITEMS_PER_THREAD, 4, _>( + input, + &mut output, + valid, + -17, + mode, + coop::WarpTransposeTimesliced, + ); + } +} + +pub fn run(trap: Option<&str>) -> Result<(), Box> { + use cuda_async::simt::device_context::{init_device_contexts, with_cuda_context}; + use cuda_async::simt::device_operation::DeviceOperation; + use cuda_core::{DeviceBuffer, LaunchConfig3D}; + init_device_contexts(0, 1)?; + let context = with_cuda_context(0, |context| context.clone())?; + let stream = context.default_stream(); + // SAFETY: the typed declarations and device kernels are compiled together. + let module = unsafe { kernels::load_async(0)? }; + if let Some(trap) = trap { + let plan = module.prepare_block_warp_timesliced::<4>(LaunchConfig3D::auto_shared( + (1, 1, 1), + (8, 4, 2), + ))?; + let source = DeviceBuffer::from_host(&stream, &vec![11_i32; 256])?; + let target = DeviceBuffer::from_host(&stream, &vec![0_i32; 256])?; + let (valid, mode) = match trap { + "block-warp-count" => (257, 1), + "block-warp-unwritten" => (256, 9), + "block-warp-different-count" => (31, 5), + _ => panic!("unknown block warp trap"), + }; + let _ = module + .block_warp_timesliced_async_owned::<4, _, _>(&plan, source, target, valid, mode) + .sync()?; + panic!("expected block warp trap returned successfully"); + } + let mut cases = 0; + macro_rules! run { + ($prepare:ident,$launch:ident,$items:literal,$scratch:literal) => {{ + let plan = + module.$prepare::<$items>(LaunchConfig3D::auto_shared((2, 1, 1), (8, 4, 2)))?; + assert_eq!(plan.__raw_config().shared_mem_bytes, $scratch); + assert!( + module + .$prepare::<$items>(LaunchConfig3D::auto_shared((2, 1, 1), (64, 1, 1))) + .is_err() + ); + for mode in [0, 1, 2, 3, 4, 6] { + for valid in [0, 1, 32 * $items + 1, 64 * $items - 1, 64 * $items] { + if mode == 0 && valid != 64 * $items { + continue; + } + let input: Vec = (0..512).map(|i| i + 19).collect(); + let source = DeviceBuffer::from_host(&stream, &input)?; + let target = DeviceBuffer::from_host(&stream, &vec![0_i32; 512])?; + let (_, target) = module + .$launch::<$items, _, _>(&plan, source, target, valid, mode) + .sync()?; + let mut result = vec![0_i32; 512]; + target.copy_to_host(&stream, &mut result)?; + stream.synchronize()?; + for (index, &value) in result.iter().enumerate() { + let within = index % 256; + let written = if mode == 0 || mode == 2 { + 64 * $items + } else if mode == 6 { + valid / 2 + } else { + valid + }; + let expected = if within < written { + if mode == 2 && within >= valid { + -17 + } else { + input[index] + } + } else { + 0 + }; + assert_eq!( + value, + expected, + "{} items={} mode={} valid={} index={}", + stringify!($launch), + $items, + mode, + valid, + index + ); + } + cases += 1; + } + } + }}; + } + // CUB maximum plus the independent counted-Store bitmap plane, alignment16. + run!(prepare_block_warp, block_warp_async_owned, 1, 272); + run!(prepare_block_warp, block_warp_async_owned, 4, 1072); + run!( + prepare_block_warp_timesliced, + block_warp_timesliced_async_owned, + 1, + 144 + ); + run!( + prepare_block_warp_timesliced, + block_warp_timesliced_async_owned, + 4, + 560 + ); + println!("PASS {cases} public block warp Load/Store cases"); + Ok(()) +} diff --git a/rust/cuda-coop/tests/load_store/src/logical_plain.rs b/rust/cuda-coop/tests/load_store/src/logical_plain.rs new file mode 100644 index 000000000000..cf3d30497db4 --- /dev/null +++ b/rust/cuda-coop/tests/load_store/src/logical_plain.rs @@ -0,0 +1,228 @@ +// Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +//! Scratch-free Warp algorithms, adjusted pointers, and shared Store preservation. +use cuda_coop as coop; +use cuda_device::{cooperative_groups as cg, kernel, launch_contract, thread}; +use cuda_host::cuda_module; + +#[cuda_module] +mod kernels { + use super::*; + use cg::ThreadGroup as _; + + #[inline(never)] + #[allow(clippy::too_many_arguments)] + fn exercise< + T: coop::Numeric + PartialEq + From + core::ops::Add, + A: coop::Algorithm + Default, + const W: usize, + const I: usize, + const N: u32, + >( + input: &[T], + output: &mut coop::DisjointSlice>, + valid: usize, + shift: usize, + mode: u32, + striped: bool, + ) where + coop::ThreadsWithinWarp: coop::LoadStoreGroup, + for<'a> &'a coop::ThreadData: coop::StoreValue, + { + let group = coop::this_warp().group_by::(); + let tile = cg::this_thread_block().tiled_partition::(); + let lane = group.rank() as usize; + let q = coop::this_block().rank_at::() as usize * (32 / W) + + group.rank_at::() as usize; + let offset = thread::blockIdx_x() as usize * 512 + q * W * (8 - I) + shift; + let count = if mode == 0 { W * I } else { valid }; + let mut values = coop::ThreadData::<_, I>::new(); + if mode == 4 { + let mut i = 0; + while i < I { + let position = if striped { lane + W * i } else { lane * I + i }; + if position < count || (lane + i) % 3 == 1 { + values.write(i, T::from(position as i8) + T::from(17)); + } + i += 1; + } + } else { + coop::load( + group, + input, + &mut values, + coop::LoadOptions::default() + .algorithm(A::default()) + .offset(offset), + ); + if mode != 0 { + let options = coop::LoadOptions::default() + .algorithm(A::default()) + .offset(offset) + .valid_items(count); + if mode == 2 { + coop::load(group, input, &mut values, options.oob_default(T::from(-7))); + } else { + coop::load(group, input, &mut values, options); + } + } + } + let before = values; + let options = coop::StoreOptions::default() + .algorithm(A::default()) + .offset(offset); + if mode == 0 { + coop::store(group, &mut *output, &values, options); + } else { + coop::store(group, &mut *output, &values, options.valid_items(count)); + } + let after = values; + tile.sync(); + coop::store( + group, + &mut *output, + &values, + coop::StoreOptions::default() + .algorithm(A::default()) + .offset(offset) + .valid_items(0), + ); + let twice = values; + tile.sync(); + coop::load( + group, + output, + &mut values, + coop::LoadOptions::default() + .algorithm(A::default()) + .offset(offset) + .valid_items(count), + ); + let mut matches = true; + let mut i = 0; + while i < I { + let position = if striped { lane + W * i } else { lane * I + i }; + let initialized = if mode == 4 { + position < count || (lane + i) % 3 == 1 + } else { + mode == 0 || mode == 2 || position < count + }; + let expected = if initialized { + Some(if mode == 4 { + T::from(position as i8) + T::from(17) + } else if mode == 2 && position >= count { + T::from(-7) + } else { + input[offset + q * W * I + position] + }) + } else { + None + }; + matches &= + before.get(i) == expected && after.get(i) == expected && twice.get(i) == expected; + matches &= values.get(i) == if position < count { expected } else { None }; + i += 1; + } + assert!(matches); + } + + #[kernel] + #[launch_contract(domain=1,block=(64,1,1),dynamic_shared=inferred)] + pub fn logical_plain< + T: coop::Numeric + PartialEq + From + core::ops::Add, + A: coop::Algorithm + Default, + const W: usize, + const I: usize, + const N: u32, + >( + input: &[T], + mut output: coop::DisjointSlice>, + valid: usize, + shift: usize, + mode: u32, + striped: bool, + ) where + coop::ThreadsWithinWarp: coop::LoadStoreGroup, + for<'a> &'a coop::ThreadData: coop::StoreValue, + { + exercise::(input, &mut output, valid, shift, mode, striped); + } +} + +pub fn run() -> Result<(), Box> { + use cuda_async::simt::device_context::{init_device_contexts, with_cuda_context}; + use cuda_async::simt::device_operation::DeviceOperation; + use cuda_core::{DeviceBuffer, LaunchConfig1D}; + init_device_contexts(0, 1)?; + let stream = with_cuda_context(0, |context| context.default_stream())?; + // SAFETY: the typed module and launchers compile together. + let module = unsafe { kernels::load_async(0)? }; + let mut cases = 0; + macro_rules! exercise { + ($ty:ty,$algorithm:ty,$striped:literal,$w:literal,$i:literal,$controls:expr) => {{ + let input: Vec<$ty> = (0..1536).map(|i| (i + 31) as $ty).collect(); + let plan = module.prepare_logical_plain::<$ty, $algorithm, $w, $i, $w>( + LaunchConfig1D::auto_shared(3, 64), + )?; + assert_eq!(plan.__raw_config().shared_mem_bytes, 0); + for shift in [0, 1] { + for (mode, count) in $controls { + let source = DeviceBuffer::from_host(&stream, &input)?; + let target = DeviceBuffer::from_host(&stream, &vec![-99 as $ty; input.len()])?; + let (_, target) = module + .logical_plain_async_owned::<$ty, $algorithm, $w, $i, $w, _, _>( + &plan, source, target, count, shift, mode, $striped, + ) + .sync()?; + let mut result = vec![0 as $ty; input.len()]; + target.copy_to_host(&stream, &mut result)?; + stream.synchronize()?; + let mut expected = vec![-99 as $ty; input.len()]; + for block in 0..3 { + for q in 0..64 / $w { + let valid = if mode == 0 { $w * $i } else { count }; + for i in 0..valid { + let index = block * 512 + q * $w * 8 + shift + i; + expected[index] = if mode == 4 { + i as $ty + 17 as $ty + } else { + input[index] + }; + } + } + } + assert_eq!(result, expected); + cases += 1; + } + } + }}; + } + let controls = |width: usize, items: usize| { + let mut counts = vec![0, 1, width * items - 1, width * items]; + counts.sort_unstable(); + counts.dedup(); + [0, 1, 2, 4] + .into_iter() + .flat_map(|mode| { + counts.iter().copied().filter_map(move |count| { + (mode != 0 || count == width * items).then_some((mode, count)) + }) + }) + .collect::>() + }; + exercise!(i32, coop::Striped, true, 1, 1, controls(1, 1)); + exercise!(i32, coop::Striped, true, 8, 4, controls(8, 4)); + exercise!(i32, coop::Striped, true, 32, 4, controls(32, 4)); + exercise!(i32, coop::Vectorize, false, 1, 1, controls(1, 1)); + exercise!(i32, coop::Vectorize, false, 8, 4, controls(8, 4)); + exercise!(i32, coop::Vectorize, false, 32, 4, controls(32, 4)); + exercise!(i8, coop::Striped, true, 8, 4, [(0, 32), (1, 31)]); + exercise!(i8, coop::Vectorize, false, 8, 4, [(0, 32), (1, 31)]); + exercise!(f64, coop::Striped, true, 8, 4, [(0, 32), (1, 31)]); + exercise!(f64, coop::Vectorize, false, 8, 4, [(0, 32), (1, 31)]); + println!( + "PASS {cases} Warp Striped/Vectorize cases: shared Store preservation, offsets0/1, zero scratch" + ); + Ok(()) +} diff --git a/rust/cuda-coop/tests/load_store/src/logical_transpose.rs b/rust/cuda-coop/tests/load_store/src/logical_transpose.rs new file mode 100644 index 000000000000..037324ec2d03 --- /dev/null +++ b/rust/cuda-coop/tests/load_store/src/logical_transpose.rs @@ -0,0 +1,513 @@ +// Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +//! Replicated implicit Warp scratch and observed Transpose payload mutation. +use cuda_coop as coop; +use cuda_device::{kernel, launch_contract, thread}; +use cuda_host::cuda_module; + +#[cuda_module] +mod kernels { + use super::*; + + #[inline(never)] + #[allow(clippy::too_many_arguments)] + fn exercise, const I: usize>( + group: coop::ThreadGroup, + input: &[i32], + output: &mut coop::DisjointSlice>, + width: usize, + q: usize, + lane: usize, + offset: usize, + valid: usize, + mode: u32, + ) { + let count = if mode == 0 { + width * I + } else if mode == 3 { + (valid + q).min(width * I) + } else { + valid + }; + let mut values = coop::ThreadData::<_, I>::new(); + if mode == 4 { + let mut i = 0; + while i < I { + if lane * I + i < count || (lane + i) % 3 == 1 { + values.write(i, (lane * I + i) as i32 + 17); + } + i += 1; + } + } else { + coop::load( + group, + input, + &mut values, + coop::LoadOptions::default() + .algorithm(coop::Transpose) + .offset(offset), + ); + if mode != 0 { + let options = coop::LoadOptions::default() + .algorithm(coop::Transpose) + .offset(offset) + .valid_items(count); + if mode == 2 { + coop::load(group, input, &mut values, options.oob_default(-7)); + } else { + coop::load(group, input, &mut values, options); + } + } + } + let before = values; + coop::store( + group, + &mut *output, + &mut values, + coop::StoreOptions::default() + .algorithm(coop::Transpose) + .offset(offset) + .valid_items(count), + ); + let after = values; + coop::store( + group, + &mut *output, + &mut values, + coop::StoreOptions::default() + .algorithm(coop::Transpose) + .offset(offset) + .valid_items(0), + ); + let twice = values; + // Managed scratch reuse supplies matching completion before this read. + coop::load( + group, + output, + &mut values, + coop::LoadOptions::default() + .algorithm(coop::Transpose) + .offset(offset) + .valid_items(count), + ); + let mut matches = true; + for i in 0..I { + let first = width * i + lane; + let second = width * (first % I) + first / I; + let mut observation = 0; + while observation < 3 { + let (snapshot, index) = match observation { + 0 => (before.get(i), lane * I + i), + 1 => (after.get(i), first), + _ => (twice.get(i), second), + }; + let initialized = if mode == 4 { + index < count || (index / I + index % I) % 3 == 1 + } else { + mode == 0 || mode == 2 || index < count + }; + let expected = if initialized { + Some(if mode == 4 { + index as i32 + 17 + } else if mode == 2 && index >= count { + -7 + } else { + input[offset + q * width * I + index] + }) + } else { + None + }; + matches &= snapshot == expected; + observation += 1; + } + let index = lane * I + i; + let expected = if index < count { + Some(if mode == 4 { + index as i32 + 17 + } else { + input[offset + q * width * I + index] + }) + } else { + None + }; + matches &= values.get(i) == expected; + } + assert!(matches); + } + + #[inline(never)] + fn selected( + input: &[i32], + output: &mut coop::DisjointSlice>, + valid: usize, + mode: u32, + ) where + coop::ThreadsWithinWarp: coop::LoadStoreGroup, + { + let group = coop::this_warp().group_by::(); + let q = coop::this_block().rank_at::() as usize * (32 / W) + + group.rank_at::() as usize; + let block = thread::blockIdx_x() as usize + + thread::gridDim_x() as usize + * (thread::blockIdx_y() as usize + + thread::gridDim_y() as usize * thread::blockIdx_z() as usize); + exercise( + group, + input, + output, + W, + q, + group.rank() as usize, + block * 64 * I, + valid, + mode, + ); + } + #[kernel] + #[launch_contract(domain=1,block=(64,1,1),dynamic_shared=inferred)] + pub fn logical_transpose( + input: &[i32], + mut output: coop::DisjointSlice>, + valid: usize, + mode: u32, + ) where + coop::ThreadsWithinWarp: coop::LoadStoreGroup, + { + selected::(input, &mut output, valid, mode); + } + #[kernel] + #[launch_contract(domain=1,block=(128,1,1),dynamic_shared=inferred)] + pub fn logical_transpose_b128( + input: &[i32], + mut output: coop::DisjointSlice>, + ) { + let group = coop::this_warp().group_by::<8>(); + let q = coop::this_block().rank_at::() as usize * 4 + + group.rank_at::() as usize; + exercise( + group, + input, + &mut output, + 8, + q, + group.rank() as usize, + thread::blockIdx_x() as usize * 512, + 17, + 4, + ); + } + #[kernel] + #[launch_contract(domain=3,block=(8,4,2),dynamic_shared=inferred)] + pub fn logical_transpose_xyz( + input: &[i32], + mut output: coop::DisjointSlice>, + valid: usize, + mode: u32, + ) { + selected::<8, 4>(input, &mut output, valid, mode); + } + #[kernel] + #[launch_contract(domain=1,block=(64,1,1),dynamic_shared=inferred)] + pub fn logical_transpose_prefix( + input: &[i32], + mut output: coop::DisjointSlice>, + valid: usize, + mode: u32, + ) { + // SAFETY: one whole fixed prefix, disjoint lane writes and Block completion. + let lane = thread::threadIdx_x() as usize; + let prefix = + unsafe { cuda_device::DynamicSharedArray::<[i32; 64], 32>::get_fixed() as *mut i32 }; + unsafe { + prefix.add(lane).write(lane as i32); + } + thread::sync_threads(); + selected::<8, 4>(input, &mut output, valid, mode); + let neighbor = (lane + 1) % 64; + assert!(unsafe { prefix.add(neighbor).read() } == neighbor as i32); + } + #[kernel] + #[launch_contract(domain=1,block=(64,1,1),dynamic_shared=inferred)] + pub fn logical_transpose_containing( + input: &[i32], + mut output: coop::DisjointSlice>, + ) { + use cuda_device::cooperative_groups::{ThreadGroup, this_thread_block}; + let warp8 = coop::this_warp().group_by::<8>(); + let warp16 = coop::this_warp().group_by::<16>(); + let tile16 = this_thread_block().tiled_partition::<16>(); + let offset = thread::blockIdx_x() as usize * 256; + let lane = thread::threadIdx_x() as usize; + let mut values = coop::ThreadData::<_, 4>::new(); + coop::load( + warp8, + input, + &mut values, + coop::LoadOptions::default() + .algorithm(coop::Transpose) + .offset(offset), + ); + coop::store( + warp8, + &mut output, + &mut values, + coop::StoreOptions::default() + .algorithm(coop::Transpose) + .offset(offset) + .valid_items(32), + ); + tile16.sync(); + coop::load( + warp16, + &mut output, + &mut values, + coop::LoadOptions::default() + .algorithm(coop::Transpose) + .offset(offset), + ); + let first = values; + coop::load( + warp16, + input, + &mut values, + coop::LoadOptions::default() + .algorithm(coop::Transpose) + .offset(offset), + ); + let second = values; + tile16.sync(); + coop::load( + warp8, + &mut output, + &mut values, + coop::LoadOptions::default() + .algorithm(coop::Transpose) + .offset(offset), + ); + let mut matches = true; + for i in 0..4 { + let expected = Some(input[offset + lane * 4 + i]); + matches &= + first.get(i) == expected && second.get(i) == expected && values.get(i) == expected; + } + assert!(matches); + } + #[kernel] + #[launch_contract(domain=1,block=(64,1,1),dynamic_shared=inferred)] + pub fn logical_transpose_mixed( + input: &[i32], + mut output: coop::DisjointSlice>, + ) { + let block = coop::this_block(); + let warp = coop::this_warp().group_by::<8>(); + let offset = thread::blockIdx_x() as usize * 256; + let lane = thread::threadIdx_x() as usize; + let mut values = coop::ThreadData::<_, 4>::new(); + coop::load( + block, + input, + &mut values, + coop::LoadOptions::default() + .algorithm(coop::Transpose) + .offset(offset), + ); + let first = values; + coop::load( + warp, + input, + &mut values, + coop::LoadOptions::default() + .algorithm(coop::Transpose) + .offset(offset), + ); + let second = values; + coop::load( + warp, + input, + &mut values, + coop::LoadOptions::default() + .algorithm(coop::Transpose) + .offset(offset), + ); + let third = values; + // W8 reuse must leave Block scratch dirty, requiring a Block barrier here. + coop::load( + block, + input, + &mut values, + coop::LoadOptions::default() + .algorithm(coop::Transpose) + .offset(offset), + ); + coop::store( + block, + &mut output, + &values, + coop::StoreOptions::default().offset(offset), + ); + let mut matches = true; + for i in 0..4 { + let expected = Some(input[offset + lane * 4 + i]); + matches &= first.get(i) == expected + && second.get(i) == expected + && third.get(i) == expected + && values.get(i) == expected; + } + assert!(matches); + } +} + +pub fn run() -> Result<(), Box> { + use cuda_async::simt::device_context::{init_device_contexts, with_cuda_context}; + use cuda_async::simt::device_operation::DeviceOperation; + use cuda_core::{DeviceBuffer, LaunchConfig1D}; + init_device_contexts(0, 1)?; + let stream = with_cuda_context(0, |context| context.default_stream())?; + // SAFETY: the typed module and its launchers compile together. + let module = unsafe { kernels::load_async(0)? }; + let mut cases = 0; + macro_rules! exercise { + ($w:literal,$i:literal,$scratch:literal,$prepare:expr,$launch:expr,$config:expr) => {{ + let plan = ($prepare)(($config)(3))?; + assert_eq!(plan.__raw_config().shared_mem_bytes, $scratch); + let input: Vec = (0..3 * 64 * $i).map(|i| i as i32 + 31).collect(); + for mode in 0..5 { + for valid in [0, 1, $w * $i - 1, $w * $i] { + let source = DeviceBuffer::from_host(&stream, &input)?; + let target = DeviceBuffer::from_host(&stream, &vec![-99_i32; input.len()])?; + let (_, target) = ($launch)(&plan, source, target, valid, mode).sync()?; + let mut result = vec![0; input.len()]; + target.copy_to_host(&stream, &mut result)?; + stream.synchronize()?; + let mut expected = vec![-99; input.len()]; + for block in 0..3 { + for q in 0..64 / $w { + let count = if mode == 0 { + $w * $i + } else if mode == 3 { + (valid + q).min($w * $i) + } else { + valid + }; + for i in 0..count { + let index = block * 64 * $i + q * $w * $i + i; + expected[index] = if mode == 4 { + i as i32 + 17 + } else { + input[index] + }; + } + } + } + assert_eq!(result, expected); + cases += 1; + } + } + }}; + } + exercise!( + 1, + 1, + 512, + |c| module.prepare_logical_transpose::<1, 1>(c), + |p, s, t, v, m| module.logical_transpose_async_owned::<1, 1, _, _>(p, s, t, v, m), + |g| LaunchConfig1D::auto_shared(g, 64) + ); + exercise!( + 8, + 4, + 1184, + |c| module.prepare_logical_transpose::<8, 4>(c), + |p, s, t, v, m| module.logical_transpose_async_owned::<8, 4, _, _>(p, s, t, v, m), + |g| LaunchConfig1D::auto_shared(g, 64) + ); + exercise!( + 32, + 4, + 1064, + |c| module.prepare_logical_transpose::<32, 4>(c), + |p, s, t, v, m| module.logical_transpose_async_owned::<32, 4, _, _>(p, s, t, v, m), + |g| LaunchConfig1D::auto_shared(g, 64) + ); + exercise!( + 8, + 4, + 1184, + |c| module.prepare_logical_transpose_xyz(c), + |p, s, t, v, m| module.logical_transpose_xyz_async_owned(p, s, t, v, m), + |g| cuda_core::LaunchConfig3D::auto_shared((1, g, 1), (8, 4, 2)) + ); + exercise!( + 8, + 4, + 1440, + |c| module.prepare_logical_transpose_prefix(c), + |p, s, t, v, m| module.logical_transpose_prefix_async_owned(p, s, t, v, m), + |g| LaunchConfig1D::auto_shared(g, 64) + ); + assert!( + module + .prepare_logical_transpose::<8, 4>(LaunchConfig1D::new(3, 64, 1183)) + .is_err() + ); + let plan = module.prepare_logical_transpose_mixed(LaunchConfig1D::auto_shared(3, 64))?; + assert_eq!(plan.__raw_config().shared_mem_bytes, 2080); + let input: Vec = (0..768).map(|i| i + 53).collect(); + let source = DeviceBuffer::from_host(&stream, &input)?; + let target = DeviceBuffer::from_host(&stream, &vec![-99_i32; 768])?; + let (_, target) = module + .logical_transpose_mixed_async_owned(&plan, source, target) + .sync()?; + let mut result = vec![0; 768]; + target.copy_to_host(&stream, &mut result)?; + stream.synchronize()?; + assert_eq!(result, input); + let plan = module.prepare_logical_transpose_containing(LaunchConfig1D::auto_shared(3, 64))?; + assert_eq!(plan.__raw_config().shared_mem_bytes, 2224); + let source = DeviceBuffer::from_host(&stream, &input)?; + let target = DeviceBuffer::from_host(&stream, &vec![-99_i32; 768])?; + let (_, target) = module + .logical_transpose_containing_async_owned(&plan, source, target) + .sync()?; + target.copy_to_host(&stream, &mut result)?; + stream.synchronize()?; + assert_eq!(result, input); + println!( + "PASS {cases} logical Transpose cases plus mixed Block/W8 and W8/W16 pool reuse, containing synchronization and undersized launch rejection" + ); + run_b128()?; + Ok(()) +} + +/// One changed-replica-count contrast using the same value/validity helper. +fn run_b128() -> Result<(), Box> { + use cuda_async::simt::device_context::with_cuda_context; + use cuda_async::simt::device_operation::DeviceOperation; + use cuda_core::{DeviceBuffer, LaunchConfig1D}; + let stream = with_cuda_context(0, |context| context.default_stream())?; + // SAFETY: the typed module and its launchers compile together. + let module = unsafe { kernels::load_async(0)? }; + let plan = module.prepare_logical_transpose_b128(LaunchConfig1D::auto_shared(3, 128))?; + assert_eq!(plan.__raw_config().shared_mem_bytes, 148 * 16); + let input: Vec = (0..1536).map(|i| i + 31).collect(); + let source = DeviceBuffer::from_host(&stream, &input)?; + let target = DeviceBuffer::from_host(&stream, &vec![-99_i32; input.len()])?; + let (_, target) = module + .logical_transpose_b128_async_owned(&plan, source, target) + .sync()?; + let mut result = vec![0; input.len()]; + target.copy_to_host(&stream, &mut result)?; + stream.synchronize()?; + let mut expected = vec![-99; input.len()]; + for block in 0..3 { + for q in 0..16 { + for i in 0..17 { + expected[block * 512 + q * 32 + i] = i as i32 + 17; + } + } + } + assert_eq!(result, expected); + println!("PASS B128/W8 Transpose: 16 replicas, 2368 bytes, observed values/validity"); + Ok(()) +} diff --git a/rust/cuda-coop/tests/load_store/src/logical_warp.rs b/rust/cuda-coop/tests/load_store/src/logical_warp.rs new file mode 100644 index 000000000000..7d9b3af76c45 --- /dev/null +++ b/rust/cuda-coop/tests/load_store/src/logical_warp.rs @@ -0,0 +1,504 @@ +// Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +//! Logical Warp Direct through inferred public calls and a generic helper. +use cuda_coop as coop; +use cuda_device::{cooperative_groups as cg, kernel, launch_contract, thread}; +use cuda_host::cuda_module; + +#[cuda_module] +mod kernels { + use super::*; + use cg::ThreadGroup as _; + + #[inline(never)] + #[allow(clippy::too_many_arguments)] + fn exercise, const I: usize, const N: u32>( + group: coop::ThreadGroup, + tile: cg::WarpTile, + input: &[i32], + output: &mut coop::DisjointSlice>, + width: usize, + q: usize, + lane: usize, + offset: usize, + valid: usize, + mode: u32, + ) { + let count = if mode == 0 { + width * I + } else if mode == 3 { + (valid + q).min(width * I) + } else { + valid + }; + let mut values = coop::ThreadData::<_, I>::new(); + if mode == 4 { + let mut other = coop::ThreadData::<_, I>::new(); + let mut i = 0; + while i < I { + if lane * I + i < count || i % 2 == 0 { + values.write(i, -7); + } + if lane * I + i < count || i % 2 != 0 { + other.write(i, -7); + } + i += 1; + } + values = if lane & 1 == 0 { values } else { other }; + } else if mode == 0 { + coop::load( + group, + input, + &mut values, + coop::LoadOptions::default().offset(offset), + ); + } else if mode == 2 { + coop::load( + group, + input, + &mut values, + coop::LoadOptions::default() + .offset(offset) + .valid_items(count) + .oob_default(-7), + ); + } else { + // A preceding full Load makes the exact masked reset observable. + coop::load( + group, + input, + &mut values, + coop::LoadOptions::default().offset(offset), + ); + coop::load( + group, + input, + &mut values, + coop::LoadOptions::default() + .offset(offset) + .valid_items(count), + ); + } + let before = values; + coop::store( + group, + &mut *output, + &values, + coop::StoreOptions::default() + .offset(offset) + .valid_items(count), + ); + let after = values; + tile.sync(); + coop::load( + group, + output, + &mut values, + coop::LoadOptions::default() + .offset(offset) + .valid_items(count), + ); + let mut matches = true; + let mut i = 0; + while i < I { + matches &= before.get(i) == after.get(i); + if lane * I + i < count { + let expected = if mode == 4 { + -7 + } else { + input[offset + q * width * I + lane * I + i] + }; + matches &= values.get(i) == Some(expected); + } else { + matches &= values.get(i).is_none(); + if mode == 1 || mode == 3 { + matches &= before.get(i).is_none(); + } + if mode == 2 { + matches &= before.get(i) == Some(-7); + } + } + i += 1; + } + assert!(matches); + } + + #[inline(never)] + fn selected( + input: &[i32], + output: &mut coop::DisjointSlice>, + valid: usize, + mode: u32, + ) where + coop::ThreadsWithinWarp: coop::LoadStoreGroup, + { + let group = coop::this_warp().group_by::(); + let q = coop::this_block().rank_at::() as usize * (32 / W) + + group.rank_at::() as usize; + let block = thread::blockIdx_x() as usize + + thread::gridDim_x() as usize + * (thread::blockIdx_y() as usize + + thread::gridDim_y() as usize * thread::blockIdx_z() as usize); + exercise( + group, + cg::this_thread_block().tiled_partition::(), + input, + output, + W, + q, + group.rank() as usize, + block * 64 * I, + valid, + mode, + ); + } + #[kernel] + #[launch_contract(domain=1,block=(64,1,1),dynamic_shared=inferred)] + pub fn logical_direct( + input: &[i32], + mut output: coop::DisjointSlice>, + valid: usize, + mode: u32, + ) where + coop::ThreadsWithinWarp: coop::LoadStoreGroup, + { + selected::(input, &mut output, valid, mode); + } + #[kernel] + #[launch_contract(domain=3,block=(8,4,2),dynamic_shared=inferred)] + pub fn logical_xyz( + input: &[i32], + mut output: coop::DisjointSlice>, + valid: usize, + mode: u32, + ) { + selected::<8, 4, 8>(input, &mut output, valid, mode); + } + #[kernel] + #[launch_contract(domain=1,block=(64,1,1),dynamic_shared=inferred)] + pub fn logical_nonexhaustive( + input: &[i32], + mut output: coop::DisjointSlice>, + valid: usize, + mode: u32, + ) { + let group = coop::this_warp().group_by_non_exhaustive::<8>(); + let q = coop::this_block().rank_at::() as usize * 4 + + group.rank_at::() as usize; + exercise( + group, + cg::this_thread_block().tiled_partition::<8>(), + input, + &mut output, + 8, + q, + group.rank() as usize, + thread::blockIdx_x() as usize * 256, + valid, + mode, + ); + } + #[kernel] + #[launch_contract(domain=1,block=(64,1,1),dynamic_shared=inferred)] + pub fn logical_grain( + mut output: coop::DisjointSlice>, + adapt: bool, + ) { + let group = coop::this_warp().group_by::<8>(); + let q = coop::this_block().rank_at::() as usize * 4 + + group.rank_at::() as usize; + let offset = thread::blockIdx_x() as usize * 256 + if adapt { q * 24 } else { 0 }; + coop::store( + group, + &mut output, + 7_i32, + coop::StoreOptions::default().offset(offset), + ); + } + + #[kernel] + #[launch_contract(domain=1,block=(64,1,1),dynamic_shared=inferred)] + pub fn logical_cycle( + input: &[i32], + mut output: coop::DisjointSlice>, + ) { + let group = coop::this_warp().group_by::<8>(); + let q = coop::this_block().rank_at::() as usize * 4 + + group.rank_at::() as usize; + let block = thread::blockIdx_x() as usize; + let lane = thread::threadIdx_x() as usize; + let base = block * 9 * 256; + let mut values = coop::ThreadData::<_, 4>::new(); + let mut sums = [0_i32; 4]; + let mut matches = true; + let mut iteration = 0; + // Siblings execute zero through seven iterations over immutable input. + while iteration < q { + let offset = base + iteration * 256; + coop::load( + group, + input, + &mut values, + coop::LoadOptions::default().offset(offset), + ); + let mut i = 0; + while i < 4 { + matches &= values.get(i) == Some(input[offset + lane * 4 + i]); + sums[i] += values.get(i).unwrap_or(0); + i += 1; + } + iteration += 1; + } + if COMPLETE { + cg::this_thread_block().tiled_partition::<32>().sync(); + } + coop::load( + coop::this_warp(), + input, + &mut values, + coop::LoadOptions::default().offset(base + 7 * 256), + ); + let mut i = 0; + while i < 4 { + matches &= values.get(i) == Some(input[base + 7 * 256 + lane * 4 + i]); + sums[i] += 3 * values.get(i).unwrap_or(0); + i += 1; + } + if COMPLETE { + thread::sync_threads(); + } + coop::load( + coop::this_block(), + input, + &mut values, + coop::LoadOptions::default().offset(base + 8 * 256), + ); + let mut i = 0; + while i < 4 { + matches &= values.get(i) == Some(input[base + 8 * 256 + lane * 4 + i]); + values.write(i, sums[i] + 5 * values.get(i).unwrap_or(0)); + i += 1; + } + coop::store( + coop::this_block(), + &mut output, + &values, + coop::StoreOptions::default().offset(block * 256), + ); + assert!(matches); + } + + #[kernel] + #[launch_contract(domain=1,block=(64,1,1),dynamic_shared=inferred)] + pub fn logical_mixed_block( + input: &[i32], + mut output: coop::DisjointSlice>, + ) { + let block = coop::this_block(); + let group = coop::this_warp().group_by::<8>(); + let offset = thread::blockIdx_x() as usize * 256; + let lane = thread::threadIdx_x() as usize; + let mut values = coop::ThreadData::<_, 4>::new(); + coop::load( + block, + input, + &mut values, + coop::LoadOptions::default() + .offset(offset) + .algorithm(coop::Transpose), + ); + coop::store( + group, + &mut output, + &values, + coop::StoreOptions::default().offset(offset), + ); + cg::this_thread_block().tiled_partition::<8>().sync(); + coop::load( + group, + &mut output, + &mut values, + coop::LoadOptions::default().offset(offset), + ); + let mut matches = (values.get(0) == Some(input[offset + lane * 4])) + & (values.get(1) == Some(input[offset + lane * 4 + 1])) + & (values.get(2) == Some(input[offset + lane * 4 + 2])) + & (values.get(3) == Some(input[offset + lane * 4 + 3])); + // Tile completion leaves the Block scratch dirty. Managed reuse must + // insert a full-block barrier before this second exchange. + coop::load( + block, + input, + &mut values, + coop::LoadOptions::default() + .offset(offset) + .algorithm(coop::Transpose), + ); + for i in 0..4 { + matches &= values.get(i) == Some(input[offset + lane * 4 + i]); + } + assert!(matches); + } +} + +pub fn run() -> Result<(), Box> { + use cuda_async::simt::device_context::{init_device_contexts, with_cuda_context}; + use cuda_async::simt::device_operation::DeviceOperation; + use cuda_core::{DeviceBuffer, LaunchConfig1D}; + init_device_contexts(0, 1)?; + let stream = with_cuda_context(0, |context| context.default_stream())?; + // SAFETY: this module and its typed launchers compile together. + let module = unsafe { kernels::load_async(0)? }; + let mut cases = 0; + macro_rules! exercise { + ($width:literal,$items:literal,$prepare:expr,$launch:expr,$config:expr) => {{ + let groups = 3; + let config = ($config)(groups); + let plan = ($prepare)(config)?; + assert_eq!(plan.__raw_config().shared_mem_bytes, 0); + let input: Vec = (0..groups as usize * 64 * $items) + .map(|i| i as i32 + 31) + .collect(); + for mode in 0..5 { + for valid in [0, 1, $width * $items - 1, $width * $items] { + let source = DeviceBuffer::from_host(&stream, &input)?; + let target = DeviceBuffer::from_host(&stream, &vec![-99_i32; input.len()])?; + let (_, target) = ($launch)(&plan, source, target, valid, mode).sync()?; + let mut result = vec![0; input.len()]; + target.copy_to_host(&stream, &mut result)?; + stream.synchronize()?; + let mut expected = vec![-99; input.len()]; + for block in 0..groups as usize { + for q in 0..64 / $width { + let count = if mode == 0 { + $width * $items + } else if mode == 3 { + (valid + q).min($width * $items) + } else { + valid + }; + for i in 0..count { + let index = block * 64 * $items + q * $width * $items + i; + expected[index] = if mode == 4 { -7 } else { input[index] }; + } + } + } + assert_eq!(result, expected); + cases += 1; + } + } + }}; + } + exercise!( + 1, + 1, + |c| module.prepare_logical_direct::<1, 1, 1>(c), + |p, s, t, v, m| module.logical_direct_async_owned::<1, 1, 1, _, _>(p, s, t, v, m), + |g| LaunchConfig1D::auto_shared(g, 64) + ); + exercise!( + 8, + 4, + |c| module.prepare_logical_direct::<8, 4, 8>(c), + |p, s, t, v, m| module.logical_direct_async_owned::<8, 4, 8, _, _>(p, s, t, v, m), + |g| LaunchConfig1D::auto_shared(g, 64) + ); + exercise!( + 16, + 4, + |c| module.prepare_logical_direct::<16, 4, 16>(c), + |p, s, t, v, m| module.logical_direct_async_owned::<16, 4, 16, _, _>(p, s, t, v, m), + |g| LaunchConfig1D::auto_shared(g, 64) + ); + exercise!( + 8, + 4, + |c| module.prepare_logical_xyz(c), + |p, s, t, v, m| module.logical_xyz_async_owned(p, s, t, v, m), + |g| cuda_core::LaunchConfig3D::auto_shared((1, g, 1), (8, 4, 2)) + ); + exercise!( + 8, + 4, + |c| module.prepare_logical_nonexhaustive(c), + |p, s, t, v, m| module.logical_nonexhaustive_async_owned(p, s, t, v, m), + |g| LaunchConfig1D::auto_shared(g, 64) + ); + let grain_plan = module.prepare_logical_grain(LaunchConfig1D::auto_shared(2, 64))?; + let target = DeviceBuffer::from_host(&stream, &vec![-99_i32; 512])?; + let adapt = !std::env::args().any(|arg| arg == "logical-grain-trap"); + let target = module + .logical_grain_async_owned(&grain_plan, target, adapt) + .sync()?; + assert!(adapt, "expected fixed-grain trap completed successfully"); + let mut result = vec![0; 512]; + target.copy_to_host(&stream, &mut result)?; + stream.synchronize()?; + let mut expected = vec![-99; 512]; + for block in 0..2 { + for q in 0..8 { + expected[block * 256 + q * 32..block * 256 + q * 32 + 8].fill(7); + } + } + assert_eq!(result, expected); + println!("PASS {cases} logical Warp mask/default/scattered cases and fixed-grain offset24"); + + let plan = module.prepare_logical_mixed_block(LaunchConfig1D::auto_shared(3, 64))?; + assert_eq!(plan.__raw_config().shared_mem_bytes, 1024); + let input: Vec = (0..768).map(|i| i + 53).collect(); + let source = DeviceBuffer::from_host(&stream, &input)?; + let target = DeviceBuffer::from_host(&stream, &vec![-99_i32; input.len()])?; + let (_, target) = module + .logical_mixed_block_async_owned(&plan, source, target) + .sync()?; + let mut result = vec![0; input.len()]; + target.copy_to_host(&stream, &mut result)?; + stream.synchronize()?; + assert_eq!(result, input); + println!("PASS mixed Block Transpose/logical Warp Direct: 768 values"); + run_cycles()?; + Ok(()) +} + +fn run_cycles() -> Result<(), Box> { + use cuda_async::simt::device_context::with_cuda_context; + use cuda_async::simt::device_operation::DeviceOperation; + use cuda_core::{DeviceBuffer, LaunchConfig1D}; + let stream = with_cuda_context(0, |context| context.default_stream())?; + // SAFETY: this module and its typed launchers compile together. + let module = unsafe { kernels::load_async(0)? }; + let input: Vec = (0..3 * 9 * 256).map(|i| i + 17).collect(); + let expected: Vec = (0..3 * 256) + .map(|index| { + let base = index / 256 * 9 * 256; + let element = index % 256; + let q = element / 4 / 8; + (0..q) + .map(|iteration| input[base + iteration * 256 + element]) + .sum::() + + 3 * input[base + 7 * 256 + element] + + 5 * input[base + 8 * 256 + element] + }) + .collect(); + macro_rules! exercise { + ($complete:literal) => {{ + let plan = + module.prepare_logical_cycle::<$complete>(LaunchConfig1D::auto_shared(3, 64))?; + assert_eq!(plan.__raw_config().shared_mem_bytes, 0); + let source = DeviceBuffer::from_host(&stream, &input)?; + let target = DeviceBuffer::from_host(&stream, &vec![-99_i32; expected.len()])?; + let (_, target) = module + .logical_cycle_async_owned::<$complete, _, _>(&plan, source, target) + .sync()?; + let mut result = vec![0; expected.len()]; + target.copy_to_host(&stream, &mut result)?; + stream.synchronize()?; + assert_eq!(result, expected); + }}; + } + exercise!(false); + exercise!(true); + println!("PASS finite sibling cycles then W32/Block: 1536 values, without/with barriers"); + Ok(()) +} diff --git a/rust/cuda-coop/tests/load_store/src/main.rs b/rust/cuda-coop/tests/load_store/src/main.rs new file mode 100644 index 000000000000..2c27de8e667f --- /dev/null +++ b/rust/cuda-coop/tests/load_store/src/main.rs @@ -0,0 +1,766 @@ +// Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +use coop::{Numeric, ThreadData, this_block}; +mod block_warp; +mod logical_plain; +mod logical_transpose; +mod logical_warp; +mod storage; +mod striped; +mod transpose; +mod vectorize; +mod warp; + +use cuda_coop as coop; +use cuda_device::{CollectiveTiles, DisjointSlice, kernel, launch_contract, thread}; +use cuda_host::cuda_module; + +#[cuda_module] +mod kernels { + use super::*; + + #[inline(never)] + fn exercise( + input: &[T], + output: &mut DisjointSlice>, + valid: usize, + shift: usize, + default: T, + mode: u32, + ) { + let block = this_block(); + let x = thread::blockDim_x() as usize; + let y = thread::blockDim_y() as usize; + let threads = x * y * thread::blockDim_z() as usize; + let lane = thread::threadIdx_x() as usize + + x * (thread::threadIdx_y() as usize + y * thread::threadIdx_z() as usize); + let group = thread::blockIdx_x() as usize + + thread::gridDim_x() as usize + * (thread::blockIdx_y() as usize + + thread::gridDim_y() as usize * thread::blockIdx_z() as usize); + let offset = group * threads * 4 + shift; + let mut values = ThreadData::<_, ITEMS_PER_THREAD>::new(); + match mode { + 0 => coop::load( + block, + input, + &mut values, + coop::LoadOptions::default().offset(offset), + ), + 1 | 5 | 6 => { + coop::load( + block, + input, + &mut values, + coop::LoadOptions::default().offset(offset), + ); + coop::load( + block, + input, + &mut values, + coop::LoadOptions::default() + .offset(offset) + .valid_items(valid), + ); + } + 2 | 4 => coop::load( + block, + input, + &mut values, + coop::LoadOptions::default() + .offset(offset) + .valid_items(valid) + .oob_default(default), + ), + 9 => {} + 10 => coop::load( + block, + input, + &mut values, + coop::LoadOptions::default() + .offset(offset) + .oob_default(default), + ), + 11 => { + coop::load( + block, + input, + &mut values, + coop::LoadOptions::default().offset(offset).valid_items(0), + ); + let _invalid = values[0]; + } + _ => { + // Real nonadjacent initialized slots beyond the stored prefix. + // The two masks differ, then reconverged lane-local selection + // moves each value together with its validity metadata. + let mut other = ThreadData::<_, ITEMS_PER_THREAD>::new(); + let mut item = 0; + while item < ITEMS_PER_THREAD { + let prefix = lane * ITEMS_PER_THREAD + item < valid; + if prefix || item & 1 == 1 { + values.write(item, default); + } + if prefix || item & 1 == 0 { + other.write(item, default); + } + item += 1; + } + values = if lane & 1 == 0 { values } else { other }; + } + } + let before = values; + if mode == 0 || mode == 2 { + coop::store( + block, + &mut *output, + &values, + coop::StoreOptions::default().offset(offset), + ); + } else { + let written = if mode == 5 { + valid + 1 + } else if mode == 6 { + valid / 2 + } else { + valid + }; + coop::store( + block, + &mut *output, + &values, + coop::StoreOptions::default() + .offset(offset) + .valid_items(written), + ); + } + let mut observations_ok = true; + let mut item = 0; + while item < ITEMS_PER_THREAD { + if let Some(value) = before.get(item) { + observations_ok &= values.get(item) == Some(value); + } else { + observations_ok &= values.get(item).is_none(); + } + if mode == 1 || mode == 6 { + observations_ok &= + values.get(item).is_some() == (lane * ITEMS_PER_THREAD + item < valid); + } + item += 1; + } + if mode == 4 { + thread::sync_threads(); + coop::load( + block, + &mut *output, + &mut values, + coop::LoadOptions::default() + .offset(offset) + .valid_items(valid) + .oob_default(default), + ); + thread::sync_threads(); + coop::store( + block, + &mut *output, + &values, + coop::StoreOptions::default() + .offset(offset) + .valid_items(valid), + ); + } + assert!(observations_ok); + // Final writes complete through the owned async host operation. + } + + #[kernel] + #[launch_contract(domain = 1, block = (128, 1, 1), dynamic_shared = inferred)] + pub fn copy_i8_128( + input: &[i8], + mut output: DisjointSlice>, + valid: usize, + shift: usize, + default: i8, + mode: u32, + ) { + exercise::<_, ITEMS_PER_THREAD>(input, &mut output, valid, shift, default, mode); + } + + #[kernel] + #[launch_contract(domain = 1, block = (128, 1, 1), dynamic_shared = inferred)] + pub fn copy_u8_128( + input: &[u8], + mut output: DisjointSlice>, + valid: usize, + shift: usize, + default: u8, + mode: u32, + ) { + exercise::<_, ITEMS_PER_THREAD>(input, &mut output, valid, shift, default, mode); + } + + #[kernel] + #[launch_contract(domain = 1, block = (128, 1, 1), dynamic_shared = inferred)] + pub fn copy_i16_128( + input: &[i16], + mut output: DisjointSlice>, + valid: usize, + shift: usize, + default: i16, + mode: u32, + ) { + exercise::<_, ITEMS_PER_THREAD>(input, &mut output, valid, shift, default, mode); + } + + #[kernel] + #[launch_contract(domain = 1, block = (128, 1, 1), dynamic_shared = inferred)] + pub fn copy_u16_128( + input: &[u16], + mut output: DisjointSlice>, + valid: usize, + shift: usize, + default: u16, + mode: u32, + ) { + exercise::<_, ITEMS_PER_THREAD>(input, &mut output, valid, shift, default, mode); + } + + #[kernel] + #[launch_contract(domain = 1, block = (128, 1, 1), dynamic_shared = inferred)] + pub fn copy_i32_128( + input: &[i32], + mut output: DisjointSlice>, + valid: usize, + shift: usize, + default: i32, + mode: u32, + ) { + exercise::<_, ITEMS_PER_THREAD>(input, &mut output, valid, shift, default, mode); + } + + #[kernel] + #[launch_contract(domain = 1, block = (128, 1, 1), dynamic_shared = inferred)] + pub fn copy_u32_128( + input: &[u32], + mut output: DisjointSlice>, + valid: usize, + shift: usize, + default: u32, + mode: u32, + ) { + exercise::<_, ITEMS_PER_THREAD>(input, &mut output, valid, shift, default, mode); + } + + #[kernel] + #[launch_contract(domain = 1, block = (128, 1, 1), dynamic_shared = inferred)] + pub fn copy_i64_128( + input: &[i64], + mut output: DisjointSlice>, + valid: usize, + shift: usize, + default: i64, + mode: u32, + ) { + exercise::<_, ITEMS_PER_THREAD>(input, &mut output, valid, shift, default, mode); + } + + #[kernel] + #[launch_contract(domain = 1, block = (128, 1, 1), dynamic_shared = inferred)] + pub fn copy_u64_128( + input: &[u64], + mut output: DisjointSlice>, + valid: usize, + shift: usize, + default: u64, + mode: u32, + ) { + exercise::<_, ITEMS_PER_THREAD>(input, &mut output, valid, shift, default, mode); + } + + #[kernel] + #[launch_contract(domain = 1, block = (128, 1, 1), dynamic_shared = inferred)] + pub fn copy_f32_128( + input: &[f32], + mut output: DisjointSlice>, + valid: usize, + shift: usize, + default: f32, + mode: u32, + ) { + exercise::<_, ITEMS_PER_THREAD>(input, &mut output, valid, shift, default, mode); + } + + #[kernel] + #[launch_contract(domain = 1, block = (128, 1, 1), dynamic_shared = inferred)] + pub fn copy_f64_128( + input: &[f64], + mut output: DisjointSlice>, + valid: usize, + shift: usize, + default: f64, + mode: u32, + ) { + exercise::<_, ITEMS_PER_THREAD>(input, &mut output, valid, shift, default, mode); + } + + #[kernel] + #[launch_contract(domain = 1, block = (64, 1, 1), dynamic_shared = inferred)] + pub fn copy_i32_64( + input: &[i32], + mut output: DisjointSlice>, + valid: usize, + shift: usize, + default: i32, + mode: u32, + ) { + exercise::<_, ITEMS_PER_THREAD>(input, &mut output, valid, shift, default, mode); + } + #[kernel] + #[launch_contract(domain = 2, block = (11, 3, 1), dynamic_shared = inferred)] + pub fn copy_i32_2d( + input: &[i32], + mut output: DisjointSlice>, + valid: usize, + shift: usize, + default: i32, + mode: u32, + ) { + exercise::<_, ITEMS_PER_THREAD>(input, &mut output, valid, shift, default, mode); + } + #[kernel] + #[launch_contract(domain = 3, block = (8, 4, 4), dynamic_shared = inferred)] + pub fn copy_i32_3d( + input: &[i32], + mut output: DisjointSlice>, + valid: usize, + shift: usize, + default: i32, + mode: u32, + ) { + exercise::<_, ITEMS_PER_THREAD>(input, &mut output, valid, shift, default, mode); + } + #[kernel] + #[launch_contract(domain = 1, block = (128, 1, 1), dynamic_shared = inferred)] + pub fn scalar(mut output: DisjointSlice>, value: i16) { + let block = this_block(); + let offset = block.rank_at::() as usize * block.count() as usize; + coop::store( + block, + &mut output, + value, + coop::StoreOptions::default().offset(offset), + ); + thread::sync_threads(); + let mut payload = ThreadData::<_, 1>::new(); + coop::load( + block, + &mut output, + &mut payload, + coop::LoadOptions::default().offset(offset), + ); + let scalar_preserved = payload.get(0) == Some(value); + thread::sync_threads(); + coop::load( + block, + &mut output, + &mut payload, + coop::LoadOptions::default().offset(offset).valid_items(1), + ); + assert!(scalar_preserved && payload.get(0).is_some() == (thread::threadIdx_x() == 0)); + } +} + +fn main() -> Result<(), Box> { + if std::env::args().any(|arg| arg == "logical-plain") { + return logical_plain::run(); + } + if std::env::args().any(|arg| arg == "logical-transpose") { + return logical_transpose::run(); + } + if matches!( + std::env::args().nth(1).as_deref(), + Some("--warp" | "--warp-grain-trap") + ) { + return warp::run(); + } + if let Some(argument) = std::env::args().nth(1) { + if matches!(argument.as_str(), "logical-warp" | "logical-grain-trap") { + return logical_warp::run(); + } + if argument == "storage" { + return storage::run(); + } + if argument == "block-warp" { + return block_warp::run(None); + } + if argument.starts_with("block-warp-") { + return block_warp::run(Some(&argument)); + } + if argument == "vectorize" { + return vectorize::run(None); + } + if argument.starts_with("vectorize-") { + return vectorize::run(Some(&argument)); + } + if argument == "striped" { + return striped::run(None); + } + if argument.starts_with("striped-") { + return striped::run(Some(&argument)); + } + if argument == "transpose" { + return transpose::run(None); + } + if argument.starts_with("transpose-") { + return transpose::run(Some(&argument)); + } + } + + use cuda_async::simt::device_context::{init_device_contexts, with_cuda_context}; + use cuda_async::simt::device_operation::DeviceOperation; + use cuda_core::{DeviceBuffer, LaunchConfig1D, LaunchConfig2D, LaunchConfig3D}; + init_device_contexts(0, 1)?; + let context = with_cuda_context(0, |context| context.clone())?; + let stream = context.default_stream(); + // SAFETY: embedded kernels and these typed declarations are compiled together. + let module = unsafe { kernels::load_async(0)? }; + let trap = std::env::args().nth(1); + let mut cases = 0; + macro_rules! run { + ($t:ty, $prepare:ident, $launch:ident, $block:literal, $items:literal, $default:expr) => { + run!( + $t, + $prepare, + $launch, + $block, + $items, + $default, + |groups| LaunchConfig1D::auto_shared(groups, $block), + |groups| LaunchConfig1D::auto_shared(groups, $block / 2) + ); + }; + ($t:ty, $prepare:ident, $launch:ident, $block:literal, $items:literal, $default:expr, + $config:expr, $wrong_shape:expr) => {{ + let default: $t = $default; + for groups in [1_u32, 3] { + let prepared = module.$prepare::<$items>(($config)(groups))?; + assert_eq!(prepared.__raw_config().shared_mem_bytes, 0); + let error = module + .$prepare::<$items>(($wrong_shape)(groups)) + .err() + .expect("a different exact block shape must reject"); + assert!(error.to_string().contains("does not match required block")); + let (gx, gy, gz) = prepared.__raw_config().grid_dim; + let groups = gx * gy * gz; + for mode in [0, 1, 2, 3, 4, 6] { + for valid in [0, 1, $block * $items - 1, $block * $items] { + if (mode == 0) && valid != $block * $items { + continue; + } + let shift = if $items == 1 { 7 } else { 0 }; + let end = (groups as usize - 1) * $block * 4 + shift; + let source_len = end + + if mode <= 1 || mode == 6 { + $block * $items + } else { + valid + }; + let target_len = end + if mode <= 2 { $block * $items } else { valid }; + let input: Vec<$t> = (0..source_len).map(|i| (i % 61 + 1) as $t).collect(); + let source = DeviceBuffer::from_host(&stream, &input)?; + let target = DeviceBuffer::from_host(&stream, &vec![0 as $t; target_len])?; + let (_, target) = module + .$launch::<$items, _, _>( + &prepared, source, target, valid, shift, default, mode, + ) + .sync()?; + let mut result = vec![0 as $t; target_len]; + target.copy_to_host(&stream, &mut result)?; + stream.synchronize()?; + for (index, &value) in result.iter().enumerate() { + let within = index % ($block * 4); + let written = if mode == 0 || mode == 2 { + $block * $items + } else if mode == 6 { + valid / 2 + } else { + valid + }; + let expected = if within >= shift && within - shift < written { + if mode == 3 || (mode == 2 && within - shift >= valid) { + default + } else { + input[index] + } + } else { + 0 as $t + }; + assert_eq!( + value.to_ne_bytes(), + expected.to_ne_bytes(), + "type={} items={} mode={} valid={} index={}", + stringify!($t), + $items, + mode, + valid, + index + ); + } + cases += 1; + } + } + } + }}; + } + if let Some(trap) = trap { + let plan = module.prepare_copy_i32_128::<4>(LaunchConfig1D::auto_shared(1, 128))?; + let len = if trap == "ownership" { 1024 } else { 512 }; + let source = DeviceBuffer::from_host(&stream, &vec![11_i32; len])?; + let target = DeviceBuffer::from_host(&stream, &vec![0_i32; len])?; + let (valid, shift, mode) = match trap.as_str() { + "unwritten" => (512, 0, 9), + "missing-default-count" => (512, 0, 10), + "invalid-read" => (0, 0, 11), + "count" => (513, 0, 1), + "bounds" => (512, 1, 1), + "ownership" => (512, 1, 0), + "different-count" => (127, 0, 5), + _ => panic!("unknown trap mode"), + }; + let _ = module + .copy_i32_128_async_owned::<4, _, _>(&plan, source, target, valid, shift, -7, mode) + .sync()?; + panic!("expected trap returned successfully"); + } + run!( + i8, + prepare_copy_i8_128, + copy_i8_128_async_owned, + 128, + 1, + i8::MIN + ); + run!( + i8, + prepare_copy_i8_128, + copy_i8_128_async_owned, + 128, + 4, + i8::MIN + ); + run!( + u8, + prepare_copy_u8_128, + copy_u8_128_async_owned, + 128, + 1, + u8::MAX + ); + run!( + u8, + prepare_copy_u8_128, + copy_u8_128_async_owned, + 128, + 4, + u8::MAX + ); + run!( + i16, + prepare_copy_i16_128, + copy_i16_128_async_owned, + 128, + 1, + i16::MIN + ); + run!( + i16, + prepare_copy_i16_128, + copy_i16_128_async_owned, + 128, + 4, + i16::MIN + ); + run!( + u16, + prepare_copy_u16_128, + copy_u16_128_async_owned, + 128, + 1, + u16::MAX + ); + run!( + u16, + prepare_copy_u16_128, + copy_u16_128_async_owned, + 128, + 4, + u16::MAX + ); + run!( + i32, + prepare_copy_i32_128, + copy_i32_128_async_owned, + 128, + 1, + i32::MIN + ); + run!( + i32, + prepare_copy_i32_128, + copy_i32_128_async_owned, + 128, + 4, + i32::MIN + ); + run!( + u32, + prepare_copy_u32_128, + copy_u32_128_async_owned, + 128, + 1, + u32::MAX + ); + run!( + u32, + prepare_copy_u32_128, + copy_u32_128_async_owned, + 128, + 4, + u32::MAX + ); + run!( + i64, + prepare_copy_i64_128, + copy_i64_128_async_owned, + 128, + 1, + i64::MIN + 17 + ); + run!( + i64, + prepare_copy_i64_128, + copy_i64_128_async_owned, + 128, + 4, + i64::MIN + 17 + ); + run!( + u64, + prepare_copy_u64_128, + copy_u64_128_async_owned, + 128, + 1, + u64::MAX - 17 + ); + run!( + u64, + prepare_copy_u64_128, + copy_u64_128_async_owned, + 128, + 4, + u64::MAX - 17 + ); + run!( + f32, + prepare_copy_f32_128, + copy_f32_128_async_owned, + 128, + 1, + -0.0_f32 + ); + run!( + f32, + prepare_copy_f32_128, + copy_f32_128_async_owned, + 128, + 4, + -0.0_f32 + ); + run!( + f64, + prepare_copy_f64_128, + copy_f64_128_async_owned, + 128, + 1, + -0.0_f64 + ); + run!( + f64, + prepare_copy_f64_128, + copy_f64_128_async_owned, + 128, + 4, + -0.0_f64 + ); + run!( + i32, + prepare_copy_i32_64, + copy_i32_64_async_owned, + 64, + 1, + -1234567 + ); + run!( + i32, + prepare_copy_i32_64, + copy_i32_64_async_owned, + 64, + 4, + -1234567 + ); + run!( + i32, + prepare_copy_i32_2d, + copy_i32_2d_async_owned, + 33, + 1, + -7, + |groups| LaunchConfig2D::auto_shared((groups, 2), (11, 3)), + |groups| LaunchConfig2D::auto_shared((groups, 2), (3, 11)) + ); + run!( + i32, + prepare_copy_i32_3d, + copy_i32_3d_async_owned, + 128, + 1, + -7, + |groups| LaunchConfig3D::auto_shared((groups, 2, 2), (8, 4, 4)), + |groups| LaunchConfig3D::auto_shared((groups, 2, 2), (16, 2, 4)) + ); + run!( + i32, + prepare_copy_i32_2d, + copy_i32_2d_async_owned, + 33, + 4, + -7, + |groups| LaunchConfig2D::auto_shared((groups, 2), (11, 3)), + |groups| LaunchConfig2D::auto_shared((groups, 2), (3, 11)) + ); + run!( + i32, + prepare_copy_i32_3d, + copy_i32_3d_async_owned, + 128, + 4, + -7, + |groups| LaunchConfig3D::auto_shared((groups, 2, 2), (8, 4, 4)), + |groups| LaunchConfig3D::auto_shared((groups, 2, 2), (16, 2, 4)) + ); + let plan = module.prepare_scalar(LaunchConfig1D::auto_shared(3, 128))?; + let target = DeviceBuffer::from_host(&stream, &vec![0_i16; 384])?; + let target = module + .scalar_async_owned::<_>(&plan, target, -17_i16) + .sync()?; + let mut result = vec![0_i16; 384]; + target.copy_to_host(&stream, &mut result)?; + stream.synchronize()?; + assert!(result.iter().all(|value| *value == -17)); + cases += 1; + println!( + "PASS: {cases} public Direct cases, ten numeric types, scattered validity, short tails, owned async" + ); + Ok(()) +} diff --git a/rust/cuda-coop/tests/load_store/src/storage.rs b/rust/cuda-coop/tests/load_store/src/storage.rs new file mode 100644 index 000000000000..36435cd6d454 --- /dev/null +++ b/rust/cuda-coop/tests/load_store/src/storage.rs @@ -0,0 +1,357 @@ +// Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +//! Public storage policies and observable cross-lane reservation access. +use cuda_coop as coop; +use cuda_device::{CollectiveTiles, DisjointSlice, kernel, launch_contract, thread}; +use cuda_host::cuda_module; + +#[cuda_module] +mod kernels { + use super::*; + + #[inline(never)] + fn load_tile( + input: &[i32], + values: &mut coop::ThreadData, + offset: usize, + valid: usize, + scratch: &mut coop::TempStorage, + ) { + coop::load( + coop::this_block(), + input, + values, + coop::LoadOptions::default() + .offset(offset) + .valid_items(valid) + .oob_default(-7) + .temp_storage(scratch) + .algorithm(coop::Transpose), + ); + } + + #[inline(never)] + fn store_tile( + output: &mut DisjointSlice>, + values: &mut coop::ThreadData, + offset: usize, + valid: usize, + scratch: &mut coop::TempStorage, + ) { + coop::store( + coop::this_block(), + output, + values, + coop::StoreOptions::default() + .offset(offset) + .valid_items(valid) + .algorithm(coop::Transpose) + .temp_storage(scratch), + ); + } + + #[inline(never)] + fn descriptor() -> coop::TempStorage { + coop::TempStorage::new() + } + + #[kernel] + #[launch_contract(domain = 1, block = (33, 1, 1), dynamic_shared = inferred)] + pub fn exercise( + input: &[i32], + mut output: DisjointSlice>, + valid: usize, + ) { + let block = coop::this_block(); + let offset = thread::blockIdx_x() as usize * 33 * 4; + let lane = thread::threadIdx_x() as usize; + let mut values = coop::ThreadData::<_, ITEMS>::new(); + match MODE { + 0 => { + coop::load( + block, + input, + &mut values, + coop::LoadOptions::default() + .offset(offset) + .valid_items(valid) + .oob_default(-7) + .algorithm(coop::Transpose), + ); + coop::store( + block, + &mut output, + &mut values, + coop::StoreOptions::default() + .offset(offset) + .valid_items(valid) + .algorithm(coop::Transpose), + ); + } + 1 => { + let mut scratch = descriptor(); + load_tile(input, &mut values, offset, valid, &mut scratch); + thread::sync_threads(); + store_tile(&mut output, &mut values, offset, valid, &mut scratch); + } + 2 => { + let mut scratch = coop::TempStorage::new().auto_sync(); + let mut iteration = 0; + while iteration < 2 + (valid & 1) { + load_tile(input, &mut values, offset, valid, &mut scratch); + store_tile(&mut output, &mut values, offset, valid, &mut scratch); + iteration += 1; + } + } + 3 => { + let mut scratch = coop::TempStorage::new().with_alignment::<64>().exclusive(); + // The same helper at two sites receives distinct exclusive slices. + load_tile(input, &mut values, offset, valid, &mut scratch); + load_tile(input, &mut values, offset, valid, &mut scratch); + store_tile(&mut output, &mut values, offset, valid, &mut scratch); + } + 4 => { + let mut first = descriptor(); + let mut second = descriptor(); + load_tile(input, &mut values, offset, valid, &mut first); + load_tile(input, &mut values, offset, valid, &mut second); + thread::sync_threads(); + store_tile(&mut output, &mut values, offset, valid, &mut first); + } + 5 => { + let mut scratch = descriptor(); + let view = scratch.reserve::(); + // SAFETY: each lane initializes its own in-bounds element, then + // all writers complete before a different lane reads it. + unsafe { core::ptr::write(view.as_mut_ptr().add(lane), lane as i32) }; + thread::sync_threads(); + let saved = unsafe { view.as_mut_ptr().add((lane + 1) % 33).read() }; + view.release(); + thread::sync_threads(); + load_tile(input, &mut values, offset, valid, &mut scratch); + let mut item = 0; + while item < ITEMS { + // Filled Load initialized every slot; the saved cross-lane + // value remains observable after the shared region is reused. + let value = values.get(item).unwrap_or(-100_000); + values.write(item, value + saved); + item += 1; + } + thread::sync_threads(); + store_tile(&mut output, &mut values, offset, valid, &mut scratch); + } + 6 => { + let mut scratch = coop::TempStorage::new().exclusive(); + let first = scratch.reserve::(); + let second = scratch.reserve_aligned::(); + // SAFETY: distinct exclusive sites do not overlap. Each lane + // writes its own slots before the full-block completion barrier. + unsafe { + core::ptr::write(first.as_mut_ptr().add(lane), lane as i32); + core::ptr::write(second.as_mut_ptr().add(lane), (2 * lane + 1) as u64); + } + thread::sync_threads(); + let neighbor = (lane + 1) % 33; + let saved = unsafe { + first.as_mut_ptr().add(neighbor).read() + + second.as_mut_ptr().add(neighbor).read() as i32 + }; + first.release(); + second.release(); + thread::sync_threads(); + load_tile(input, &mut values, offset, valid, &mut scratch); + let mut item = 0; + while item < ITEMS { + let value = values.get(item).unwrap_or(-100_000); + values.write(item, value + saved); + item += 1; + } + store_tile(&mut output, &mut values, offset, valid, &mut scratch); + } + 7 => { + let mut scratch = coop::TempStorage::new().auto_sync(); + coop::load( + block, + input, + &mut values, + coop::LoadOptions::default() + .offset(offset) + .valid_items(valid) + .oob_default(-7) + .temp_storage(&mut scratch), + ); + coop::store( + block, + &mut output, + &values, + coop::StoreOptions::default() + .offset(offset) + .valid_items(valid) + .temp_storage(&mut scratch), + ); + } + _ => { + let mut scratch = coop::TempStorage::new() + .with_capacity::<4096>() + .with_alignment::<64>() + .auto_sync(); + load_tile(input, &mut values, offset, valid, &mut scratch); + store_tile(&mut output, &mut values, offset, valid, &mut scratch); + } + } + } + + #[kernel] + #[launch_contract(domain = 1, block = (64, 1, 1), dynamic_shared = inferred)] + pub fn timesliced_reuse( + input: &[i32], + mut output: DisjointSlice>, + valid: usize, + ) { + let block = coop::this_block(); + let offset = thread::blockIdx_x() as usize * 256; + let lane = thread::threadIdx_x() as usize; + let mut values = coop::ThreadData::<_, 4>::new(); + let mut scratch = coop::TempStorage::new().auto_sync(); + let mut correct = true; + let mut iteration = 0; + while iteration < 2 { + coop::load( + block, + input, + &mut values, + coop::LoadOptions::default() + .offset(offset) + .valid_items(valid) + .algorithm(coop::WarpTransposeTimesliced) + .temp_storage(&mut scratch), + ); + coop::store( + block, + &mut output, + &mut values, + coop::StoreOptions::default() + .offset(offset) + .valid_items(valid) + .algorithm(coop::WarpTransposeTimesliced) + .temp_storage(&mut scratch), + ); + let mut item = 0; + while item < 4 { + let index = (lane / 32) * 128 + 32 * item + lane % 32; + let expected = if index < valid { + Some((offset + index) as i32 * 7 + 3) + } else { + None + }; + correct &= values.get(item) == expected; + item += 1; + } + iteration += 1; + } + assert!(correct); + } +} + +pub fn run() -> Result<(), Box> { + use cuda_async::simt::device_context::{init_device_contexts, with_cuda_context}; + use cuda_async::simt::device_operation::DeviceOperation; + use cuda_core::{DeviceBuffer, LaunchConfig1D}; + + init_device_contexts(0, 1)?; + let context = with_cuda_context(0, |context| context.clone())?; + let stream = context.default_stream(); + // SAFETY: the embedded module and typed declarations are built together. + let module = unsafe { kernels::load_async(0)? }; + let input: Vec = (0..3 * 33 * 4).map(|i| i * 7 + 3).collect(); + let mut cases = 0; + macro_rules! run { + ($items:literal, $mode:literal) => {{ + let plan = + module.prepare_exercise::<$items, $mode>(LaunchConfig1D::auto_shared(3, 33))?; + let scratch = plan.__raw_config().shared_mem_bytes; + if $mode == 7 { + assert_eq!(scratch, 0); + } else if $mode == 8 { + assert_eq!(scratch, 4096); + } else { + assert!(scratch > 0); + } + println!("storage mode={} items={} shared={scratch}", $mode, $items); + for valid in [0, 1, 33 * $items - 1, 33 * $items] { + let source = DeviceBuffer::from_host(&stream, &input)?; + let target = DeviceBuffer::from_host(&stream, &vec![-19_i32; input.len()])?; + let (_, target) = module + .exercise_async_owned::<$items, $mode, _, _>(&plan, source, target, valid) + .sync()?; + let mut output = vec![0; input.len()]; + target.copy_to_host(&stream, &mut output)?; + stream.synchronize()?; + for (index, actual) in output.into_iter().enumerate() { + let local = index % (33 * 4); + let expected = if local < valid { + let neighbor = (local / $items + 1) % 33; + input[index] + + if $mode == 5 { + neighbor as i32 + } else if $mode == 6 { + (3 * neighbor + 1) as i32 + } else { + 0 + } + } else { + -19 + }; + assert_eq!( + actual, expected, + "mode={} items={} valid={valid} index={index}", + $mode, $items + ); + } + cases += 1; + } + }}; + } + macro_rules! modes { + ($items:literal) => { + run!($items, 0); + run!($items, 1); + run!($items, 2); + run!($items, 3); + run!($items, 4); + run!($items, 5); + run!($items, 6); + run!($items, 7); + run!($items, 8); + }; + } + modes!(1); + modes!(4); + // Counted timesliced Store needs its full, non-timesliced validity plane: + // 528 CUB bytes plus 32 mask bytes for two warps at four items per thread. + let plan = module.prepare_timesliced_reuse(LaunchConfig1D::auto_shared(3, 64))?; + assert_eq!(plan.__raw_config().shared_mem_bytes, 560); + let input: Vec = (0..3 * 256).map(|i| i * 7 + 3).collect(); + for valid in [0, 1, 255, 256] { + let source = DeviceBuffer::from_host(&stream, &input)?; + let target = DeviceBuffer::from_host(&stream, &vec![-19_i32; input.len()])?; + let (_, target) = module + .timesliced_reuse_async_owned(&plan, source, target, valid) + .sync()?; + let mut output = vec![0; input.len()]; + target.copy_to_host(&stream, &mut output)?; + stream.synchronize()?; + for (index, actual) in output.into_iter().enumerate() { + let expected = if index % 256 < valid { + input[index] + } else { + -19 + }; + assert_eq!(actual, expected, "timesliced valid={valid} index={index}"); + } + cases += 1; + } + println!("PASS: {cases} public storage cases"); + Ok(()) +} diff --git a/rust/cuda-coop/tests/load_store/src/striped.rs b/rust/cuda-coop/tests/load_store/src/striped.rs new file mode 100644 index 000000000000..003de60cc877 --- /dev/null +++ b/rust/cuda-coop/tests/load_store/src/striped.rs @@ -0,0 +1,415 @@ +// Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +//! Striped payload mapping and preservation, observed independently of roundtrip I/O. +use cuda_coop as coop; +use cuda_device::{CollectiveTiles, DisjointSlice, kernel, launch_contract, thread}; +use cuda_host::cuda_module; + +#[cuda_module] +mod kernels { + use super::*; + + #[kernel] + #[launch_contract(domain = 1, block = (33, 1, 1), dynamic_shared = inferred)] + pub fn copy_33( + input: &[i32], + mut output: DisjointSlice>, + ) { + let block = coop::this_block(); + let lane = thread::threadIdx_x() as usize; + let offset = thread::blockIdx_x() as usize * 33 * ITEMS_PER_THREAD; + let mut values = coop::ThreadData::<_, ITEMS_PER_THREAD>::new(); + coop::load( + block, + input, + &mut values, + coop::LoadOptions::default() + .offset(offset) + .algorithm(coop::Striped), + ); + let before = values; + coop::store( + block, + &mut output, + &values, + coop::StoreOptions::default() + .algorithm(coop::Striped) + .offset(offset), + ); + let mut observations_ok = true; + let mut item = 0; + while item < ITEMS_PER_THREAD { + let expected = input[offset + lane + 33 * item]; + observations_ok &= before.get(item) == Some(expected); + observations_ok &= values.get(item) == before.get(item); + item += 1; + } + assert!(observations_ok); + } + + fn exercise( + input: &[T], + output: &mut DisjointSlice>, + valid: usize, + shift: usize, + default: T, + mode: u32, + ) { + let block = coop::this_block(); + let x = thread::blockDim_x() as usize; + let y = thread::blockDim_y() as usize; + let threads = x * y * thread::blockDim_z() as usize; + let lane = thread::threadIdx_x() as usize + + x * (thread::threadIdx_y() as usize + y * thread::threadIdx_z() as usize); + let group = thread::blockIdx_x() as usize + + thread::gridDim_x() as usize + * (thread::blockIdx_y() as usize + + thread::gridDim_y() as usize * thread::blockIdx_z() as usize); + let offset = group * threads * 4 + shift; + let mut values = coop::ThreadData::<_, ITEMS>::new(); + match mode { + 0 => coop::load( + block, + input, + &mut values, + coop::LoadOptions::default() + .algorithm(coop::Striped) + .offset(offset), + ), + 1 | 5 | 6 => { + // A second masked Load must clear validity left by a full Load. + coop::load( + block, + input, + &mut values, + coop::LoadOptions::default() + .algorithm(coop::Striped) + .offset(offset), + ); + coop::load( + block, + input, + &mut values, + coop::LoadOptions::default() + .algorithm(coop::Striped) + .offset(offset) + .valid_items(valid), + ); + } + 2 | 4 => coop::load( + block, + input, + &mut values, + coop::LoadOptions::default() + .algorithm(coop::Striped) + .offset(offset) + .valid_items(valid) + .oob_default(default), + ), + 7 => {} + _ => { + let mut item = 0; + while item < ITEMS { + let selected = lane + threads * item < valid; + if (selected || (lane + item) & 1 == 0) + && !(mode == 8 && lane + threads * item == valid - 1) + { + values.write(item, default); + } + item += 1; + } + } + } + let before = values; + if mode == 0 || mode == 2 { + coop::store( + block, + &mut *output, + &values, + coop::StoreOptions::default() + .algorithm(coop::Striped) + .offset(offset), + ); + } else { + let written = if mode == 5 { + valid / 2 + } else if mode == 6 { + valid + 1 + } else { + valid + }; + coop::store( + block, + &mut *output, + &values, + coop::StoreOptions::default() + .algorithm(coop::Striped) + .offset(offset) + .valid_items(written), + ); + } + let mut observations_ok = true; + let mut item = 0; + while item < ITEMS { + let position = lane + threads * item; + let expected = match mode { + 3 | 8 => { + if position < valid || (lane + item) & 1 == 0 { + Some(default) + } else { + None + } + } + 7 => None, + _ => { + if position < valid || mode == 0 { + Some(input[offset + position]) + } else if mode == 2 || mode == 4 { + Some(default) + } else { + None + } + } + }; + observations_ok &= before.get(item) == expected; + observations_ok &= values.get(item) == before.get(item); + item += 1; + } + if mode == 3 || mode == 4 { + thread::sync_threads(); + // A repeated zero Store preserves arbitrary initialized/invalid slots. + coop::store( + block, + &mut *output, + &values, + coop::StoreOptions::default() + .algorithm(coop::Striped) + .offset(offset) + .valid_items(0), + ); + let mut item = 0; + while item < ITEMS { + observations_ok &= values.get(item) == before.get(item); + item += 1; + } + } + assert!(observations_ok); + } + + #[kernel] + #[launch_contract(domain = 1, block = (33, 1, 1), dynamic_shared = inferred)] + pub fn partial_33( + input: &[T], + mut output: DisjointSlice>, + valid: usize, + shift: usize, + default: T, + mode: u32, + ) { + exercise::<_, ITEMS>(input, &mut output, valid, shift, default, mode); + } + + #[kernel] + #[launch_contract(domain = 3, block = (8, 4, 4), dynamic_shared = inferred)] + pub fn partial_3d( + input: &[T], + mut output: DisjointSlice>, + valid: usize, + shift: usize, + default: T, + mode: u32, + ) { + exercise::<_, ITEMS>(input, &mut output, valid, shift, default, mode); + } + + #[kernel] + #[launch_contract(domain = 1, block = (33, 1, 1), dynamic_shared = inferred)] + pub fn striped_scalar(mut output: DisjointSlice>, valid: usize) { + let offset = thread::blockIdx_x() as usize * 33; + coop::store( + coop::this_block(), + &mut output, + 17_i32, + coop::StoreOptions::default() + .algorithm(coop::Striped) + .offset(offset) + .valid_items(valid), + ); + } +} + +pub fn run(trap: Option<&str>) -> Result<(), Box> { + use cuda_async::simt::device_context::{init_device_contexts, with_cuda_context}; + use cuda_async::simt::device_operation::DeviceOperation; + use cuda_core::{DeviceBuffer, LaunchConfig1D, LaunchConfig3D}; + init_device_contexts(0, 1)?; + let context = with_cuda_context(0, |context| context.clone())?; + let stream = context.default_stream(); + // SAFETY: embedded kernels and typed declarations compile together. + let module = unsafe { kernels::load_async(0)? }; + if let Some(trap) = trap { + let mode = if trap == "striped-larger" { 6 } else { 8 }; + let plan = module.prepare_partial_33::(LaunchConfig1D::auto_shared(1, 33))?; + let source = DeviceBuffer::from_host(&stream, &vec![7_i32; 132])?; + let target = DeviceBuffer::from_host(&stream, &vec![0_i32; 132])?; + module + .partial_33_async_owned::(&plan, source, target, 34, 0, -7, mode) + .sync()?; + panic!("invalid Striped Store unexpectedly completed"); + } + let input: Vec = (0..2 * 33 * 4).map(|i| i * 7 + 3).collect(); + let source = DeviceBuffer::from_host(&stream, &input)?; + let target = DeviceBuffer::from_host(&stream, &vec![-1_i32; input.len()])?; + let plan = module.prepare_copy_33::<4>(LaunchConfig1D::auto_shared(2, 33))?; + assert_eq!(plan.__raw_config().shared_mem_bytes, 0); + let (_, target) = module + .copy_33_async_owned::<4, _, _>(&plan, source, target) + .sync()?; + let mut output = vec![0_i32; input.len()]; + target.copy_to_host(&stream, &mut output)?; + stream.synchronize()?; + assert_eq!(output, input); + let mut cases = 1; + macro_rules! run { + ($t:ty, $items:literal, $threads:literal, $prepare:ident, $launch:ident, $config:expr, $extra:expr) => {{ + let tile = $threads * $items; + let plan = module.$prepare::<$t, $items>($config)?; + assert_eq!(plan.__raw_config().shared_mem_bytes, 0); + let prefix = ($threads + 1).min(tile); + let mut selected = vec![ + (0, tile, 0), + (1, prefix, 0), + (2, 1, 0), + (3, prefix, 0), + (5, prefix, 0), + ]; + if $extra { + for valid in [0, 1, $threads - 1, $threads, prefix, tile] { + for mode in [1, 2, 3, 4, 5] { + selected.push((mode, valid, 0)); + } + } + for mode in [1, 3, 4, 5] { + selected.push((mode, tile - 7, 7)); + } + selected.push((7, 0, 0)); + } + selected.sort_unstable(); + selected.dedup(); + for (mode, valid, shift) in selected { + let written = if mode == 0 || mode == 2 { + tile + } else if mode == 5 { + valid / 2 + } else { + valid + }; + // Keep the final allocation shorter than the full tile when + // counted Store writes a prefix. Striped per-lane iterator + // formation must not create an out-of-allocation pointer. + let output_len = $threads * 4 + shift + written; + let input: Vec<$t> = (0..2 * $threads * 4 + shift) + .map(|i| (i % 101) as $t) + .collect(); + let default = 111 as $t; + let sentinel = 123 as $t; + let source = DeviceBuffer::from_host(&stream, &input)?; + let target = DeviceBuffer::from_host(&stream, &vec![sentinel; output_len])?; + let (_, target) = module + .$launch::<$t, $items, _, _>(&plan, source, target, valid, shift, default, mode) + .sync()?; + let mut output = vec![sentinel; output_len]; + target.copy_to_host(&stream, &mut output)?; + stream.synchronize()?; + for (index, actual) in output.into_iter().enumerate() { + let local = index % ($threads * 4); + let expected = if local >= shift && local - shift < written { + if mode == 3 || (mode == 2 && local - shift >= valid) { + default + } else { + input[index] + } + } else { + sentinel + }; + assert_eq!( + actual, + expected, + "type={} items={} mode={mode} valid={valid} shift={shift} index={index}", + stringify!($t), + $items + ); + } + cases += 1; + } + }}; + } + macro_rules! width { + ($t:ty) => { + run!( + $t, + 1, + 33, + prepare_partial_33, + partial_33_async_owned, + LaunchConfig1D::auto_shared(2, 33), + false + ); + run!( + $t, + 4, + 33, + prepare_partial_33, + partial_33_async_owned, + LaunchConfig1D::auto_shared(2, 33), + false + ); + }; + } + width!(i8); + width!(u8); + width!(i16); + width!(u16); + width!(i32); + width!(u32); + width!(i64); + width!(u64); + width!(f32); + width!(f64); + run!( + i32, + 4, + 33, + prepare_partial_33, + partial_33_async_owned, + LaunchConfig1D::auto_shared(2, 33), + true + ); + run!( + i32, + 4, + 128, + prepare_partial_3d, + partial_3d_async_owned, + LaunchConfig3D::auto_shared((1, 2, 1), (8, 4, 4)), + true + ); + for valid in [0, 1, 32, 33] { + let plan = module.prepare_striped_scalar(LaunchConfig1D::auto_shared(2, 33))?; + assert_eq!(plan.__raw_config().shared_mem_bytes, 0); + let target = DeviceBuffer::from_host(&stream, &vec![-1_i32; 33 + valid])?; + let target = module + .striped_scalar_async_owned(&plan, target, valid) + .sync()?; + let mut output = vec![0_i32; 33 + valid]; + target.copy_to_host(&stream, &mut output)?; + stream.synchronize()?; + for (i, value) in output.into_iter().enumerate() { + assert_eq!(value, if i % 33 < valid { 17 } else { -1 }); + } + cases += 1; + } + println!("PASS: {cases} Striped cases with observed payload mapping and validity"); + Ok(()) +} diff --git a/rust/cuda-coop/tests/load_store/src/transpose.rs b/rust/cuda-coop/tests/load_store/src/transpose.rs new file mode 100644 index 000000000000..6fc2cb9fbadb --- /dev/null +++ b/rust/cuda-coop/tests/load_store/src/transpose.rs @@ -0,0 +1,798 @@ +// Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +//! Public Transpose checks observe values and arbitrary validity after each Store. +use cuda_coop as coop; +use cuda_device::{CollectiveTiles, DisjointSlice, kernel, launch_contract, thread}; +use cuda_host::cuda_module; + +#[cuda_module] +mod kernels { + use super::*; + + fn expected( + input: &[T], + offset: usize, + index: usize, + items: usize, + valid: usize, + default: T, + mode: u32, + ) -> Option { + if mode == 1 || mode == 5 || mode == 6 { + if index < valid { + Some(input[offset + index]) + } else { + None + } + } else if mode == 2 || mode == 4 { + if index < valid { + Some(input[offset + index]) + } else { + Some(default) + } + } else if mode == 3 { + let lane = index / items; + let item = index % items; + if index < valid || (lane & 1 == 0 && item & 1 == 1) || (lane & 1 == 1 && item & 1 == 0) + { + Some(input[offset + index]) + } else { + None + } + } else { + Some(input[offset + index]) + } + } + + #[inline(never)] + fn exercise( + input: &[T], + output: &mut DisjointSlice>, + valid: usize, + default: T, + mode: u32, + ) { + let block = coop::this_block(); + let x = thread::blockDim_x() as usize; + let y = thread::blockDim_y() as usize; + let threads = x * y * thread::blockDim_z() as usize; + let lane = thread::threadIdx_x() as usize + + x * (thread::threadIdx_y() as usize + y * thread::threadIdx_z() as usize); + let group = thread::blockIdx_x() as usize + + thread::gridDim_x() as usize + * (thread::blockIdx_y() as usize + + thread::gridDim_y() as usize * thread::blockIdx_z() as usize); + let offset = group * threads * CAPACITY; + let mut values = coop::ThreadData::<_, ITEMS>::new(); + if mode != 9 { + coop::load( + block, + input, + &mut values, + coop::LoadOptions::default() + .offset(offset) + .algorithm(coop::Transpose), + ); + } + if mode == 1 || mode == 5 || mode == 6 { + coop::load( + block, + input, + &mut values, + coop::LoadOptions::default() + .offset(offset) + .valid_items(valid) + .algorithm(coop::Transpose), + ); + } else if mode == 2 || mode == 4 { + coop::load( + block, + input, + &mut values, + coop::LoadOptions::default() + .offset(offset) + .valid_items(valid) + .oob_default(default) + .algorithm(coop::Transpose), + ); + } else if mode == 3 { + let incoming = values; + values = coop::ThreadData::new(); + let mut item = 0; + while item < ITEMS { + let index = lane * ITEMS + item; + let initialize = index < valid + || (lane & 1 == 0 && item & 1 == 1) + || (lane & 1 == 1 && item & 1 == 0); + if let (true, Some(value)) = (initialize, incoming.get(item)) { + values.write(item, value); + } + item += 1; + } + } + if mode == 0 || mode == 2 { + coop::store( + block, + &mut *output, + &mut values, + coop::StoreOptions::default() + .offset(offset) + .algorithm(coop::Transpose), + ); + } else { + let written = if mode == 5 { + valid + 1 + } else if mode == 6 { + valid / 2 + } else { + valid + }; + coop::store( + block, + &mut *output, + &mut values, + coop::StoreOptions::default() + .offset(offset) + .valid_items(written) + .algorithm(coop::Transpose), + ); + } + let mut correct = true; + let mut item = 0; + while item < ITEMS { + let source = threads * item + lane; + correct &= + values.get(item) == expected(input, offset, source, ITEMS, valid, default, mode); + item += 1; + } + // Managed scratch reuse inserts completion before this second operation. + // A zero count still permutes all values and bits, with no global writes. + coop::store( + block, + &mut *output, + &mut values, + coop::StoreOptions::default() + .offset(offset) + .valid_items(0) + .algorithm(coop::Transpose), + ); + item = 0; + while item < ITEMS { + let first = threads * item + lane; + let source = threads * (first % ITEMS) + first / ITEMS; + correct &= + values.get(item) == expected(input, offset, source, ITEMS, valid, default, mode); + item += 1; + } + if mode == 4 { + thread::sync_threads(); + coop::load( + block, + &mut *output, + &mut values, + coop::LoadOptions::default() + .offset(offset) + .valid_items(valid) + .oob_default(default) + .algorithm(coop::Transpose), + ); + coop::store( + block, + &mut *output, + &mut values, + coop::StoreOptions::default() + .offset(offset) + .valid_items(valid) + .algorithm(coop::Transpose), + ); + item = 0; + while item < ITEMS { + correct &= values.get(item) + == expected( + input, + offset, + threads * item + lane, + ITEMS, + valid, + default, + mode, + ); + item += 1; + } + } + assert!(correct); + } + #[kernel] + #[launch_contract(domain=1,block=(33,1,1),dynamic_shared=inferred)] + pub fn copy_i8_33_4( + input: &[i8], + mut output: DisjointSlice>, + valid: usize, + default: i8, + mode: u32, + ) { + exercise::<_, ITEMS_PER_THREAD, 4>(input, &mut output, valid, default, mode); + } + #[kernel] + #[launch_contract(domain=1,block=(33,1,1),dynamic_shared=inferred)] + pub fn copy_u8_33_4( + input: &[u8], + mut output: DisjointSlice>, + valid: usize, + default: u8, + mode: u32, + ) { + exercise::<_, ITEMS_PER_THREAD, 4>(input, &mut output, valid, default, mode); + } + #[kernel] + #[launch_contract(domain=1,block=(33,1,1),dynamic_shared=inferred)] + pub fn copy_i16_33_4( + input: &[i16], + mut output: DisjointSlice>, + valid: usize, + default: i16, + mode: u32, + ) { + exercise::<_, ITEMS_PER_THREAD, 4>(input, &mut output, valid, default, mode); + } + #[kernel] + #[launch_contract(domain=1,block=(33,1,1),dynamic_shared=inferred)] + pub fn copy_u16_33_4( + input: &[u16], + mut output: DisjointSlice>, + valid: usize, + default: u16, + mode: u32, + ) { + exercise::<_, ITEMS_PER_THREAD, 4>(input, &mut output, valid, default, mode); + } + #[kernel] + #[launch_contract(domain=1,block=(33,1,1),dynamic_shared=inferred)] + pub fn copy_i32_33_4( + input: &[i32], + mut output: DisjointSlice>, + valid: usize, + default: i32, + mode: u32, + ) { + exercise::<_, ITEMS_PER_THREAD, 4>(input, &mut output, valid, default, mode); + } + #[kernel] + #[launch_contract(domain=1,block=(33,1,1),dynamic_shared=inferred)] + pub fn copy_u32_33_4( + input: &[u32], + mut output: DisjointSlice>, + valid: usize, + default: u32, + mode: u32, + ) { + exercise::<_, ITEMS_PER_THREAD, 4>(input, &mut output, valid, default, mode); + } + #[kernel] + #[launch_contract(domain=1,block=(33,1,1),dynamic_shared=inferred)] + pub fn copy_i64_33_4( + input: &[i64], + mut output: DisjointSlice>, + valid: usize, + default: i64, + mode: u32, + ) { + exercise::<_, ITEMS_PER_THREAD, 4>(input, &mut output, valid, default, mode); + } + #[kernel] + #[launch_contract(domain=1,block=(33,1,1),dynamic_shared=inferred)] + pub fn copy_u64_33_4( + input: &[u64], + mut output: DisjointSlice>, + valid: usize, + default: u64, + mode: u32, + ) { + exercise::<_, ITEMS_PER_THREAD, 4>(input, &mut output, valid, default, mode); + } + #[kernel] + #[launch_contract(domain=1,block=(33,1,1),dynamic_shared=inferred)] + pub fn copy_f32_33_4( + input: &[f32], + mut output: DisjointSlice>, + valid: usize, + default: f32, + mode: u32, + ) { + exercise::<_, ITEMS_PER_THREAD, 4>(input, &mut output, valid, default, mode); + } + #[kernel] + #[launch_contract(domain=1,block=(33,1,1),dynamic_shared=inferred)] + pub fn copy_f64_33_4( + input: &[f64], + mut output: DisjointSlice>, + valid: usize, + default: f64, + mode: u32, + ) { + exercise::<_, ITEMS_PER_THREAD, 4>(input, &mut output, valid, default, mode); + } + #[kernel] + #[launch_contract(domain=1,block=(1,1,1),dynamic_shared=inferred)] + pub fn copy_i32_1_4( + input: &[i32], + mut output: DisjointSlice>, + valid: usize, + default: i32, + mode: u32, + ) { + exercise::<_, ITEMS_PER_THREAD, 4>(input, &mut output, valid, default, mode); + } + #[kernel] + #[launch_contract(domain=1,block=(128,1,1),dynamic_shared=inferred)] + pub fn copy_i32_128_4( + input: &[i32], + mut output: DisjointSlice>, + valid: usize, + default: i32, + mode: u32, + ) { + exercise::<_, ITEMS_PER_THREAD, 4>(input, &mut output, valid, default, mode); + } + #[kernel] + #[launch_contract(domain=1,block=(130,1,1),dynamic_shared=inferred)] + pub fn copy_i32_130_4( + input: &[i32], + mut output: DisjointSlice>, + valid: usize, + default: i32, + mode: u32, + ) { + exercise::<_, ITEMS_PER_THREAD, 4>(input, &mut output, valid, default, mode); + } + #[kernel] + #[launch_contract(domain=1,block=(33,1,1),dynamic_shared=inferred)] + pub fn copy_i32_33_9( + input: &[i32], + mut output: DisjointSlice>, + valid: usize, + default: i32, + mode: u32, + ) { + exercise::<_, ITEMS_PER_THREAD, 9>(input, &mut output, valid, default, mode); + } + #[kernel] + #[launch_contract(domain=2,block=(11,3,1),dynamic_shared=inferred)] + pub fn copy_i32_2d( + input: &[i32], + mut output: DisjointSlice>, + valid: usize, + default: i32, + mode: u32, + ) { + exercise::<_, ITEMS_PER_THREAD, 4>(input, &mut output, valid, default, mode); + } + #[kernel] + #[launch_contract(domain=3,block=(8,4,4),dynamic_shared=inferred)] + pub fn copy_i32_3d( + input: &[i32], + mut output: DisjointSlice>, + valid: usize, + default: i32, + mode: u32, + ) { + exercise::<_, ITEMS_PER_THREAD, 4>(input, &mut output, valid, default, mode); + } + #[kernel] + #[launch_contract(domain=1,block=(33,1,1),dynamic_shared=inferred)] + pub fn transpose_scalar(mut output: DisjointSlice>) { + coop::store( + coop::this_block(), + &mut output, + -17_i16, + coop::StoreOptions::default().algorithm(coop::Transpose), + ); + } +} + +pub fn run(trap: Option<&str>) -> Result<(), Box> { + use cuda_async::simt::device_context::{init_device_contexts, with_cuda_context}; + use cuda_async::simt::device_operation::DeviceOperation; + use cuda_core::{DeviceBuffer, LaunchConfig1D, LaunchConfig2D, LaunchConfig3D}; + init_device_contexts(0, 1)?; + let context = with_cuda_context(0, |context| context.clone())?; + let stream = context.default_stream(); + // SAFETY: these typed declarations and native kernels are compiled together. + let module = unsafe { kernels::load_async(0)? }; + if let Some(trap) = trap { + let plan = module.prepare_copy_i32_33_4::<4>(LaunchConfig1D::auto_shared(1, 33))?; + let source = DeviceBuffer::from_host(&stream, &vec![11_i32; 132])?; + let target = DeviceBuffer::from_host(&stream, &vec![0_i32; 132])?; + let (valid, mode) = match trap { + "transpose-count" => (133, 1), + "transpose-unwritten" => (132, 9), + "transpose-different-count" => (31, 5), + _ => panic!("unknown Transpose trap"), + }; + let _ = module + .copy_i32_33_4_async_owned::<4, _, _>(&plan, source, target, valid, -7, mode) + .sync()?; + panic!("expected Transpose trap returned successfully"); + } + let mut cases = 0; + macro_rules! run { + ($t:ty,$prepare:ident,$launch:ident,$block:literal,$capacity:literal,$items:literal,$default:expr) => { + run!( + $t, + $prepare, + $launch, + $block, + $capacity, + $items, + $default, + |groups| LaunchConfig1D::auto_shared(groups, $block), + |groups| LaunchConfig1D::auto_shared(groups, $block + 1) + ); + }; + ($t:ty,$prepare:ident,$launch:ident,$block:literal,$capacity:literal,$items:literal,$default:expr, + $config:expr,$wrong_shape:expr) => {{ + let default: $t = $default; + for groups in [1_u32, 3] { + let plan = module.$prepare::<$items>(($config)(groups))?; + assert!(plan.__raw_config().shared_mem_bytes > 0); + let error = module + .$prepare::<$items>(($wrong_shape)(groups)) + .err() + .expect("a different exact block shape must reject"); + assert!(error.to_string().contains("does not match required block")); + let (gx, gy, gz) = plan.__raw_config().grid_dim; + let groups = gx * gy * gz; + for mode in [0, 1, 2, 3, 4, 6] { + for valid in [0, 1, $block * $items - 1, $block * $items] { + if mode == 0 && valid != $block * $items { + continue; + } + let size = groups as usize * $block * $capacity; + let input: Vec<$t> = (0..size).map(|i| (i % 61 + 1) as $t).collect(); + let source = DeviceBuffer::from_host(&stream, &input)?; + let target = DeviceBuffer::from_host(&stream, &vec![0 as $t; size])?; + let (_, target) = module + .$launch::<$items, _, _>(&plan, source, target, valid, default, mode) + .sync()?; + let mut result = vec![0 as $t; size]; + target.copy_to_host(&stream, &mut result)?; + stream.synchronize()?; + for (index, &value) in result.iter().enumerate() { + let within = index % ($block * $capacity); + let written = if mode == 0 || mode == 2 { + $block * $items + } else if mode == 6 { + valid / 2 + } else { + valid + }; + let expected = if within < written { + if mode == 2 && within >= valid { + default + } else { + input[index] + } + } else { + 0 as $t + }; + assert_eq!( + value.to_ne_bytes(), + expected.to_ne_bytes(), + "Transpose type={} B={} I={} mode={} valid={} index={}", + stringify!($t), + $block, + $items, + mode, + valid, + index + ); + } + cases += 1; + } + } + } + }}; + } + run!( + i8, + prepare_copy_i8_33_4, + copy_i8_33_4_async_owned, + 33, + 4, + 1, + i8::MIN + ); + run!( + i8, + prepare_copy_i8_33_4, + copy_i8_33_4_async_owned, + 33, + 4, + 4, + i8::MIN + ); + run!( + u8, + prepare_copy_u8_33_4, + copy_u8_33_4_async_owned, + 33, + 4, + 1, + u8::MAX + ); + run!( + u8, + prepare_copy_u8_33_4, + copy_u8_33_4_async_owned, + 33, + 4, + 4, + u8::MAX + ); + run!( + i16, + prepare_copy_i16_33_4, + copy_i16_33_4_async_owned, + 33, + 4, + 1, + i16::MIN + ); + run!( + i16, + prepare_copy_i16_33_4, + copy_i16_33_4_async_owned, + 33, + 4, + 4, + i16::MIN + ); + run!( + u16, + prepare_copy_u16_33_4, + copy_u16_33_4_async_owned, + 33, + 4, + 1, + u16::MAX + ); + run!( + u16, + prepare_copy_u16_33_4, + copy_u16_33_4_async_owned, + 33, + 4, + 4, + u16::MAX + ); + run!( + i32, + prepare_copy_i32_33_4, + copy_i32_33_4_async_owned, + 33, + 4, + 1, + i32::MIN + ); + run!( + i32, + prepare_copy_i32_33_4, + copy_i32_33_4_async_owned, + 33, + 4, + 4, + i32::MIN + ); + run!( + u32, + prepare_copy_u32_33_4, + copy_u32_33_4_async_owned, + 33, + 4, + 1, + u32::MAX + ); + run!( + u32, + prepare_copy_u32_33_4, + copy_u32_33_4_async_owned, + 33, + 4, + 4, + u32::MAX + ); + run!( + i64, + prepare_copy_i64_33_4, + copy_i64_33_4_async_owned, + 33, + 4, + 1, + i64::MIN + 17 + ); + run!( + i64, + prepare_copy_i64_33_4, + copy_i64_33_4_async_owned, + 33, + 4, + 4, + i64::MIN + 17 + ); + run!( + u64, + prepare_copy_u64_33_4, + copy_u64_33_4_async_owned, + 33, + 4, + 1, + u64::MAX - 17 + ); + run!( + u64, + prepare_copy_u64_33_4, + copy_u64_33_4_async_owned, + 33, + 4, + 4, + u64::MAX - 17 + ); + run!( + f32, + prepare_copy_f32_33_4, + copy_f32_33_4_async_owned, + 33, + 4, + 1, + -0.0_f32 + ); + run!( + f32, + prepare_copy_f32_33_4, + copy_f32_33_4_async_owned, + 33, + 4, + 4, + -0.0_f32 + ); + run!( + f64, + prepare_copy_f64_33_4, + copy_f64_33_4_async_owned, + 33, + 4, + 1, + -0.0_f64 + ); + run!( + f64, + prepare_copy_f64_33_4, + copy_f64_33_4_async_owned, + 33, + 4, + 4, + -0.0_f64 + ); + run!( + i32, + prepare_copy_i32_1_4, + copy_i32_1_4_async_owned, + 1, + 4, + 1, + i32::MIN + ); + run!( + i32, + prepare_copy_i32_1_4, + copy_i32_1_4_async_owned, + 1, + 4, + 4, + i32::MIN + ); + run!( + i32, + prepare_copy_i32_128_4, + copy_i32_128_4_async_owned, + 128, + 4, + 1, + i32::MIN + ); + run!( + i32, + prepare_copy_i32_128_4, + copy_i32_128_4_async_owned, + 128, + 4, + 4, + i32::MIN + ); + run!( + i32, + prepare_copy_i32_130_4, + copy_i32_130_4_async_owned, + 130, + 4, + 1, + i32::MIN + ); + run!( + i32, + prepare_copy_i32_130_4, + copy_i32_130_4_async_owned, + 130, + 4, + 4, + i32::MIN + ); + run!( + i32, + prepare_copy_i32_33_9, + copy_i32_33_9_async_owned, + 33, + 9, + 9, + i32::MIN + ); + run!( + i32, + prepare_copy_i32_2d, + copy_i32_2d_async_owned, + 33, + 4, + 1, + -7, + |groups| LaunchConfig2D::auto_shared((groups, 2), (11, 3)), + |groups| LaunchConfig2D::auto_shared((groups, 2), (3, 11)) + ); + run!( + i32, + prepare_copy_i32_3d, + copy_i32_3d_async_owned, + 128, + 4, + 1, + -7, + |groups| LaunchConfig3D::auto_shared((groups, 2, 2), (8, 4, 4)), + |groups| LaunchConfig3D::auto_shared((groups, 2, 2), (16, 2, 4)) + ); + run!( + i32, + prepare_copy_i32_2d, + copy_i32_2d_async_owned, + 33, + 4, + 4, + -7, + |groups| LaunchConfig2D::auto_shared((groups, 2), (11, 3)), + |groups| LaunchConfig2D::auto_shared((groups, 2), (3, 11)) + ); + run!( + i32, + prepare_copy_i32_3d, + copy_i32_3d_async_owned, + 128, + 4, + 4, + -7, + |groups| LaunchConfig3D::auto_shared((groups, 2, 2), (8, 4, 4)), + |groups| LaunchConfig3D::auto_shared((groups, 2, 2), (16, 2, 4)) + ); + let plan = module.prepare_transpose_scalar(LaunchConfig1D::auto_shared(1, 33))?; + let target = DeviceBuffer::from_host(&stream, &[0_i16; 33])?; + let target = module + .transpose_scalar_async_owned::<_>(&plan, target) + .sync()?; + let mut result = vec![0_i16; 33]; + target.copy_to_host(&stream, &mut result)?; + stream.synchronize()?; + assert!(result.iter().all(|&value| value == -17)); + cases += 1; + println!( + "PASS: {cases} public Transpose cases, all ten types, arbitrary masks, repeated Store, partial warps and N9" + ); + Ok(()) +} diff --git a/rust/cuda-coop/tests/load_store/src/vectorize.rs b/rust/cuda-coop/tests/load_store/src/vectorize.rs new file mode 100644 index 000000000000..c4708797870d --- /dev/null +++ b/rust/cuda-coop/tests/load_store/src/vectorize.rs @@ -0,0 +1,380 @@ +// Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +//! Vectorize full/guarded access and post-Store validity, including misaligned offsets. +use cuda_coop as coop; +use cuda_device::{CollectiveTiles, DisjointSlice, kernel, launch_contract, thread}; +use cuda_host::cuda_module; + +#[cuda_module] +mod kernels { + use super::*; + + fn exercise( + input: &[T], + output: &mut DisjointSlice>, + valid: usize, + shift: usize, + default: T, + mode: u32, + ) { + let block = coop::this_block(); + let x = thread::blockDim_x() as usize; + let y = thread::blockDim_y() as usize; + let threads = x * y * thread::blockDim_z() as usize; + let lane = thread::threadIdx_x() as usize + + x * (thread::threadIdx_y() as usize + y * thread::threadIdx_z() as usize); + let group = thread::blockIdx_x() as usize + + thread::gridDim_x() as usize + * (thread::blockIdx_y() as usize + + thread::gridDim_y() as usize * thread::blockIdx_z() as usize); + let offset = group * threads * 8 + shift; + let mut values = coop::ThreadData::<_, ITEMS>::new(); + match mode { + 0 => coop::load( + block, + input, + &mut values, + coop::LoadOptions::default() + .algorithm(coop::Vectorize) + .offset(offset), + ), + 1 | 5 | 6 | 9 => { + // A second masked Load must clear validity left by a full Load. + coop::load( + block, + input, + &mut values, + coop::LoadOptions::default() + .algorithm(coop::Vectorize) + .offset(offset), + ); + coop::load( + block, + input, + &mut values, + coop::LoadOptions::default() + .algorithm(coop::Vectorize) + .offset(offset) + .valid_items(valid), + ); + } + 2 | 4 => coop::load( + block, + input, + &mut values, + coop::LoadOptions::default() + .algorithm(coop::Vectorize) + .offset(offset) + .valid_items(valid) + .oob_default(default), + ), + 7 => {} + _ => { + let mut item = 0; + while item < ITEMS { + let selected = lane * ITEMS + item < valid; + if (selected || (lane + item) & 1 == 0) + && !(mode == 8 && lane * ITEMS + item == valid - 1) + { + values.write(item, default); + } + item += 1; + } + } + } + let before = values; + if mode == 0 || mode == 2 || mode == 9 { + coop::store( + block, + &mut *output, + &values, + coop::StoreOptions::default() + .algorithm(coop::Vectorize) + .offset(offset), + ); + } else { + let written = if mode == 5 { + valid / 2 + } else if mode == 6 { + valid + 1 + } else { + valid + }; + coop::store( + block, + &mut *output, + &values, + coop::StoreOptions::default() + .algorithm(coop::Vectorize) + .offset(offset) + .valid_items(written), + ); + } + let mut observations_ok = true; + let mut item = 0; + while item < ITEMS { + let position = lane * ITEMS + item; + let expected = match mode { + 3 | 8 => { + if position < valid || (lane + item) & 1 == 0 { + Some(default) + } else { + None + } + } + 7 => None, + _ => { + if position < valid || mode == 0 { + Some(input[offset + position]) + } else if mode == 2 || mode == 4 { + Some(default) + } else { + None + } + } + }; + observations_ok &= before.get(item) == expected; + observations_ok &= values.get(item) == before.get(item); + item += 1; + } + if mode == 3 || mode == 4 { + thread::sync_threads(); + // A repeated zero Store preserves arbitrary initialized/invalid slots. + coop::store( + block, + &mut *output, + &values, + coop::StoreOptions::default() + .algorithm(coop::Vectorize) + .offset(offset) + .valid_items(0), + ); + let mut item = 0; + while item < ITEMS { + observations_ok &= values.get(item) == before.get(item); + item += 1; + } + } + assert!(observations_ok); + } + + #[kernel] + #[launch_contract(domain = 1, block = (33, 1, 1), dynamic_shared = inferred)] + pub fn partial_33( + input: &[T], + mut output: DisjointSlice>, + valid: usize, + shift: usize, + default: T, + mode: u32, + ) { + exercise::<_, ITEMS>(input, &mut output, valid, shift, default, mode); + } + + #[kernel] + #[launch_contract(domain = 3, block = (8, 4, 4), dynamic_shared = inferred)] + pub fn partial_3d( + input: &[T], + mut output: DisjointSlice>, + valid: usize, + shift: usize, + default: T, + mode: u32, + ) { + exercise::<_, ITEMS>(input, &mut output, valid, shift, default, mode); + } + + #[kernel] + #[launch_contract(domain = 1, block = (33, 1, 1), dynamic_shared = inferred)] + pub fn vectorize_scalar(mut output: DisjointSlice>, valid: usize) { + let offset = thread::blockIdx_x() as usize * 33; + coop::store( + coop::this_block(), + &mut output, + 17_i32, + coop::StoreOptions::default() + .algorithm(coop::Vectorize) + .offset(offset) + .valid_items(valid), + ); + } +} + +pub fn run(trap: Option<&str>) -> Result<(), Box> { + use cuda_async::simt::device_context::{init_device_contexts, with_cuda_context}; + use cuda_async::simt::device_operation::DeviceOperation; + use cuda_core::{DeviceBuffer, LaunchConfig1D, LaunchConfig3D}; + init_device_contexts(0, 1)?; + let context = with_cuda_context(0, |context| context.clone())?; + let stream = context.default_stream(); + // SAFETY: embedded kernels and typed declarations compile together. + let module = unsafe { kernels::load_async(0)? }; + if let Some(trap) = trap { + let mode = match trap { + "vectorize-larger" => 6, + "vectorize-full-hole" => 9, + _ => 8, + }; + let plan = module.prepare_partial_33::(LaunchConfig1D::auto_shared(1, 33))?; + let source = DeviceBuffer::from_host(&stream, &vec![7_i32; 132])?; + let target = DeviceBuffer::from_host(&stream, &vec![0_i32; 132])?; + module + .partial_33_async_owned::(&plan, source, target, 34, 0, -7, mode) + .sync()?; + panic!("invalid Vectorize Store unexpectedly completed"); + } + let mut cases = 0; + macro_rules! run { + ($t:ty, $items:literal, $threads:literal, $prepare:ident, $launch:ident, $config:expr, $extra:expr) => {{ + let tile = $threads * $items; + let plan = module.$prepare::<$t, $items>($config)?; + assert_eq!(plan.__raw_config().shared_mem_bytes, 0); + let prefix = ($threads + 1).min(tile); + let mut selected = vec![ + (0, tile, 0), + (0, tile, 1), + (0, tile, 4), + (1, prefix, 0), + (2, 1, 0), + (3, prefix, 0), + (5, prefix, 0), + ]; + if $extra { + for valid in [0, 1, $threads - 1, $threads, prefix, tile] { + for mode in [1, 2, 3, 4, 5] { + selected.push((mode, valid, 0)); + } + } + for mode in [1, 3, 4, 5] { + selected.push((mode, tile - 7, 7)); + } + selected.push((7, 0, 0)); + } + selected.sort_unstable(); + selected.dedup(); + for (mode, valid, shift) in selected { + let written = if mode == 0 || mode == 2 || mode == 9 { + tile + } else if mode == 5 { + valid / 2 + } else { + valid + }; + // Keep the final allocation shorter than the full tile when + // counted Store writes a prefix. Its guarded iterator must not + // form an out-of-allocation pointer. + let output_len = $threads * 8 + shift + written; + let input: Vec<$t> = (0..2 * $threads * 8 + shift) + .map(|i| (i % 101) as $t) + .collect(); + let default = 111 as $t; + let sentinel = 123 as $t; + let source = DeviceBuffer::from_host(&stream, &input)?; + let target = DeviceBuffer::from_host(&stream, &vec![sentinel; output_len])?; + let (_, target) = module + .$launch::<$t, $items, _, _>(&plan, source, target, valid, shift, default, mode) + .sync()?; + let mut output = vec![sentinel; output_len]; + target.copy_to_host(&stream, &mut output)?; + stream.synchronize()?; + for (index, actual) in output.into_iter().enumerate() { + let local = index % ($threads * 8); + let expected = if local >= shift && local - shift < written { + if mode == 3 || (mode == 2 && local - shift >= valid) { + default + } else { + input[index] + } + } else { + sentinel + }; + assert_eq!( + actual, + expected, + "type={} items={} mode={mode} valid={valid} shift={shift} index={index}", + stringify!($t), + $items + ); + } + cases += 1; + } + }}; + } + macro_rules! width { + ($t:ty) => { + run!( + $t, + 1, + 33, + prepare_partial_33, + partial_33_async_owned, + LaunchConfig1D::auto_shared(2, 33), + false + ); + run!( + $t, + 4, + 33, + prepare_partial_33, + partial_33_async_owned, + LaunchConfig1D::auto_shared(2, 33), + false + ); + }; + } + width!(i8); + width!(u8); + width!(i16); + width!(u16); + width!(i32); + width!(u32); + width!(i64); + width!(u64); + width!(f32); + width!(f64); + run!( + i32, + 4, + 33, + prepare_partial_33, + partial_33_async_owned, + LaunchConfig1D::auto_shared(2, 33), + true + ); + run!( + i32, + 4, + 128, + prepare_partial_3d, + partial_3d_async_owned, + LaunchConfig3D::auto_shared((1, 2, 1), (8, 4, 4)), + true + ); + run!( + i32, + 3, + 33, + prepare_partial_33, + partial_33_async_owned, + LaunchConfig1D::auto_shared(2, 33), + false + ); + for valid in [0, 1, 32, 33] { + let plan = module.prepare_vectorize_scalar(LaunchConfig1D::auto_shared(2, 33))?; + assert_eq!(plan.__raw_config().shared_mem_bytes, 0); + let target = DeviceBuffer::from_host(&stream, &vec![-1_i32; 33 + valid])?; + let target = module + .vectorize_scalar_async_owned(&plan, target, valid) + .sync()?; + let mut output = vec![0_i32; 33 + valid]; + target.copy_to_host(&stream, &mut output)?; + stream.synchronize()?; + for (i, value) in output.into_iter().enumerate() { + assert_eq!(value, if i % 33 < valid { 17 } else { -1 }); + } + cases += 1; + } + println!("PASS: {cases} Vectorize cases with observed payload mapping and validity"); + Ok(()) +} diff --git a/rust/cuda-coop/tests/load_store/src/warp.rs b/rust/cuda-coop/tests/load_store/src/warp.rs new file mode 100644 index 000000000000..eec19289f87b --- /dev/null +++ b/rust/cuda-coop/tests/load_store/src/warp.rs @@ -0,0 +1,334 @@ +// Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +//! Physical Warp Direct uses the ordinary public API and owned async launch. +use cuda_coop as coop; +use cuda_device::{kernel, launch_contract, library, thread}; +use cuda_host::cuda_module; + +#[cuda_module] +mod kernels { + use super::*; + + #[inline(never)] + fn exercise( + input: &[i32], + output: &mut coop::DisjointSlice>, + valid: usize, + mode: u32, + ) { + let warp = coop::this_warp(); + let q = library::warp_group_index_v1::<32>() as usize; + let block = thread::blockIdx_x() as usize; + let offset = block * 64 * I; + let lane = (thread::threadIdx_x() as usize + + thread::blockDim_x() as usize + * (thread::threadIdx_y() as usize + + thread::blockDim_y() as usize * thread::threadIdx_z() as usize)) + % 32; + let count = if mode == 0 { + 32 * I + } else if mode == 3 { + (valid + q).min(32 * I) + } else { + valid + }; + let mut values = coop::ThreadData::<_, I>::new(); + if mode == 4 { + let mut other = coop::ThreadData::<_, I>::new(); + let mut i = 0; + while i < I { + if lane * I + i < count || i % 2 == 0 { + values.write(i, -7); + } + if lane * I + i < count || i % 2 != 0 { + other.write(i, -7); + } + i += 1; + } + values = if lane & 1 == 0 { values } else { other }; + } else if mode == 0 { + coop::load( + warp, + input, + &mut values, + coop::LoadOptions::default().offset(offset), + ); + } else if mode == 2 { + coop::load( + warp, + input, + &mut values, + coop::LoadOptions::default() + .offset(offset) + .valid_items(count) + .oob_default(-7), + ); + } else { + // A preceding full Load makes the exact masked reset observable. + coop::load( + warp, + input, + &mut values, + coop::LoadOptions::default().offset(offset), + ); + coop::load( + warp, + input, + &mut values, + coop::LoadOptions::default() + .offset(offset) + .valid_items(count), + ); + } + let before = values; + coop::store( + warp, + &mut *output, + &values, + coop::StoreOptions::default() + .offset(offset) + .valid_items(count), + ); + let after = values; + thread::sync_threads(); + coop::load( + warp, + output, + &mut values, + coop::LoadOptions::default() + .offset(offset) + .valid_items(count), + ); + let mut matches = true; + let mut i = 0; + while i < I { + matches &= before.get(i) == after.get(i); + if lane * I + i < count { + let expected = if mode == 4 { + -7 + } else { + input[offset + q * 32 * I + lane * I + i] + }; + matches &= values.get(i) == Some(expected); + } else { + matches &= values.get(i).is_none(); + if mode == 1 || mode == 3 { + matches &= before.get(i).is_none(); + } + if mode == 2 { + matches &= before.get(i) == Some(-7); + } + } + i += 1; + } + assert!(matches); + } + + #[kernel] + #[launch_contract(domain=1,block=(64,1,1),dynamic_shared=inferred)] + pub fn warp_direct_64( + input: &[i32], + mut output: coop::DisjointSlice>, + valid: usize, + mode: u32, + ) { + exercise::(input, &mut output, valid, mode); + } + #[kernel] + #[launch_contract(domain=3,block=(8,4,2),dynamic_shared=inferred)] + pub fn warp_direct_xyz( + input: &[i32], + mut output: coop::DisjointSlice>, + valid: usize, + mode: u32, + ) { + exercise::(input, &mut output, valid, mode); + } + + #[kernel] + #[launch_contract(domain=1,block=(64,1,1),dynamic_shared=inferred)] + pub fn warp_mixed_block_transpose( + input: &[i32], + mut output: coop::DisjointSlice>, + ) { + let offset = thread::blockIdx_x() as usize * 256; + let mut values = coop::ThreadData::<_, 4>::new(); + coop::load( + coop::this_block(), + input, + &mut values, + coop::LoadOptions::default() + .offset(offset) + .algorithm(coop::Transpose), + ); + coop::store( + coop::this_warp(), + &mut output, + &values, + coop::StoreOptions::default().offset(offset), + ); + thread::sync_threads(); + coop::load( + coop::this_warp(), + &mut output, + &mut values, + coop::LoadOptions::default().offset(offset).valid_items(127), + ); + let lane = thread::threadIdx_x() as usize; + let mut matches = true; + let mut i = 0; + while i < 4 { + let expected = if (lane % 32) * 4 + i < 127 { + Some(input[offset + lane * 4 + i]) + } else { + None + }; + matches &= values.get(i) == expected; + i += 1; + } + assert!(matches); + } + + #[kernel] + #[launch_contract(domain=1,block=(64,1,1),dynamic_shared=inferred)] + pub fn warp_fixed_grain( + mut output: coop::DisjointSlice>, + adapt: bool, + ) { + let q = library::warp_group_index_v1::<32>() as usize; + let offset = thread::blockIdx_x() as usize * 256 + if adapt { q * 96 } else { 0 }; + // Automatic origin is q*32 for this one-item scalar. Carrier capacity + // stays four, so warp one needs caller offset96 to reach its region128. + coop::store( + coop::this_warp(), + &mut output, + 7_i32, + coop::StoreOptions::default().offset(offset), + ); + } +} + +pub fn run() -> Result<(), Box> { + use cuda_async::simt::device_context::{init_device_contexts, with_cuda_context}; + use cuda_async::simt::device_operation::DeviceOperation; + use cuda_core::{DeviceBuffer, LaunchConfig1D, LaunchConfig3D}; + init_device_contexts(0, 1)?; + let stream = with_cuda_context(0, |context| context.default_stream())?; + // SAFETY: the typed launcher and embedded native kernels compile together. + let module = unsafe { kernels::load_async(0)? }; + if std::env::args().any(|argument| argument == "--warp-grain-trap") { + let plan = module.prepare_warp_fixed_grain(LaunchConfig1D::auto_shared(1, 64))?; + let target = DeviceBuffer::from_host(&stream, &vec![-99_i32; 256])?; + module + .warp_fixed_grain_async_owned(&plan, target, false) + .sync()?; + panic!("expected ownership trap returned successfully"); + } + for groups in [1_u32, 3] { + let plan = module.prepare_warp_fixed_grain(LaunchConfig1D::auto_shared(groups, 64))?; + let target = DeviceBuffer::from_host(&stream, &vec![-99_i32; groups as usize * 256])?; + let target = module + .warp_fixed_grain_async_owned(&plan, target, true) + .sync()?; + let mut result = vec![0; groups as usize * 256]; + target.copy_to_host(&stream, &mut result)?; + stream.synchronize()?; + let mut expected = vec![-99_i32; result.len()]; + for block in 0..groups as usize { + for warp in 0..2 { + expected[block * 256 + warp * 128..block * 256 + warp * 128 + 32].fill(7); + } + } + assert_eq!(result, expected); + } + // A storage-bearing Block operation and scratch-free Warp operations + // share one root plan and the same collective allocation. + let plan = module.prepare_warp_mixed_block_transpose(LaunchConfig1D::auto_shared(2, 64))?; + assert!(plan.__raw_config().shared_mem_bytes > 0); + let input: Vec = (0..512).map(|i| i + 17).collect(); + let source = DeviceBuffer::from_host(&stream, &input)?; + let target = DeviceBuffer::from_host(&stream, &vec![-99_i32; input.len()])?; + let (_, target) = module + .warp_mixed_block_transpose_async_owned(&plan, source, target) + .sync()?; + let mut result = vec![0; input.len()]; + target.copy_to_host(&stream, &mut result)?; + stream.synchronize()?; + assert_eq!(result, input); + let mut cases = 0; + macro_rules! run { + ($i:literal,$prepare:ident,$launch:ident,$config:expr,$wrong:expr) => {{ + assert!(module.$prepare::<$i>($wrong).is_err()); + for groups in [1_u32, 3] { + let input: Vec = (0..groups as usize * 64 * $i) + .map(|i| i as i32 + 11) + .collect(); + let plan = module.$prepare::<$i>($config(groups))?; + assert_eq!(plan.__raw_config().shared_mem_bytes, 0); + for mode in 0..5 { + for valid in [0, 1, 32 * $i - 1, 32 * $i] { + let source = DeviceBuffer::from_host(&stream, &input)?; + let target = DeviceBuffer::from_host(&stream, &vec![-99_i32; input.len()])?; + let (_, target) = module + .$launch::<$i, _, _>(&plan, source, target, valid, mode) + .sync()?; + let mut result = vec![0; input.len()]; + target.copy_to_host(&stream, &mut result)?; + stream.synchronize()?; + let mut expected = vec![-99; input.len()]; + for block in 0..groups as usize { + for q in 0..2 { + let count = if mode == 0 { + 32 * $i + } else if mode == 3 { + (valid + q).min(32 * $i) + } else { + valid + }; + for i in 0..count { + let index = block * 64 * $i + q * 32 * $i + i; + expected[index] = if mode == 4 { -7 } else { input[index] }; + } + } + } + assert_eq!(result, expected); + cases += 1; + } + } + } + }}; + } + run!( + 1, + prepare_warp_direct_64, + warp_direct_64_async_owned, + |g| LaunchConfig1D::auto_shared(g, 64), + LaunchConfig1D::auto_shared(1, 32) + ); + run!( + 4, + prepare_warp_direct_64, + warp_direct_64_async_owned, + |g| LaunchConfig1D::auto_shared(g, 64), + LaunchConfig1D::auto_shared(1, 32) + ); + run!( + 1, + prepare_warp_direct_xyz, + warp_direct_xyz_async_owned, + |g| LaunchConfig3D::auto_shared((g, 1, 1), (8, 4, 2)), + LaunchConfig3D::auto_shared((1, 1, 1), (4, 4, 4)) + ); + run!( + 4, + prepare_warp_direct_xyz, + warp_direct_xyz_async_owned, + |g| LaunchConfig3D::auto_shared((g, 1, 1), (8, 4, 2)), + LaunchConfig3D::auto_shared((1, 1, 1), (4, 4, 4)) + ); + println!( + "PASS {cases} physical Warp Direct cases and 2 fixed-grain cases; mixed Block Transpose/Warp Direct passed" + ); + Ok(()) +} diff --git a/rust/cuda-coop/tests/payload.rs b/rust/cuda-coop/tests/payload.rs new file mode 100644 index 000000000000..88e8c24b4a9e --- /dev/null +++ b/rust/cuda-coop/tests/payload.rs @@ -0,0 +1,178 @@ +// Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +//! Public payload behavior, numeric inference, validity and element alignment. + +use core::mem::{align_of, size_of}; +use std::panic::{AssertUnwindSafe, catch_unwind}; + +use cuda_coop as coop; + +#[test] +fn typed_writes_infer_every_numeric_width_at_one_and_four_items() { + macro_rules! check { + ($($ty:ty),* $(,)?) => { + $( + for value in [<$ty>::MIN, <$ty>::MAX] { + let mut one = coop::ThreadData::<_, 1>::new(); + one.write(0, value); + assert_eq!(one.get(0), Some(value)); + assert_eq!(size_of_val(&one[0]), size_of::<$ty>()); + + let mut four = coop::ThreadData::<_, 4>::new(); + four.write(3, value); + assert_eq!(four.get(3), Some(value)); + assert_eq!(four.get(0), None); + assert_eq!(four.get(1), None); + assert_eq!(four.get(2), None); + } + )* + }; + } + check!(i8, u8, i16, u16, i32, u32, i64, u64, f32, f64); + + let mut single = coop::ThreadData::<_, 1>::new(); + single.write(0, f32::from_bits(0x7fc0_1234)); + assert_eq!(single[0].to_bits(), 0x7fc0_1234); + let mut double = coop::ThreadData::<_, 1>::new(); + double.write(0, f64::from_bits(0xfff8_0000_0000_1234)); + assert_eq!(double[0].to_bits(), 0xfff8_0000_0000_1234); +} + +#[test] +fn arbitrary_writes_and_checked_mutation_preserve_unwritten_slots() { + let mut items = coop::ThreadData::<_, 4>::new(); + items.write(2, 8_i32); + items.write(0, 3_i32); + items[2] *= 2; + *items.get_mut(0).unwrap() += 1; + assert_eq!( + [items.get(0), items.get(1), items.get(2), items.get(3)], + [Some(4), None, Some(16), None] + ); + assert_eq!(items.len(), 4); + assert!(!items.is_empty()); + assert!(items.get_mut(1).is_none()); + for index in [4, usize::MAX] { + assert_eq!(items.get(index), None); + assert!(items.get_mut(index).is_none()); + } + items.write(2, -1); + assert_eq!(items.get(2), Some(-1)); +} + +#[test] +fn sparse_writes_across_byte_boundaries_preserve_neighbors_and_copies() { + let mut items = coop::ThreadData::<_, 33>::new(); + for index in [0, 7, 8, 16, 31, 32] { + items.write(index, index as u16 + 1); + } + let copy = items; + *items.get_mut(8).unwrap() += 10; + items[31] += 20; + for index in 0..33 { + let expected = [0, 7, 8, 16, 31, 32] + .contains(&index) + .then_some(index as u16 + 1); + assert_eq!(copy.get(index), expected); + assert_eq!( + items.get(index), + expected.map(|value| value + + match index { + 8 => 10, + 31 => 20, + _ => 0, + }) + ); + } + for index in [33, usize::MAX] { + assert_eq!(items.get(index), None); + assert!(items.get_mut(index).is_none()); + } +} + +#[test] +fn invalid_indexing_and_writes_panic_without_initializing_other_slots() { + let mut items = coop::ThreadData::<_, 4>::new(); + items.write(2, 9_i32); + for index in [0, 4, usize::MAX] { + assert!(catch_unwind(|| items[index]).is_err()); + assert!(catch_unwind(AssertUnwindSafe(|| items[index] = 17)).is_err()); + } + for index in [4, usize::MAX] { + assert!(catch_unwind(AssertUnwindSafe(|| items.write(index, 17))).is_err()); + } + assert_eq!( + [items.get(0), items.get(1), items.get(2), items.get(3)], + [None, None, Some(9), None] + ); +} + +#[test] +fn copies_moves_swaps_and_selection_keep_exact_masks() { + let mut first = coop::ThreadData::<_, 4>::new(); + first.write(0, 1_u16); + first.write(3, 4_u16); + let original = first; + let mut second = coop::ThreadData::<_, 4>::new(); + second.write(1, 2_u16); + core::mem::swap(&mut first, &mut second); + assert_eq!( + (first.get(0), first.get(1), first.get(3)), + (None, Some(2), None) + ); + assert_eq!( + (second.get(0), second.get(1), second.get(3)), + (Some(1), None, Some(4)) + ); + for choose_first in [true, false] { + let mut selected = if choose_first { first } else { second }; + assert_eq!(selected.get(1), if choose_first { Some(2) } else { None }); + selected.write(2, 99); + assert_eq!(first.get(2), None); + assert_eq!(second.get(2), None); + } + assert_eq!( + (original.get(0), original.get(1), original.get(3)), + (Some(1), None, Some(4)) + ); +} + +#[test] +fn sealed_generic_interface_preserves_inference_and_extent() { + fn fill_last(items: &mut D, value: D::Item) { + assert_eq!(items.len(), D::EXTENT); + assert!(!items.is_empty()); + assert!(items.get(D::EXTENT - 1).is_none()); + items.write(D::EXTENT - 1, value); + } + let mut items = coop::ThreadData::<_, 4, 16>::new(); + fill_last(&mut items, 255_u8); + assert_eq!(items.get(3), Some(255_u8)); + assert_eq!(items.get(0), None); +} + +#[test] +fn alignment_applies_to_the_values_array() { + let mut items = coop::ThreadData::<_, 4, 64>::new(); + items.write(0, 13_u16); + items.write(1, 27_u16); + let values = &items[0] as *const u16 as usize; + assert_eq!(values, &items as *const _ as usize); + assert_eq!(values % 64, 0); + assert_eq!(&items[1] as *const u16 as usize - values, size_of::()); + assert_eq!(align_of::>(), align_of::()); + + // Query, rather than instantiate, the largest layouts: an over-aligned + // object can have a large padded size even with one element. + macro_rules! check { + ($($align:literal),* $(,)?) => { + $(assert_eq!(align_of::>(), $align);)* + }; + } + check!( + 1, 2, 4, 8, 16, 32, 64, 128, 256, 512, 1024, 2048, 4096, 8192, 16384, 32768, 65536, 131072, + 262144, 524288, 1048576, 2097152, 4194304, 8388608, 16777216, 33554432, 67108864, + 134217728, 268435456, 536870912 + ); +} diff --git a/rust/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..ee1b29d8e5f5 --- /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=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "libnvvm-sys", + "nvjitlink-sys", + "serde", + "sha2", + "thiserror 2.0.21", +] + +[[package]] +name = "cuda-async" +version = "0.5.0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "anyhow", + "cuda-bindings", + "cuda-core", + "dashmap", + "futures", + "half", + "once_cell", + "rustc-hash 2.1.3", + "thiserror 1.0.69", +] + +[[package]] +name = "cuda-bindings" +version = "0.5.0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "bindgen", + "libloading", + "prettyplease", + "proc-macro2", + "quote", + "syn 2.0.119", +] + +[[package]] +name = "cuda-coop" +version = "0.0.0" +dependencies = [ + "cuda-coop-codegen", + "cuda-device", +] + +[[package]] +name = "cuda-coop-codegen" +version = "0.0.0" +dependencies = [ + "libloading", + "serde", + "serde_json", + "sha2", +] + +[[package]] +name = "cuda-core" +version = "0.5.0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "anyhow", + "cuda-bindings", + "cuda-core-derive", + "half", + "oxide-artifacts", +] + +[[package]] +name = "cuda-core-derive" +version = "0.5.0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", +] + +[[package]] +name = "cuda-device" +version = "0.2.1" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "cuda-macros", +] + +[[package]] +name = "cuda-host" +version = "0.2.1" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "cuda-artifact-finalizer", + "cuda-async", + "cuda-core", + "cuda-device", + "cuda-macros", + "half", + "libc", + "oxide-artifacts", + "ptx-parse", + "sha2", + "thiserror 2.0.21", +] + +[[package]] +name = "cuda-macros" +version = "0.2.1" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "oxide-tile-export", + "proc-macro2", + "quote", + "reserved-oxide-symbols", + "serde", + "serde_json", + "syn 2.0.119", +] + +[[package]] +name = "cuda-target-spec" +version = "0.1.0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" + +[[package]] +name = "dashmap" +version = "6.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e6361d5c062261c78a176addb82d4c821ae42bed6089de0e12603cd25de2059c" +dependencies = [ + "cfg-if", + "crossbeam-utils", + "hashbrown 0.14.5", + "lock_api", + "once_cell", + "parking_lot_core", +] + +[[package]] +name = "digest" +version = "0.10.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9ed9a281f7bc9b7576e61468ba615a66a5c8cfdff42420a70aa82701a3b1e292" +dependencies = [ + "block-buffer", + "crypto-common", +] + +[[package]] +name = "either" +version = "1.19.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0e9c71c2167ca323c882b99918929403426e2373ea17242ff5653e0d5e1058be" + +[[package]] +name = "equivalent" +version = "1.0.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "00d174d5400e5e8fd687ad1049e2f578285fa914201b1af7e8b112a4546bd826" + +[[package]] +name = "errno" +version = "0.3.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "39cab71617ae0d63f51a36d69f866391735b51691dbda63cf6f96d042b63efeb" +dependencies = [ + "libc", + "windows-sys 0.61.2", +] + +[[package]] +name = "foldhash" +version = "0.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d9c4f5dac5e15c24eb999c26181a6ca40b39fe946cbe4c263c7209467bc83af2" + +[[package]] +name = "futures" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9a31d2a3fbaaeb2af2368bbdd904aa8e812d3c04a1ee10d3171f52d556e5d0a3" +dependencies = [ + "futures-channel", + "futures-core", + "futures-executor", + "futures-io", + "futures-sink", + "futures-task", + "futures-util", +] + +[[package]] +name = "futures-channel" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b1f9e3d69d39e4862ffed03ed071a76f9a13ba1d9109d355b0f0aa6b15e393c4" +dependencies = [ + "futures-core", + "futures-sink", +] + +[[package]] +name = "futures-core" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "92d699e522242e69e3003b94ecc1f960f3a5e015aa7c5d7486e65ad01dd94f5e" + +[[package]] +name = "futures-executor" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "031b47cf1a3c6cc8bc2fc76cd437f521619387907d469316e7c0bc278f1f5432" +dependencies = [ + "futures-core", + "futures-task", + "futures-util", +] + +[[package]] +name = "futures-io" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "53c0fa8157de1303bfffdaa1cc2a673bfffb60102f76b0ef4441659124373fed" + +[[package]] +name = "futures-macro" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9fb9654ba8355388abeb8dcb4fc62f511300867002afc858860463bdd9fe0c44" +dependencies = [ + "proc-macro2", + "quote", + "syn 3.0.9", +] + +[[package]] +name = "futures-sink" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1944426bf7d03f1d14f708785e4b33efd750b36d48a157b836b3efc15ede8e1d" + +[[package]] +name = "futures-task" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cd417de3d1d015fc3bfd2b1ea46dfc7bab72ef86f1cc7cc9c78e728b34a6d1fd" + +[[package]] +name = "futures-util" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0d50a92467f8ba5dd6e3ee5d4bd04d73ab2e4e1c44474a0674821dfce14b79bc" +dependencies = [ + "futures-channel", + "futures-core", + "futures-io", + "futures-macro", + "futures-sink", + "futures-task", + "memchr", + "pin-project-lite", + "slab", +] + +[[package]] +name = "generic-array" +version = "0.14.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "85649ca51fd72272d7821adaf274ad91c288277713d9c18820d8499a7ff69e9a" +dependencies = [ + "typenum", + "version_check", +] + +[[package]] +name = "glob" +version = "0.3.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e4eba85ea1d0a966a983acd07deee566e67395d2d96b6fb39e62b5a833f1eb0b" + +[[package]] +name = "half" +version = "2.7.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" +dependencies = [ + "cfg-if", + "crunchy", + "zerocopy", +] + +[[package]] +name = "hashbrown" +version = "0.14.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e5274423e17b7c9fc20b6e7e208532f9b19825d82dfd615708b70edd83df41f1" + +[[package]] +name = "hashbrown" +version = "0.15.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9229cfe53dfd69f0609a49f65461bd93001ea1ef889cd5529dd176593f5338a1" +dependencies = [ + "foldhash", +] + +[[package]] +name = "hashbrown" +version = "0.17.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ed5909b6e89a2db4456e54cd5f673791d7eca6732202bbf2a9cc504fe2f9b84a" + +[[package]] +name = "home" +version = "0.5.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cc627f471c528ff0c4a49e1d5e60450c8f6461dd6d10ba9dcd3a61d3dff7728d" +dependencies = [ + "windows-sys 0.61.2", +] + +[[package]] +name = "indexmap" +version = "2.14.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cc4e190f5d26ca7051642629da2c52fc03bde85a03197c99408dcd291734c855" +dependencies = [ + "equivalent", + "hashbrown 0.17.1", +] + +[[package]] +name = "itertools" +version = "0.12.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ba291022dbbd398a455acf126c1e341954079855bc60dfdda641363bd6922569" +dependencies = [ + "either", +] + +[[package]] +name = "itoa" +version = "1.0.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8f42a60cbdf9a97f5d2305f08a87dc4e09308d1276d28c869c684d7777685682" + +[[package]] +name = "lazy_static" +version = "1.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "20870f649af7073d53e38067b2a84312175d56ea15217e1b15bc83506ec50afb" + +[[package]] +name = "lazycell" +version = "1.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "830d08ce1d1d941e6b30645f1a0eb5643013d835ce3779a5fc208261dbe10f55" + +[[package]] +name = "libc" +version = "0.2.191" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "65cd87c9a6d60536580b316ad985726104730a5151c99845f2490428ef91240a" + +[[package]] +name = "libloading" +version = "0.8.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d7c4b02199fee7c5d21a5ae7d8cfa79a6ef5bb2fc834d6e9058e89c825efdc55" +dependencies = [ + "cfg-if", + "windows-link", +] + +[[package]] +name = "libnvvm-sys" +version = "0.2.1" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "cuda-target-spec", + "libloading", + "thiserror 2.0.21", +] + +[[package]] +name = "linux-raw-sys" +version = "0.4.15" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d26c52dbd32dccf2d10cac7725f8eae5296885fb5703b261f7d0a0739ec807ab" + +[[package]] +name = "lock_api" +version = "0.4.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "224399e74b87b5f3557511d98dff8b14089b3dadafcab6bb93eab67d3aace965" +dependencies = [ + "scopeguard", +] + +[[package]] +name = "log" +version = "0.4.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9f8bd3e56ce4dfc153cf470fffbfa98c7620958b312ca5c3a4b8d5181fd13c6" + +[[package]] +name = "memchr" +version = "2.8.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cf8baf1c55e62ffcace7a9f06f4bd9cd3f0c4beb022d3b367256b91b87513d98" + +[[package]] +name = "minimal-lexical" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "68354c5c6bd36d73ff3feceb05efa59b6acb7626617f4962be322a825e61f79a" + +[[package]] +name = "nom" +version = "7.1.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d273983c5a657a70a3e8f2a01329822f3b8c8172b73826411a55751e404a0a4a" +dependencies = [ + "memchr", + "minimal-lexical", +] + +[[package]] +name = "nvjitlink-sys" +version = "0.2.1" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "libloading", + "thiserror 2.0.21", +] + +[[package]] +name = "object" +version = "0.36.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "62948e14d923ea95ea2c7c86c71013138b66525b86bdc08d2dcc262bdb497b87" +dependencies = [ + "crc32fast", + "hashbrown 0.15.5", + "indexmap", + "memchr", +] + +[[package]] +name = "once_cell" +version = "1.21.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9f7c3e4beb33f85d45ae3e3a1792185706c8e16d043238c593331cc7cd313b50" + +[[package]] +name = "oxide-artifacts" +version = "0.2.1" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "object", +] + +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "serde", + "serde_json", + "thiserror 2.0.21", +] + +[[package]] +name = "parking_lot_core" +version = "0.9.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2621685985a2ebf1c516881c026032ac7deafcda1a2c9b7850dc81e3dfcb64c1" +dependencies = [ + "cfg-if", + "libc", + "redox_syscall", + "smallvec", + "windows-link", +] + +[[package]] +name = "pin-project-lite" +version = "0.2.17" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a89322df9ebe1c1578d689c92318e070967d1042b512afbe49518723f4e6d5cd" + +[[package]] +name = "prettyplease" +version = "0.2.37" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "479ca8adacdd7ce8f1fb39ce9ecccbfe93a3f1344b3d0d97f20bc0196208f62b" +dependencies = [ + "proc-macro2", + "syn 2.0.119", +] + +[[package]] +name = "proc-macro2" +version = "1.0.107" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "985e7ec9bb745e6ce6535b544d84d6cd6f7ad8bd711c398938ae983b91a766d9" +dependencies = [ + "unicode-ident", +] + +[[package]] +name = "ptx-parse" +version = "0.1.0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" + +[[package]] +name = "quote" +version = "1.0.47" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1fbf4db142a473a8d80c26bbf18454ed458bf8d26c8219c331daecfdbd079001" +dependencies = [ + "proc-macro2", +] + +[[package]] +name = "redox_syscall" +version = "0.5.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ed2bf2547551a7053d6fdfafda3f938979645c44812fbfcda098faae3f1a362d" +dependencies = [ + "bitflags", +] + +[[package]] +name = "regex" +version = "1.13.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f020237b6c8eed93db2e2cb53c00c60a8e1bc73da7d073199a1180401450218d" +dependencies = [ + "aho-corasick", + "memchr", + "regex-automata", + "regex-syntax", +] + +[[package]] +name = "regex-automata" +version = "0.4.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ad8553b9b26413251cbf30e620595c7a41b3887f03da04579c0e6b0d6a06b4b2" +dependencies = [ + "aho-corasick", + "memchr", + "regex-syntax", +] + +[[package]] +name = "regex-syntax" +version = "0.8.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d6f6ff9a378485b298a5286656da665ba74413d36db0979633275d2e708145d4" + +[[package]] +name = "reserved-oxide-symbols" +version = "0.2.1" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" + +[[package]] +name = "rustc-hash" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "08d43f7aa6b08d49f382cde6a7982047c3426db949b1424bc4b7ec9ae12c6ce2" + +[[package]] +name = "rustc-hash" +version = "2.1.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6b1e7f9a428571be2dc5bc0505c13fb6bf936822b894ec87abf8a08a4e51742d" + +[[package]] +name = "rustix" +version = "0.38.44" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fdb5bc1ae2baa591800df16c9ca78619bf65c0488b41b96ccec5d11220d8c154" +dependencies = [ + "bitflags", + "errno", + "libc", + "linux-raw-sys", + "windows-sys 0.59.0", +] + +[[package]] +name = "scopeguard" +version = "1.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "94143f37725109f92c262ed2cf5e59bce7498c01bcc1502d7b9afe439a4e9f49" + +[[package]] +name = "serde" +version = "1.0.229" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4148590afebada386688f18773da617792bf2ef03ffc1e4cbd2b1d45b023e0ba" +dependencies = [ + "serde_core", + "serde_derive", +] + +[[package]] +name = "serde_core" +version = "1.0.229" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "67dca2c9c51e58a4791a4b1ed58308b39c64224d349a935ab5039aa360942a48" +dependencies = [ + "serde_derive", +] + +[[package]] +name = "serde_derive" +version = "1.0.229" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e7a5d71263a5a7d47b41f6b3f06ba276f10cc18b0931f1799f710578e2309348" +dependencies = [ + "proc-macro2", + "quote", + "syn 3.0.9", +] + +[[package]] +name = "serde_json" +version = "1.0.154" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e7e9cc8b1b85264074fbcc02a88680c4096b1e47df8f739dceb03bf482f04bd6" +dependencies = [ + "itoa", + "memchr", + "serde", + "serde_core", + "zmij", +] + +[[package]] +name = "sha2" +version = "0.10.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a7507d819769d01a365ab707794a4084392c824f54a7a6a7862f8c3d0892b283" +dependencies = [ + "cfg-if", + "cpufeatures", + "digest", +] + +[[package]] +name = "shlex" +version = "1.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0fda2ff0d084019ba4d7c6f371c95d8fd75ce3524c3cb8fb653a3023f6323e64" + +[[package]] +name = "slab" +version = "0.4.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0c790de23124f9ab44544d7ac05d60440adc586479ce501c1d6d7da3cd8c9cf5" + +[[package]] +name = "smallvec" +version = "1.16.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5b3dc8af474f516a851ff4bd12db780f948b9250ad37211e4eec0bccea54e01b" + +[[package]] +name = "syn" +version = "2.0.119" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "872831b642d1a07999a962a351ed35b955ea2cfc8f3862091e2a240a84f17297" +dependencies = [ + "proc-macro2", + "quote", + "unicode-ident", +] + +[[package]] +name = "syn" +version = "3.0.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d78c8dee4c7bf0e14673097256fed6142ce9d3b85a408189d07482442145823b" +dependencies = [ + "proc-macro2", + "quote", + "unicode-ident", +] + +[[package]] +name = "thiserror" +version = "1.0.69" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b6aaf5339b578ea85b50e080feb250a3e8ae8cfcdff9a461c9ec2904bc923f52" +dependencies = [ + "thiserror-impl 1.0.69", +] + +[[package]] +name = "thiserror" +version = "2.0.21" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "09e52cb86a36cede5cb101bf8908837b3e4c6e5e59fe7fd85c23fb56200d189e" +dependencies = [ + "thiserror-impl 2.0.21", +] + +[[package]] +name = "thiserror-impl" +version = "1.0.69" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4fee6c4efc90059e10f81e6d42c60a18f76588c3d74cb83a0b242a2b6c7504c1" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", +] + +[[package]] +name = "thiserror-impl" +version = "2.0.21" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fe5197923287db20a58125f0bc85c062f7f2c892de97b18c356f9efb14b28524" +dependencies = [ + "proc-macro2", + "quote", + "syn 3.0.9", +] + +[[package]] +name = "typenum" +version = "1.20.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b6f5e870be6c3b371b77fe0ee0bafb859fa4964b4404c27de1d380043c4dda20" + +[[package]] +name = "unicode-ident" +version = "1.0.27" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a2c754d6c33795a1c324727428e5a7dedb5b06195f9890bdbcba760d3e246563" + +[[package]] +name = "version_check" +version = "0.9.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0b928f33d975fc6ad9f86c8f283853ad26bdd5b10b7f1542aa2fa15e2289105a" + +[[package]] +name = "which" +version = "4.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "87ba24419a2078cd2b0f2ede2691b6c66d8e47836da3b6db8265ebad47afbfc7" +dependencies = [ + "either", + "home", + "once_cell", + "rustix", +] + +[[package]] +name = "windows-link" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f0805222e57f7521d6a62e36fa9163bc891acd422f971defe97d64e70d0a4fe5" + +[[package]] +name = "windows-sys" +version = "0.59.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1e38bc4d79ed67fd075bcc251a1c39b32a1776bbe92e5bef1f0bf1f8c531853b" +dependencies = [ + "windows-targets", +] + +[[package]] +name = "windows-sys" +version = "0.61.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ae137229bcbd6cdf0f7b80a31df61766145077ddf49416a728b02cb3921ff3fc" +dependencies = [ + "windows-link", +] + +[[package]] +name = "windows-targets" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9b724f72796e036ab90c1021d4780d4d3d648aca59e491e6b98e725b84e99973" +dependencies = [ + "windows_aarch64_gnullvm", + "windows_aarch64_msvc", + "windows_i686_gnu", + "windows_i686_gnullvm", + "windows_i686_msvc", + "windows_x86_64_gnu", + "windows_x86_64_gnullvm", + "windows_x86_64_msvc", +] + +[[package]] +name = "windows_aarch64_gnullvm" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "32a4622180e7a0ec044bb555404c800bc9fd9ec262ec147edd5989ccd0c02cd3" + +[[package]] +name = "windows_aarch64_msvc" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "09ec2a7bb152e2252b53fa7803150007879548bc709c039df7627cabbd05d469" + +[[package]] +name = "windows_i686_gnu" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8e9b5ad5ab802e97eb8e295ac6720e509ee4c243f69d781394014ebfe8bbfa0b" + +[[package]] +name = "windows_i686_gnullvm" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0eee52d38c090b3caa76c563b86c3a4bd71ef1a819287c19d586d7334ae8ed66" + +[[package]] +name = "windows_i686_msvc" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "240948bc05c5e7c6dabba28bf89d89ffce3e303022809e73deaefe4f6ec56c66" + +[[package]] +name = "windows_x86_64_gnu" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "147a5c80aabfbf0c7d901cb5895d1de30ef2907eb21fbbab29ca94c5b08b1a78" + +[[package]] +name = "windows_x86_64_gnullvm" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "24d5b23dc417412679681396f2b49f3de8c1473deb516bd34410872eff51ed0d" + +[[package]] +name = "windows_x86_64_msvc" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "589f6da84c646204747d1270a2a5661ea66ed1cced2631d546fdfb155959f9ec" + +[[package]] +name = "zerocopy" +version = "0.8.63" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e5fe1f8f1b06191a00962174c61aa5005e0bb391a6d80d07e24d115c01a92ed8" +dependencies = [ + "zerocopy-derive", +] + +[[package]] +name = "zerocopy-derive" +version = "0.8.63" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "863ad3ac83293fb4d740aedbfdc9240dd8d1a50c1099acd76ce80ce7c7230c7f" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", +] + +[[package]] +name = "zmij" +version = "1.0.23" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "29666d0abbfad1e3dc4dcf6144730dd3a3ab225bbbdac83319345b1b44ccfc1b" diff --git a/rust/cuda-coop/tests/reduce/Cargo.toml b/rust/cuda-coop/tests/reduce/Cargo.toml new file mode 100644 index 000000000000..2f41f1164a28 --- /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 = "96caed46c776eb49226137eac849e795b83b7ce0" } +cuda-host = { git = "https://github.com/tpn/cuda-rust.git", rev = "96caed46c776eb49226137eac849e795b83b7ce0", features = ["async"] } +cuda-core = { git = "https://github.com/tpn/cuda-rust.git", rev = "96caed46c776eb49226137eac849e795b83b7ce0" } +cuda-async = { git = "https://github.com/tpn/cuda-rust.git", rev = "96caed46c776eb49226137eac849e795b83b7ce0" } diff --git a/rust/cuda-coop/tests/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..f4f5bdf7a06f --- /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=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "libnvvm-sys", + "nvjitlink-sys", + "serde", + "sha2", + "thiserror 2.0.21", +] + +[[package]] +name = "cuda-async" +version = "0.5.0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "anyhow", + "cuda-bindings", + "cuda-core", + "dashmap", + "futures", + "half", + "once_cell", + "rustc-hash 2.1.3", + "thiserror 1.0.69", +] + +[[package]] +name = "cuda-bindings" +version = "0.5.0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "bindgen", + "libloading", + "prettyplease", + "proc-macro2", + "quote", + "syn 2.0.119", +] + +[[package]] +name = "cuda-coop" +version = "0.0.0" +dependencies = [ + "cuda-coop-codegen", + "cuda-device", +] + +[[package]] +name = "cuda-coop-codegen" +version = "0.0.0" +dependencies = [ + "libloading", + "serde", + "serde_json", + "sha2", +] + +[[package]] +name = "cuda-core" +version = "0.5.0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "anyhow", + "cuda-bindings", + "cuda-core-derive", + "half", + "oxide-artifacts", +] + +[[package]] +name = "cuda-core-derive" +version = "0.5.0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", +] + +[[package]] +name = "cuda-device" +version = "0.2.1" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "cuda-macros", +] + +[[package]] +name = "cuda-host" +version = "0.2.1" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "cuda-artifact-finalizer", + "cuda-async", + "cuda-core", + "cuda-device", + "cuda-macros", + "half", + "libc", + "oxide-artifacts", + "ptx-parse", + "sha2", + "thiserror 2.0.21", +] + +[[package]] +name = "cuda-macros" +version = "0.2.1" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "oxide-tile-export", + "proc-macro2", + "quote", + "reserved-oxide-symbols", + "serde", + "serde_json", + "syn 2.0.119", +] + +[[package]] +name = "cuda-target-spec" +version = "0.1.0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" + +[[package]] +name = "dashmap" +version = "6.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e6361d5c062261c78a176addb82d4c821ae42bed6089de0e12603cd25de2059c" +dependencies = [ + "cfg-if", + "crossbeam-utils", + "hashbrown 0.14.5", + "lock_api", + "once_cell", + "parking_lot_core", +] + +[[package]] +name = "digest" +version = "0.10.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9ed9a281f7bc9b7576e61468ba615a66a5c8cfdff42420a70aa82701a3b1e292" +dependencies = [ + "block-buffer", + "crypto-common", +] + +[[package]] +name = "either" +version = "1.19.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0e9c71c2167ca323c882b99918929403426e2373ea17242ff5653e0d5e1058be" + +[[package]] +name = "equivalent" +version = "1.0.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "00d174d5400e5e8fd687ad1049e2f578285fa914201b1af7e8b112a4546bd826" + +[[package]] +name = "errno" +version = "0.3.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "39cab71617ae0d63f51a36d69f866391735b51691dbda63cf6f96d042b63efeb" +dependencies = [ + "libc", + "windows-sys 0.61.2", +] + +[[package]] +name = "foldhash" +version = "0.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d9c4f5dac5e15c24eb999c26181a6ca40b39fe946cbe4c263c7209467bc83af2" + +[[package]] +name = "futures" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9a31d2a3fbaaeb2af2368bbdd904aa8e812d3c04a1ee10d3171f52d556e5d0a3" +dependencies = [ + "futures-channel", + "futures-core", + "futures-executor", + "futures-io", + "futures-sink", + "futures-task", + "futures-util", +] + +[[package]] +name = "futures-channel" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b1f9e3d69d39e4862ffed03ed071a76f9a13ba1d9109d355b0f0aa6b15e393c4" +dependencies = [ + "futures-core", + "futures-sink", +] + +[[package]] +name = "futures-core" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "92d699e522242e69e3003b94ecc1f960f3a5e015aa7c5d7486e65ad01dd94f5e" + +[[package]] +name = "futures-executor" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "031b47cf1a3c6cc8bc2fc76cd437f521619387907d469316e7c0bc278f1f5432" +dependencies = [ + "futures-core", + "futures-task", + "futures-util", +] + +[[package]] +name = "futures-io" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "53c0fa8157de1303bfffdaa1cc2a673bfffb60102f76b0ef4441659124373fed" + +[[package]] +name = "futures-macro" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9fb9654ba8355388abeb8dcb4fc62f511300867002afc858860463bdd9fe0c44" +dependencies = [ + "proc-macro2", + "quote", + "syn 3.0.9", +] + +[[package]] +name = "futures-sink" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1944426bf7d03f1d14f708785e4b33efd750b36d48a157b836b3efc15ede8e1d" + +[[package]] +name = "futures-task" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cd417de3d1d015fc3bfd2b1ea46dfc7bab72ef86f1cc7cc9c78e728b34a6d1fd" + +[[package]] +name = "futures-util" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0d50a92467f8ba5dd6e3ee5d4bd04d73ab2e4e1c44474a0674821dfce14b79bc" +dependencies = [ + "futures-channel", + "futures-core", + "futures-io", + "futures-macro", + "futures-sink", + "futures-task", + "memchr", + "pin-project-lite", + "slab", +] + +[[package]] +name = "generic-array" +version = "0.14.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "85649ca51fd72272d7821adaf274ad91c288277713d9c18820d8499a7ff69e9a" +dependencies = [ + "typenum", + "version_check", +] + +[[package]] +name = "glob" +version = "0.3.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e4eba85ea1d0a966a983acd07deee566e67395d2d96b6fb39e62b5a833f1eb0b" + +[[package]] +name = "half" +version = "2.7.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" +dependencies = [ + "cfg-if", + "crunchy", + "zerocopy", +] + +[[package]] +name = "hashbrown" +version = "0.14.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e5274423e17b7c9fc20b6e7e208532f9b19825d82dfd615708b70edd83df41f1" + +[[package]] +name = "hashbrown" +version = "0.15.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9229cfe53dfd69f0609a49f65461bd93001ea1ef889cd5529dd176593f5338a1" +dependencies = [ + "foldhash", +] + +[[package]] +name = "hashbrown" +version = "0.17.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ed5909b6e89a2db4456e54cd5f673791d7eca6732202bbf2a9cc504fe2f9b84a" + +[[package]] +name = "home" +version = "0.5.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cc627f471c528ff0c4a49e1d5e60450c8f6461dd6d10ba9dcd3a61d3dff7728d" +dependencies = [ + "windows-sys 0.61.2", +] + +[[package]] +name = "indexmap" +version = "2.14.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cc4e190f5d26ca7051642629da2c52fc03bde85a03197c99408dcd291734c855" +dependencies = [ + "equivalent", + "hashbrown 0.17.1", +] + +[[package]] +name = "itertools" +version = "0.12.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ba291022dbbd398a455acf126c1e341954079855bc60dfdda641363bd6922569" +dependencies = [ + "either", +] + +[[package]] +name = "itoa" +version = "1.0.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8f42a60cbdf9a97f5d2305f08a87dc4e09308d1276d28c869c684d7777685682" + +[[package]] +name = "lazy_static" +version = "1.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "20870f649af7073d53e38067b2a84312175d56ea15217e1b15bc83506ec50afb" + +[[package]] +name = "lazycell" +version = "1.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "830d08ce1d1d941e6b30645f1a0eb5643013d835ce3779a5fc208261dbe10f55" + +[[package]] +name = "libc" +version = "0.2.191" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "65cd87c9a6d60536580b316ad985726104730a5151c99845f2490428ef91240a" + +[[package]] +name = "libloading" +version = "0.8.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d7c4b02199fee7c5d21a5ae7d8cfa79a6ef5bb2fc834d6e9058e89c825efdc55" +dependencies = [ + "cfg-if", + "windows-link", +] + +[[package]] +name = "libnvvm-sys" +version = "0.2.1" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "cuda-target-spec", + "libloading", + "thiserror 2.0.21", +] + +[[package]] +name = "linux-raw-sys" +version = "0.4.15" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d26c52dbd32dccf2d10cac7725f8eae5296885fb5703b261f7d0a0739ec807ab" + +[[package]] +name = "lock_api" +version = "0.4.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "224399e74b87b5f3557511d98dff8b14089b3dadafcab6bb93eab67d3aace965" +dependencies = [ + "scopeguard", +] + +[[package]] +name = "log" +version = "0.4.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9f8bd3e56ce4dfc153cf470fffbfa98c7620958b312ca5c3a4b8d5181fd13c6" + +[[package]] +name = "memchr" +version = "2.8.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cf8baf1c55e62ffcace7a9f06f4bd9cd3f0c4beb022d3b367256b91b87513d98" + +[[package]] +name = "minimal-lexical" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "68354c5c6bd36d73ff3feceb05efa59b6acb7626617f4962be322a825e61f79a" + +[[package]] +name = "nom" +version = "7.1.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d273983c5a657a70a3e8f2a01329822f3b8c8172b73826411a55751e404a0a4a" +dependencies = [ + "memchr", + "minimal-lexical", +] + +[[package]] +name = "nvjitlink-sys" +version = "0.2.1" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "libloading", + "thiserror 2.0.21", +] + +[[package]] +name = "object" +version = "0.36.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "62948e14d923ea95ea2c7c86c71013138b66525b86bdc08d2dcc262bdb497b87" +dependencies = [ + "crc32fast", + "hashbrown 0.15.5", + "indexmap", + "memchr", +] + +[[package]] +name = "once_cell" +version = "1.21.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9f7c3e4beb33f85d45ae3e3a1792185706c8e16d043238c593331cc7cd313b50" + +[[package]] +name = "oxide-artifacts" +version = "0.2.1" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "object", +] + +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "serde", + "serde_json", + "thiserror 2.0.21", +] + +[[package]] +name = "parking_lot_core" +version = "0.9.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2621685985a2ebf1c516881c026032ac7deafcda1a2c9b7850dc81e3dfcb64c1" +dependencies = [ + "cfg-if", + "libc", + "redox_syscall", + "smallvec", + "windows-link", +] + +[[package]] +name = "pin-project-lite" +version = "0.2.17" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a89322df9ebe1c1578d689c92318e070967d1042b512afbe49518723f4e6d5cd" + +[[package]] +name = "prettyplease" +version = "0.2.37" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "479ca8adacdd7ce8f1fb39ce9ecccbfe93a3f1344b3d0d97f20bc0196208f62b" +dependencies = [ + "proc-macro2", + "syn 2.0.119", +] + +[[package]] +name = "proc-macro2" +version = "1.0.107" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "985e7ec9bb745e6ce6535b544d84d6cd6f7ad8bd711c398938ae983b91a766d9" +dependencies = [ + "unicode-ident", +] + +[[package]] +name = "ptx-parse" +version = "0.1.0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" + +[[package]] +name = "quote" +version = "1.0.47" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1fbf4db142a473a8d80c26bbf18454ed458bf8d26c8219c331daecfdbd079001" +dependencies = [ + "proc-macro2", +] + +[[package]] +name = "redox_syscall" +version = "0.5.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ed2bf2547551a7053d6fdfafda3f938979645c44812fbfcda098faae3f1a362d" +dependencies = [ + "bitflags", +] + +[[package]] +name = "regex" +version = "1.13.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f020237b6c8eed93db2e2cb53c00c60a8e1bc73da7d073199a1180401450218d" +dependencies = [ + "aho-corasick", + "memchr", + "regex-automata", + "regex-syntax", +] + +[[package]] +name = "regex-automata" +version = "0.4.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ad8553b9b26413251cbf30e620595c7a41b3887f03da04579c0e6b0d6a06b4b2" +dependencies = [ + "aho-corasick", + "memchr", + "regex-syntax", +] + +[[package]] +name = "regex-syntax" +version = "0.8.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d6f6ff9a378485b298a5286656da665ba74413d36db0979633275d2e708145d4" + +[[package]] +name = "reserved-oxide-symbols" +version = "0.2.1" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" + +[[package]] +name = "rustc-hash" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "08d43f7aa6b08d49f382cde6a7982047c3426db949b1424bc4b7ec9ae12c6ce2" + +[[package]] +name = "rustc-hash" +version = "2.1.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6b1e7f9a428571be2dc5bc0505c13fb6bf936822b894ec87abf8a08a4e51742d" + +[[package]] +name = "rustix" +version = "0.38.44" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fdb5bc1ae2baa591800df16c9ca78619bf65c0488b41b96ccec5d11220d8c154" +dependencies = [ + "bitflags", + "errno", + "libc", + "linux-raw-sys", + "windows-sys 0.59.0", +] + +[[package]] +name = "scopeguard" +version = "1.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "94143f37725109f92c262ed2cf5e59bce7498c01bcc1502d7b9afe439a4e9f49" + +[[package]] +name = "serde" +version = "1.0.229" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4148590afebada386688f18773da617792bf2ef03ffc1e4cbd2b1d45b023e0ba" +dependencies = [ + "serde_core", + "serde_derive", +] + +[[package]] +name = "serde_core" +version = "1.0.229" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "67dca2c9c51e58a4791a4b1ed58308b39c64224d349a935ab5039aa360942a48" +dependencies = [ + "serde_derive", +] + +[[package]] +name = "serde_derive" +version = "1.0.229" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e7a5d71263a5a7d47b41f6b3f06ba276f10cc18b0931f1799f710578e2309348" +dependencies = [ + "proc-macro2", + "quote", + "syn 3.0.9", +] + +[[package]] +name = "serde_json" +version = "1.0.154" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e7e9cc8b1b85264074fbcc02a88680c4096b1e47df8f739dceb03bf482f04bd6" +dependencies = [ + "itoa", + "memchr", + "serde", + "serde_core", + "zmij", +] + +[[package]] +name = "sha2" +version = "0.10.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a7507d819769d01a365ab707794a4084392c824f54a7a6a7862f8c3d0892b283" +dependencies = [ + "cfg-if", + "cpufeatures", + "digest", +] + +[[package]] +name = "shlex" +version = "1.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0fda2ff0d084019ba4d7c6f371c95d8fd75ce3524c3cb8fb653a3023f6323e64" + +[[package]] +name = "slab" +version = "0.4.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0c790de23124f9ab44544d7ac05d60440adc586479ce501c1d6d7da3cd8c9cf5" + +[[package]] +name = "smallvec" +version = "1.16.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5b3dc8af474f516a851ff4bd12db780f948b9250ad37211e4eec0bccea54e01b" + +[[package]] +name = "syn" +version = "2.0.119" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "872831b642d1a07999a962a351ed35b955ea2cfc8f3862091e2a240a84f17297" +dependencies = [ + "proc-macro2", + "quote", + "unicode-ident", +] + +[[package]] +name = "syn" +version = "3.0.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d78c8dee4c7bf0e14673097256fed6142ce9d3b85a408189d07482442145823b" +dependencies = [ + "proc-macro2", + "quote", + "unicode-ident", +] + +[[package]] +name = "thiserror" +version = "1.0.69" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b6aaf5339b578ea85b50e080feb250a3e8ae8cfcdff9a461c9ec2904bc923f52" +dependencies = [ + "thiserror-impl 1.0.69", +] + +[[package]] +name = "thiserror" +version = "2.0.21" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "09e52cb86a36cede5cb101bf8908837b3e4c6e5e59fe7fd85c23fb56200d189e" +dependencies = [ + "thiserror-impl 2.0.21", +] + +[[package]] +name = "thiserror-impl" +version = "1.0.69" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4fee6c4efc90059e10f81e6d42c60a18f76588c3d74cb83a0b242a2b6c7504c1" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", +] + +[[package]] +name = "thiserror-impl" +version = "2.0.21" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fe5197923287db20a58125f0bc85c062f7f2c892de97b18c356f9efb14b28524" +dependencies = [ + "proc-macro2", + "quote", + "syn 3.0.9", +] + +[[package]] +name = "typenum" +version = "1.20.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b6f5e870be6c3b371b77fe0ee0bafb859fa4964b4404c27de1d380043c4dda20" + +[[package]] +name = "unicode-ident" +version = "1.0.27" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a2c754d6c33795a1c324727428e5a7dedb5b06195f9890bdbcba760d3e246563" + +[[package]] +name = "version_check" +version = "0.9.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0b928f33d975fc6ad9f86c8f283853ad26bdd5b10b7f1542aa2fa15e2289105a" + +[[package]] +name = "which" +version = "4.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "87ba24419a2078cd2b0f2ede2691b6c66d8e47836da3b6db8265ebad47afbfc7" +dependencies = [ + "either", + "home", + "once_cell", + "rustix", +] + +[[package]] +name = "windows-link" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f0805222e57f7521d6a62e36fa9163bc891acd422f971defe97d64e70d0a4fe5" + +[[package]] +name = "windows-sys" +version = "0.59.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1e38bc4d79ed67fd075bcc251a1c39b32a1776bbe92e5bef1f0bf1f8c531853b" +dependencies = [ + "windows-targets", +] + +[[package]] +name = "windows-sys" +version = "0.61.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ae137229bcbd6cdf0f7b80a31df61766145077ddf49416a728b02cb3921ff3fc" +dependencies = [ + "windows-link", +] + +[[package]] +name = "windows-targets" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9b724f72796e036ab90c1021d4780d4d3d648aca59e491e6b98e725b84e99973" +dependencies = [ + "windows_aarch64_gnullvm", + "windows_aarch64_msvc", + "windows_i686_gnu", + "windows_i686_gnullvm", + "windows_i686_msvc", + "windows_x86_64_gnu", + "windows_x86_64_gnullvm", + "windows_x86_64_msvc", +] + +[[package]] +name = "windows_aarch64_gnullvm" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "32a4622180e7a0ec044bb555404c800bc9fd9ec262ec147edd5989ccd0c02cd3" + +[[package]] +name = "windows_aarch64_msvc" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "09ec2a7bb152e2252b53fa7803150007879548bc709c039df7627cabbd05d469" + +[[package]] +name = "windows_i686_gnu" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8e9b5ad5ab802e97eb8e295ac6720e509ee4c243f69d781394014ebfe8bbfa0b" + +[[package]] +name = "windows_i686_gnullvm" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0eee52d38c090b3caa76c563b86c3a4bd71ef1a819287c19d586d7334ae8ed66" + +[[package]] +name = "windows_i686_msvc" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "240948bc05c5e7c6dabba28bf89d89ffce3e303022809e73deaefe4f6ec56c66" + +[[package]] +name = "windows_x86_64_gnu" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "147a5c80aabfbf0c7d901cb5895d1de30ef2907eb21fbbab29ca94c5b08b1a78" + +[[package]] +name = "windows_x86_64_gnullvm" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "24d5b23dc417412679681396f2b49f3de8c1473deb516bd34410872eff51ed0d" + +[[package]] +name = "windows_x86_64_msvc" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "589f6da84c646204747d1270a2a5661ea66ed1cced2631d546fdfb155959f9ec" + +[[package]] +name = "zerocopy" +version = "0.8.63" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e5fe1f8f1b06191a00962174c61aa5005e0bb391a6d80d07e24d115c01a92ed8" +dependencies = [ + "zerocopy-derive", +] + +[[package]] +name = "zerocopy-derive" +version = "0.8.63" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "863ad3ac83293fb4d740aedbfdc9240dd8d1a50c1099acd76ce80ce7c7230c7f" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", +] + +[[package]] +name = "zmij" +version = "1.0.23" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "29666d0abbfad1e3dc4dcf6144730dd3a3ab225bbbdac83319345b1b44ccfc1b" diff --git a/rust/cuda-coop/tests/reduce_batched/Cargo.toml b/rust/cuda-coop/tests/reduce_batched/Cargo.toml new file mode 100644 index 000000000000..39ad33ee47d1 --- /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 = "96caed46c776eb49226137eac849e795b83b7ce0" } +cuda-host = { git = "https://github.com/tpn/cuda-rust.git", rev = "96caed46c776eb49226137eac849e795b83b7ce0", features = ["async"] } +cuda-core = { git = "https://github.com/tpn/cuda-rust.git", rev = "96caed46c776eb49226137eac849e795b83b7ce0" } +cuda-async = { git = "https://github.com/tpn/cuda-rust.git", rev = "96caed46c776eb49226137eac849e795b83b7ce0" } diff --git a/rust/cuda-coop/tests/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..0809eccb39af --- /dev/null +++ b/rust/cuda-coop/tests/scan/Cargo.lock @@ -0,0 +1,1011 @@ +# This file is automatically @generated by Cargo. +# It is not intended for manual editing. +version = 4 + +[[package]] +name = "aho-corasick" +version = "1.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c982642fa9e8606056828ee9a8505737230110bb1099153c79efe865c59d12ba" +dependencies = [ + "memchr", +] + +[[package]] +name = "anyhow" +version = "1.0.104" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "330a5ed07fa54e4702c9d6c4174f74427fc0ef6e214bbd677ae50a5099946470" + +[[package]] +name = "bindgen" +version = "0.69.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "271383c67ccabffb7381723dea0672a673f292304fcb45c01cc648c7a8d58088" +dependencies = [ + "bitflags", + "cexpr", + "clang-sys", + "itertools", + "lazy_static", + "lazycell", + "log", + "prettyplease", + "proc-macro2", + "quote", + "regex", + "rustc-hash 1.1.0", + "shlex", + "syn 2.0.119", + "which", +] + +[[package]] +name = "bitflags" +version = "2.13.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3ded4057c258ba199e2d26386d3af3780957ecaee6c4ef4041c6b4b8b97c0b06" + +[[package]] +name = "block-buffer" +version = "0.10.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3078c7629b62d3f0439517fa394996acacc5cbc91c5a20d8c658e77abd503a71" +dependencies = [ + "generic-array", +] + +[[package]] +name = "cexpr" +version = "0.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6fac387a98bb7c37292057cffc56d62ecb629900026402633ae9160df93a8766" +dependencies = [ + "nom", +] + +[[package]] +name = "cfg-if" +version = "1.0.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4e7648175b45a9a48536d676f68d918270699102aa8dab5496df06904c914600" + +[[package]] +name = "clang-sys" +version = "1.9.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "157a8ba7b480713b56f4c09fd13fc3e0a22a5dfab8097ba61cbc5feef950788a" +dependencies = [ + "glob", + "libc", + "libloading", +] + +[[package]] +name = "coop-scan-tests" +version = "0.0.0" +dependencies = [ + "cuda-async", + "cuda-coop", + "cuda-core", + "cuda-device", + "cuda-host", +] + +[[package]] +name = "cpufeatures" +version = "0.2.17" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "59ed5838eebb26a2bb2e58f6d5b5316989ae9d08bab10e0e6d103e656d1b0280" +dependencies = [ + "libc", +] + +[[package]] +name = "crc32fast" +version = "1.5.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0ba106e844045ddc02c20768ac4b0b7cb45f2cc2ecbd833ac9f5dacb183c2055" +dependencies = [ + "cfg-if", +] + +[[package]] +name = "crossbeam-utils" +version = "0.8.23" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a31eee39dddec8330830986fcd7625edb5a24ec90ea038215273bbc3adb08ac6" + +[[package]] +name = "crunchy" +version = "0.2.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "460fbee9c2c2f33933d720630a6a0bac33ba7053db5344fac858d4b8952d77d5" + +[[package]] +name = "crypto-common" +version = "0.1.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "78c8292055d1c1df0cce5d180393dc8cce0abec0a7102adb6c7b1eef6016d60a" +dependencies = [ + "generic-array", + "typenum", +] + +[[package]] +name = "cuda-artifact-finalizer" +version = "0.2.1" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "libnvvm-sys", + "nvjitlink-sys", + "serde", + "sha2", + "thiserror 2.0.21", +] + +[[package]] +name = "cuda-async" +version = "0.5.0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "anyhow", + "cuda-bindings", + "cuda-core", + "dashmap", + "futures", + "half", + "once_cell", + "rustc-hash 2.1.3", + "thiserror 1.0.69", +] + +[[package]] +name = "cuda-bindings" +version = "0.5.0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "bindgen", + "libloading", + "prettyplease", + "proc-macro2", + "quote", + "syn 2.0.119", +] + +[[package]] +name = "cuda-coop" +version = "0.0.0" +dependencies = [ + "cuda-coop-codegen", + "cuda-device", +] + +[[package]] +name = "cuda-coop-codegen" +version = "0.0.0" +dependencies = [ + "libloading", + "serde", + "serde_json", + "sha2", +] + +[[package]] +name = "cuda-core" +version = "0.5.0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "anyhow", + "cuda-bindings", + "cuda-core-derive", + "half", + "oxide-artifacts", +] + +[[package]] +name = "cuda-core-derive" +version = "0.5.0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", +] + +[[package]] +name = "cuda-device" +version = "0.2.1" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "cuda-macros", +] + +[[package]] +name = "cuda-host" +version = "0.2.1" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "cuda-artifact-finalizer", + "cuda-async", + "cuda-core", + "cuda-device", + "cuda-macros", + "half", + "libc", + "oxide-artifacts", + "ptx-parse", + "sha2", + "thiserror 2.0.21", +] + +[[package]] +name = "cuda-macros" +version = "0.2.1" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "oxide-tile-export", + "proc-macro2", + "quote", + "reserved-oxide-symbols", + "serde", + "serde_json", + "syn 2.0.119", +] + +[[package]] +name = "cuda-target-spec" +version = "0.1.0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" + +[[package]] +name = "dashmap" +version = "6.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e6361d5c062261c78a176addb82d4c821ae42bed6089de0e12603cd25de2059c" +dependencies = [ + "cfg-if", + "crossbeam-utils", + "hashbrown 0.14.5", + "lock_api", + "once_cell", + "parking_lot_core", +] + +[[package]] +name = "digest" +version = "0.10.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9ed9a281f7bc9b7576e61468ba615a66a5c8cfdff42420a70aa82701a3b1e292" +dependencies = [ + "block-buffer", + "crypto-common", +] + +[[package]] +name = "either" +version = "1.19.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0e9c71c2167ca323c882b99918929403426e2373ea17242ff5653e0d5e1058be" + +[[package]] +name = "equivalent" +version = "1.0.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "00d174d5400e5e8fd687ad1049e2f578285fa914201b1af7e8b112a4546bd826" + +[[package]] +name = "errno" +version = "0.3.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "39cab71617ae0d63f51a36d69f866391735b51691dbda63cf6f96d042b63efeb" +dependencies = [ + "libc", + "windows-sys 0.61.2", +] + +[[package]] +name = "foldhash" +version = "0.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d9c4f5dac5e15c24eb999c26181a6ca40b39fe946cbe4c263c7209467bc83af2" + +[[package]] +name = "futures" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9a31d2a3fbaaeb2af2368bbdd904aa8e812d3c04a1ee10d3171f52d556e5d0a3" +dependencies = [ + "futures-channel", + "futures-core", + "futures-executor", + "futures-io", + "futures-sink", + "futures-task", + "futures-util", +] + +[[package]] +name = "futures-channel" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b1f9e3d69d39e4862ffed03ed071a76f9a13ba1d9109d355b0f0aa6b15e393c4" +dependencies = [ + "futures-core", + "futures-sink", +] + +[[package]] +name = "futures-core" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "92d699e522242e69e3003b94ecc1f960f3a5e015aa7c5d7486e65ad01dd94f5e" + +[[package]] +name = "futures-executor" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "031b47cf1a3c6cc8bc2fc76cd437f521619387907d469316e7c0bc278f1f5432" +dependencies = [ + "futures-core", + "futures-task", + "futures-util", +] + +[[package]] +name = "futures-io" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "53c0fa8157de1303bfffdaa1cc2a673bfffb60102f76b0ef4441659124373fed" + +[[package]] +name = "futures-macro" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9fb9654ba8355388abeb8dcb4fc62f511300867002afc858860463bdd9fe0c44" +dependencies = [ + "proc-macro2", + "quote", + "syn 3.0.9", +] + +[[package]] +name = "futures-sink" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1944426bf7d03f1d14f708785e4b33efd750b36d48a157b836b3efc15ede8e1d" + +[[package]] +name = "futures-task" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cd417de3d1d015fc3bfd2b1ea46dfc7bab72ef86f1cc7cc9c78e728b34a6d1fd" + +[[package]] +name = "futures-util" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0d50a92467f8ba5dd6e3ee5d4bd04d73ab2e4e1c44474a0674821dfce14b79bc" +dependencies = [ + "futures-channel", + "futures-core", + "futures-io", + "futures-macro", + "futures-sink", + "futures-task", + "memchr", + "pin-project-lite", + "slab", +] + +[[package]] +name = "generic-array" +version = "0.14.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "85649ca51fd72272d7821adaf274ad91c288277713d9c18820d8499a7ff69e9a" +dependencies = [ + "typenum", + "version_check", +] + +[[package]] +name = "glob" +version = "0.3.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e4eba85ea1d0a966a983acd07deee566e67395d2d96b6fb39e62b5a833f1eb0b" + +[[package]] +name = "half" +version = "2.7.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" +dependencies = [ + "cfg-if", + "crunchy", + "zerocopy", +] + +[[package]] +name = "hashbrown" +version = "0.14.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e5274423e17b7c9fc20b6e7e208532f9b19825d82dfd615708b70edd83df41f1" + +[[package]] +name = "hashbrown" +version = "0.15.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9229cfe53dfd69f0609a49f65461bd93001ea1ef889cd5529dd176593f5338a1" +dependencies = [ + "foldhash", +] + +[[package]] +name = "hashbrown" +version = "0.17.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ed5909b6e89a2db4456e54cd5f673791d7eca6732202bbf2a9cc504fe2f9b84a" + +[[package]] +name = "home" +version = "0.5.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cc627f471c528ff0c4a49e1d5e60450c8f6461dd6d10ba9dcd3a61d3dff7728d" +dependencies = [ + "windows-sys 0.61.2", +] + +[[package]] +name = "indexmap" +version = "2.14.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cc4e190f5d26ca7051642629da2c52fc03bde85a03197c99408dcd291734c855" +dependencies = [ + "equivalent", + "hashbrown 0.17.1", +] + +[[package]] +name = "itertools" +version = "0.12.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ba291022dbbd398a455acf126c1e341954079855bc60dfdda641363bd6922569" +dependencies = [ + "either", +] + +[[package]] +name = "itoa" +version = "1.0.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8f42a60cbdf9a97f5d2305f08a87dc4e09308d1276d28c869c684d7777685682" + +[[package]] +name = "lazy_static" +version = "1.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "20870f649af7073d53e38067b2a84312175d56ea15217e1b15bc83506ec50afb" + +[[package]] +name = "lazycell" +version = "1.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "830d08ce1d1d941e6b30645f1a0eb5643013d835ce3779a5fc208261dbe10f55" + +[[package]] +name = "libc" +version = "0.2.191" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "65cd87c9a6d60536580b316ad985726104730a5151c99845f2490428ef91240a" + +[[package]] +name = "libloading" +version = "0.8.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d7c4b02199fee7c5d21a5ae7d8cfa79a6ef5bb2fc834d6e9058e89c825efdc55" +dependencies = [ + "cfg-if", + "windows-link", +] + +[[package]] +name = "libnvvm-sys" +version = "0.2.1" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "cuda-target-spec", + "libloading", + "thiserror 2.0.21", +] + +[[package]] +name = "linux-raw-sys" +version = "0.4.15" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d26c52dbd32dccf2d10cac7725f8eae5296885fb5703b261f7d0a0739ec807ab" + +[[package]] +name = "lock_api" +version = "0.4.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "224399e74b87b5f3557511d98dff8b14089b3dadafcab6bb93eab67d3aace965" +dependencies = [ + "scopeguard", +] + +[[package]] +name = "log" +version = "0.4.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9f8bd3e56ce4dfc153cf470fffbfa98c7620958b312ca5c3a4b8d5181fd13c6" + +[[package]] +name = "memchr" +version = "2.8.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cf8baf1c55e62ffcace7a9f06f4bd9cd3f0c4beb022d3b367256b91b87513d98" + +[[package]] +name = "minimal-lexical" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "68354c5c6bd36d73ff3feceb05efa59b6acb7626617f4962be322a825e61f79a" + +[[package]] +name = "nom" +version = "7.1.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d273983c5a657a70a3e8f2a01329822f3b8c8172b73826411a55751e404a0a4a" +dependencies = [ + "memchr", + "minimal-lexical", +] + +[[package]] +name = "nvjitlink-sys" +version = "0.2.1" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "libloading", + "thiserror 2.0.21", +] + +[[package]] +name = "object" +version = "0.36.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "62948e14d923ea95ea2c7c86c71013138b66525b86bdc08d2dcc262bdb497b87" +dependencies = [ + "crc32fast", + "hashbrown 0.15.5", + "indexmap", + "memchr", +] + +[[package]] +name = "once_cell" +version = "1.21.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9f7c3e4beb33f85d45ae3e3a1792185706c8e16d043238c593331cc7cd313b50" + +[[package]] +name = "oxide-artifacts" +version = "0.2.1" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "object", +] + +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "serde", + "serde_json", + "thiserror 2.0.21", +] + +[[package]] +name = "parking_lot_core" +version = "0.9.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2621685985a2ebf1c516881c026032ac7deafcda1a2c9b7850dc81e3dfcb64c1" +dependencies = [ + "cfg-if", + "libc", + "redox_syscall", + "smallvec", + "windows-link", +] + +[[package]] +name = "pin-project-lite" +version = "0.2.17" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a89322df9ebe1c1578d689c92318e070967d1042b512afbe49518723f4e6d5cd" + +[[package]] +name = "prettyplease" +version = "0.2.37" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "479ca8adacdd7ce8f1fb39ce9ecccbfe93a3f1344b3d0d97f20bc0196208f62b" +dependencies = [ + "proc-macro2", + "syn 2.0.119", +] + +[[package]] +name = "proc-macro2" +version = "1.0.107" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "985e7ec9bb745e6ce6535b544d84d6cd6f7ad8bd711c398938ae983b91a766d9" +dependencies = [ + "unicode-ident", +] + +[[package]] +name = "ptx-parse" +version = "0.1.0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" + +[[package]] +name = "quote" +version = "1.0.47" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1fbf4db142a473a8d80c26bbf18454ed458bf8d26c8219c331daecfdbd079001" +dependencies = [ + "proc-macro2", +] + +[[package]] +name = "redox_syscall" +version = "0.5.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ed2bf2547551a7053d6fdfafda3f938979645c44812fbfcda098faae3f1a362d" +dependencies = [ + "bitflags", +] + +[[package]] +name = "regex" +version = "1.13.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f020237b6c8eed93db2e2cb53c00c60a8e1bc73da7d073199a1180401450218d" +dependencies = [ + "aho-corasick", + "memchr", + "regex-automata", + "regex-syntax", +] + +[[package]] +name = "regex-automata" +version = "0.4.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ad8553b9b26413251cbf30e620595c7a41b3887f03da04579c0e6b0d6a06b4b2" +dependencies = [ + "aho-corasick", + "memchr", + "regex-syntax", +] + +[[package]] +name = "regex-syntax" +version = "0.8.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d6f6ff9a378485b298a5286656da665ba74413d36db0979633275d2e708145d4" + +[[package]] +name = "reserved-oxide-symbols" +version = "0.2.1" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" + +[[package]] +name = "rustc-hash" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "08d43f7aa6b08d49f382cde6a7982047c3426db949b1424bc4b7ec9ae12c6ce2" + +[[package]] +name = "rustc-hash" +version = "2.1.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6b1e7f9a428571be2dc5bc0505c13fb6bf936822b894ec87abf8a08a4e51742d" + +[[package]] +name = "rustix" +version = "0.38.44" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fdb5bc1ae2baa591800df16c9ca78619bf65c0488b41b96ccec5d11220d8c154" +dependencies = [ + "bitflags", + "errno", + "libc", + "linux-raw-sys", + "windows-sys 0.59.0", +] + +[[package]] +name = "scopeguard" +version = "1.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "94143f37725109f92c262ed2cf5e59bce7498c01bcc1502d7b9afe439a4e9f49" + +[[package]] +name = "serde" +version = "1.0.229" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4148590afebada386688f18773da617792bf2ef03ffc1e4cbd2b1d45b023e0ba" +dependencies = [ + "serde_core", + "serde_derive", +] + +[[package]] +name = "serde_core" +version = "1.0.229" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "67dca2c9c51e58a4791a4b1ed58308b39c64224d349a935ab5039aa360942a48" +dependencies = [ + "serde_derive", +] + +[[package]] +name = "serde_derive" +version = "1.0.229" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e7a5d71263a5a7d47b41f6b3f06ba276f10cc18b0931f1799f710578e2309348" +dependencies = [ + "proc-macro2", + "quote", + "syn 3.0.9", +] + +[[package]] +name = "serde_json" +version = "1.0.154" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e7e9cc8b1b85264074fbcc02a88680c4096b1e47df8f739dceb03bf482f04bd6" +dependencies = [ + "itoa", + "memchr", + "serde", + "serde_core", + "zmij", +] + +[[package]] +name = "sha2" +version = "0.10.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a7507d819769d01a365ab707794a4084392c824f54a7a6a7862f8c3d0892b283" +dependencies = [ + "cfg-if", + "cpufeatures", + "digest", +] + +[[package]] +name = "shlex" +version = "1.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0fda2ff0d084019ba4d7c6f371c95d8fd75ce3524c3cb8fb653a3023f6323e64" + +[[package]] +name = "slab" +version = "0.4.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0c790de23124f9ab44544d7ac05d60440adc586479ce501c1d6d7da3cd8c9cf5" + +[[package]] +name = "smallvec" +version = "1.16.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5b3dc8af474f516a851ff4bd12db780f948b9250ad37211e4eec0bccea54e01b" + +[[package]] +name = "syn" +version = "2.0.119" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "872831b642d1a07999a962a351ed35b955ea2cfc8f3862091e2a240a84f17297" +dependencies = [ + "proc-macro2", + "quote", + "unicode-ident", +] + +[[package]] +name = "syn" +version = "3.0.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d78c8dee4c7bf0e14673097256fed6142ce9d3b85a408189d07482442145823b" +dependencies = [ + "proc-macro2", + "quote", + "unicode-ident", +] + +[[package]] +name = "thiserror" +version = "1.0.69" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b6aaf5339b578ea85b50e080feb250a3e8ae8cfcdff9a461c9ec2904bc923f52" +dependencies = [ + "thiserror-impl 1.0.69", +] + +[[package]] +name = "thiserror" +version = "2.0.21" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "09e52cb86a36cede5cb101bf8908837b3e4c6e5e59fe7fd85c23fb56200d189e" +dependencies = [ + "thiserror-impl 2.0.21", +] + +[[package]] +name = "thiserror-impl" +version = "1.0.69" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4fee6c4efc90059e10f81e6d42c60a18f76588c3d74cb83a0b242a2b6c7504c1" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", +] + +[[package]] +name = "thiserror-impl" +version = "2.0.21" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fe5197923287db20a58125f0bc85c062f7f2c892de97b18c356f9efb14b28524" +dependencies = [ + "proc-macro2", + "quote", + "syn 3.0.9", +] + +[[package]] +name = "typenum" +version = "1.20.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b6f5e870be6c3b371b77fe0ee0bafb859fa4964b4404c27de1d380043c4dda20" + +[[package]] +name = "unicode-ident" +version = "1.0.27" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a2c754d6c33795a1c324727428e5a7dedb5b06195f9890bdbcba760d3e246563" + +[[package]] +name = "version_check" +version = "0.9.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0b928f33d975fc6ad9f86c8f283853ad26bdd5b10b7f1542aa2fa15e2289105a" + +[[package]] +name = "which" +version = "4.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "87ba24419a2078cd2b0f2ede2691b6c66d8e47836da3b6db8265ebad47afbfc7" +dependencies = [ + "either", + "home", + "once_cell", + "rustix", +] + +[[package]] +name = "windows-link" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f0805222e57f7521d6a62e36fa9163bc891acd422f971defe97d64e70d0a4fe5" + +[[package]] +name = "windows-sys" +version = "0.59.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1e38bc4d79ed67fd075bcc251a1c39b32a1776bbe92e5bef1f0bf1f8c531853b" +dependencies = [ + "windows-targets", +] + +[[package]] +name = "windows-sys" +version = "0.61.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ae137229bcbd6cdf0f7b80a31df61766145077ddf49416a728b02cb3921ff3fc" +dependencies = [ + "windows-link", +] + +[[package]] +name = "windows-targets" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9b724f72796e036ab90c1021d4780d4d3d648aca59e491e6b98e725b84e99973" +dependencies = [ + "windows_aarch64_gnullvm", + "windows_aarch64_msvc", + "windows_i686_gnu", + "windows_i686_gnullvm", + "windows_i686_msvc", + "windows_x86_64_gnu", + "windows_x86_64_gnullvm", + "windows_x86_64_msvc", +] + +[[package]] +name = "windows_aarch64_gnullvm" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "32a4622180e7a0ec044bb555404c800bc9fd9ec262ec147edd5989ccd0c02cd3" + +[[package]] +name = "windows_aarch64_msvc" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "09ec2a7bb152e2252b53fa7803150007879548bc709c039df7627cabbd05d469" + +[[package]] +name = "windows_i686_gnu" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8e9b5ad5ab802e97eb8e295ac6720e509ee4c243f69d781394014ebfe8bbfa0b" + +[[package]] +name = "windows_i686_gnullvm" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0eee52d38c090b3caa76c563b86c3a4bd71ef1a819287c19d586d7334ae8ed66" + +[[package]] +name = "windows_i686_msvc" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "240948bc05c5e7c6dabba28bf89d89ffce3e303022809e73deaefe4f6ec56c66" + +[[package]] +name = "windows_x86_64_gnu" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "147a5c80aabfbf0c7d901cb5895d1de30ef2907eb21fbbab29ca94c5b08b1a78" + +[[package]] +name = "windows_x86_64_gnullvm" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "24d5b23dc417412679681396f2b49f3de8c1473deb516bd34410872eff51ed0d" + +[[package]] +name = "windows_x86_64_msvc" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "589f6da84c646204747d1270a2a5661ea66ed1cced2631d546fdfb155959f9ec" + +[[package]] +name = "zerocopy" +version = "0.8.63" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e5fe1f8f1b06191a00962174c61aa5005e0bb391a6d80d07e24d115c01a92ed8" +dependencies = [ + "zerocopy-derive", +] + +[[package]] +name = "zerocopy-derive" +version = "0.8.63" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "863ad3ac83293fb4d740aedbfdc9240dd8d1a50c1099acd76ce80ce7c7230c7f" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", +] + +[[package]] +name = "zmij" +version = "1.0.23" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "29666d0abbfad1e3dc4dcf6144730dd3a3ab225bbbdac83319345b1b44ccfc1b" diff --git a/rust/cuda-coop/tests/scan/Cargo.toml b/rust/cuda-coop/tests/scan/Cargo.toml new file mode 100644 index 000000000000..05aab44b5af6 --- /dev/null +++ b/rust/cuda-coop/tests/scan/Cargo.toml @@ -0,0 +1,17 @@ +# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. +# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +[package] +name = "coop-scan-tests" +version = "0.0.0" +edition = "2024" +publish = false + +[workspace] + +[dependencies] +cuda-coop = { path = "../.." } +cuda-device = { git = "https://github.com/tpn/cuda-rust.git", rev = "96caed46c776eb49226137eac849e795b83b7ce0" } +cuda-host = { git = "https://github.com/tpn/cuda-rust.git", rev = "96caed46c776eb49226137eac849e795b83b7ce0", features = ["async"] } +cuda-core = { git = "https://github.com/tpn/cuda-rust.git", rev = "96caed46c776eb49226137eac849e795b83b7ce0" } +cuda-async = { git = "https://github.com/tpn/cuda-rust.git", rev = "96caed46c776eb49226137eac849e795b83b7ce0" } diff --git a/rust/cuda-coop/tests/scan/src/main.rs b/rust/cuda-coop/tests/scan/src/main.rs new file mode 100644 index 000000000000..ceba29760c27 --- /dev/null +++ b/rust/cuda-coop/tests/scan/src/main.rs @@ -0,0 +1,540 @@ +// Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +//! Public Scan values, validity, preservation, groups, and scratch composition. +use cuda_coop as coop; +use cuda_device::{kernel, launch_contract}; +use cuda_host::cuda_module; + +#[cuda_module] +mod kernels { + use super::*; + + #[inline(never)] + fn observe( + input: &coop::ThreadData, + before: &coop::ThreadData, + result: &coop::ThreadData, + ) -> bool { + let mut matches = true; + let mut i = 0; + while i < ITEMS_PER_THREAD { + matches &= input.get(i) == before.get(i) && result.get(i).is_some(); + i += 1; + } + matches && result.get(ITEMS_PER_THREAD).is_none() + } + + #[kernel] + #[launch_contract(domain = 3, block = (4, 4, 4), dynamic_shared = inferred)] + pub fn scan_sum( + input: &[T], + mut output: coop::DisjointSlice>, + valid: usize, + zero: T, + invalid: bool, + ) { + let block = coop::this_block(); + let offset = block.rank_at::() as usize * 64 * ITEMS_PER_THREAD; + let mut values = coop::ThreadData::<_, ITEMS_PER_THREAD>::new(); + let options = coop::LoadOptions::default() + .offset(offset) + .valid_items(valid); + if invalid { + coop::load(block, input, &mut values, options); + } else { + coop::load(block, input, &mut values, options.oob_default(zero)); + } + let before = values; + let mut result = coop::scan(block, &values, Default::default()); + let matches = observe(&values, &before, &result); + // Store every result, including prefixes whose input was default-filled. + coop::store( + block, + &mut output, + &result, + coop::StoreOptions::default().offset(offset), + ); + result.write(0, zero); + assert!(matches && observe(&values, &before, &result)); + } + + #[kernel] + #[launch_contract(domain = 1, block = (64, 1, 1), dynamic_shared = inferred)] + pub fn scan_seeded( + input: &[T], + mut output: coop::DisjointSlice>, + valid: usize, + zero: T, + seed: T, + ) { + let block = coop::this_block(); + let mut values = coop::ThreadData::<_, ITEMS_PER_THREAD>::new(); + coop::load( + block, + input, + &mut values, + coop::LoadOptions::default() + .valid_items(valid) + .oob_default(zero), + ); + let before = values; + let result = coop::exclusive_scan( + block, + &mut values, + coop::ScanOptions::default().initial_value(seed), + ); + let matches = observe(&values, &before, &result); + coop::store(block, &mut output, &result, Default::default()); + assert!(matches && observe(&values, &before, &result)); + } + + #[kernel] + #[launch_contract(domain = 3, block = (4, 4, 4), dynamic_shared = inferred)] + pub fn scan_operators( + input: &[i32], + mut output: coop::DisjointSlice>, + mode: u32, + ) { + let block = coop::this_block(); + let mut values = coop::ThreadData::<_, 4>::new(); + coop::load(block, input, &mut values, Default::default()); + let before = values; + let result = match mode { + 0 => coop::exclusive_sum(block, &values, Default::default()), + 1 => coop::inclusive_sum(block, &values, Default::default()), + 2 => coop::exclusive_scan( + block, + &values, + coop::ScanOptions::default() + .scan_op(coop::Multiplies) + .initial_value(1), + ), + 3 => coop::inclusive_scan( + block, + &values, + coop::ScanOptions::default().scan_op(coop::Min), + ), + 4 => coop::scan( + block, + &values, + coop::ScanOptions::default() + .mode(coop::Inclusive) + .scan_op(coop::Max), + ), + 5 => coop::exclusive_scan( + block, + &values, + coop::ScanOptions::default() + .scan_op(coop::BitAnd) + .initial_value(255), + ), + 6 => coop::inclusive_scan( + block, + &values, + coop::ScanOptions::default().scan_op(coop::BitOr), + ), + 7 => coop::inclusive_scan( + block, + &values, + coop::ScanOptions::default().scan_op(coop::BitXor), + ), + 8 => coop::exclusive_scan( + block, + &values, + coop::ScanOptions::default() + .initial_value(3) + .algorithm(coop::RakingMemoize), + ), + _ => coop::exclusive_scan( + block, + &values, + coop::ScanOptions::default() + .initial_value(3) + .algorithm(coop::WarpScans), + ), + }; + let matches = observe(&values, &before, &result); + coop::store(block, &mut output, &result, Default::default()); + assert!(matches && observe(&values, &before, &result)); + } + + #[inline(never)] + fn with_storage( + input: &[i32], + output: &mut coop::DisjointSlice>, + valid: usize, + scratch: &mut coop::TempStorage<0, 0, AUTO, false>, + ) { + let block = coop::this_block(); + let mut values = coop::ThreadData::<_, 4>::new(); + coop::load( + block, + input, + &mut values, + coop::LoadOptions::default() + .valid_items(valid) + .oob_default(0) + .algorithm(coop::Transpose) + .temp_storage(scratch), + ); + if !AUTO { + cuda_device::thread::sync_threads(); + } + let before = values; + let mut matches = true; + let mut iteration = 0; + while iteration < 2 { + let result = coop::exclusive_scan( + block, + &values, + coop::ScanOptions::default() + .initial_value(5) + .temp_storage(scratch), + ); + matches &= observe(&values, &before, &result); + if !AUTO { + cuda_device::thread::sync_threads(); + } + coop::store(block, &mut *output, &result, Default::default()); + // Complete the global Store before repeating the same output region. + cuda_device::thread::sync_threads(); + iteration += 1; + } + assert!(matches); + } + + #[kernel] + #[launch_contract(domain = 1, block = (64, 1, 1), dynamic_shared = inferred)] + pub fn scan_manual( + input: &[i32], + mut output: coop::DisjointSlice>, + valid: usize, + ) { + let mut scratch = coop::TempStorage::new(); + with_storage(input, &mut output, valid, &mut scratch); + } + + #[kernel] + #[launch_contract(domain = 1, block = (64, 1, 1), dynamic_shared = inferred)] + pub fn scan_auto( + input: &[i32], + mut output: coop::DisjointSlice>, + valid: usize, + ) { + let mut scratch = coop::TempStorage::new().auto_sync(); + with_storage(input, &mut output, valid, &mut scratch); + } + + #[kernel] + #[launch_contract(domain = 1, block = (64, 1, 1), dynamic_shared = inferred)] + pub fn scan_warp( + input: &[i32], + mut output: coop::DisjointSlice>, + seeded: bool, + ) { + let group = coop::this_warp().group_by::<8>(); + let block = coop::this_block(); + let offset = block.rank_at::() as usize * 64; + let mut values = coop::ThreadData::<_, 1>::new(); + coop::load( + group, + input, + &mut values, + coop::LoadOptions::default().offset(offset), + ); + let value = values.get(0).unwrap_or(0); + let q = block.rank_at::() as i32 * 4 + group.rank_at::() as i32; + let result = if seeded { + coop::exclusive_scan( + group, + value, + coop::ScanOptions::default().initial_value(q + 5), + ) + } else { + coop::exclusive_sum(group, value, Default::default()) + }; + let matches = values.get(0) == Some(value); + coop::store( + group, + &mut output, + result, + coop::StoreOptions::default().offset(offset), + ); + assert!(matches && values.get(0) == Some(value)); + } + + #[kernel] + #[launch_contract(domain = 1, block = (64, 1, 1), dynamic_shared = inferred)] + pub fn scan_physical_warp( + input: &[i32], + mut output: coop::DisjointSlice>, + ) { + let group = coop::this_warp(); + let offset = coop::this_block().rank_at::() as usize * 64; + let mut values = coop::ThreadData::<_, 1>::new(); + coop::load( + group, + input, + &mut values, + coop::LoadOptions::default().offset(offset), + ); + let value = values.get(0).unwrap_or(0); + let result = coop::exclusive_sum(group, value, Default::default()); + coop::store( + group, + &mut output, + result, + coop::StoreOptions::default().offset(offset), + ); + assert!(values.get(0) == Some(value)); + } + + #[kernel] + #[launch_contract(domain = 1, block = (64, 1, 1), dynamic_shared = inferred)] + pub fn scan_block_scalar( + input: &[i32], + mut output: coop::DisjointSlice>, + ) { + let group = coop::this_block(); + let offset = group.rank_at::() as usize * 64; + let mut values = coop::ThreadData::<_, 1>::new(); + coop::load( + group, + input, + &mut values, + coop::LoadOptions::default().offset(offset), + ); + let value = values.get(0).unwrap_or(0); + let result = coop::scan(group, value, Default::default()); + coop::store( + group, + &mut output, + result, + coop::StoreOptions::default().offset(offset), + ); + assert!(values.get(0) == Some(value)); + } +} + +fn main() -> Result<(), Box> { + use cuda_async::simt::device_context::{init_device_contexts, with_cuda_context}; + use cuda_async::simt::device_operation::DeviceOperation; + use cuda_core::{DeviceBuffer, LaunchConfig1D, LaunchConfig3D}; + + init_device_contexts(0, 1)?; + let context = with_cuda_context(0, |context| context.clone())?; + let stream = context.default_stream(); + // SAFETY: the typed declarations and embedded kernels compile together. + let module = unsafe { kernels::load_async(0)? }; + if std::env::args().any(|arg| arg == "invalid-input") { + let source = DeviceBuffer::from_host(&stream, &[1_i32; 256])?; + let target = DeviceBuffer::from_host(&stream, &[0_i32; 256])?; + let plan = + module.prepare_scan_sum::(LaunchConfig3D::auto_shared((1, 1, 1), (4, 4, 4)))?; + module + .scan_sum_async_owned::(&plan, source, target, 1, 0, true) + .sync()?; + stream.synchronize()?; + panic!("unwritten Scan input unexpectedly completed"); + } + let mut cases = 0; + macro_rules! sum_cases { + ($ty:ty, $items:literal) => {{ + let tile = 64 * $items; + let input: Vec<$ty> = (0..2 * tile).map(|i| (i % 8 == 0) as u8 as $ty).collect(); + let plan = module.prepare_scan_sum::<$ty, $items>(LaunchConfig3D::auto_shared( + (2, 1, 1), + (4, 4, 4), + ))?; + for valid in [0, 1, tile - 1, tile] { + let source = DeviceBuffer::from_host(&stream, &input)?; + let target = DeviceBuffer::from_host(&stream, &vec![0 as $ty; input.len()])?; + let (_, target) = module + .scan_sum_async_owned::<$ty, $items, _, _>( + &plan, source, target, valid, 0 as $ty, false, + ) + .sync()?; + let mut actual = vec![0 as $ty; input.len()]; + target.copy_to_host(&stream, &mut actual)?; + stream.synchronize()?; + let mut expected = Vec::with_capacity(input.len()); + for chunk in input.chunks(tile) { + let mut prefix = 0 as $ty; + for (i, &value) in chunk.iter().enumerate() { + expected.push(prefix); + if i < valid { + prefix += value; + } + } + } + assert_eq!( + actual, + expected, + "{} N{} valid{}", + stringify!($ty), + $items, + valid + ); + cases += 1; + } + }}; + } + sum_cases!(i8, 1); + sum_cases!(i8, 4); + sum_cases!(u8, 1); + sum_cases!(u8, 4); + sum_cases!(i16, 1); + sum_cases!(i16, 4); + sum_cases!(u16, 1); + sum_cases!(u16, 4); + sum_cases!(i32, 1); + sum_cases!(i32, 4); + sum_cases!(u32, 1); + sum_cases!(u32, 4); + sum_cases!(i64, 1); + sum_cases!(i64, 4); + sum_cases!(u64, 1); + sum_cases!(u64, 4); + sum_cases!(f32, 1); + sum_cases!(f32, 4); + sum_cases!(f64, 1); + sum_cases!(f64, 4); + sum_cases!(i32, 9); + macro_rules! seeded_cases { + ($ty:ty, $items:literal, $seed:expr) => {{ + let tile = 64 * $items; + let input: Vec<$ty> = (0..tile).map(|i| (i % 8 == 0) as u8 as $ty).collect(); + let plan = + module.prepare_scan_seeded::<$ty, $items>(LaunchConfig1D::auto_shared(1, 64))?; + for valid in [tile - 1, tile] { + let source = DeviceBuffer::from_host(&stream, &input)?; + let target = DeviceBuffer::from_host(&stream, &vec![0 as $ty; tile])?; + let (_, target) = module + .scan_seeded_async_owned::<$ty, $items, _, _>( + &plan, source, target, valid, 0 as $ty, $seed, + ) + .sync()?; + let mut actual = vec![0 as $ty; tile]; + target.copy_to_host(&stream, &mut actual)?; + stream.synchronize()?; + let mut prefix: $ty = $seed; + for (i, (&value, &got)) in input.iter().zip(&actual).enumerate() { + assert_eq!(got, prefix); + if i < valid { + prefix += value; + } + } + cases += 1; + } + }}; + } + seeded_cases!(i8, 4, -3_i8); + seeded_cases!(u16, 1, 7_u16); + let plan = module.prepare_scan_operators(LaunchConfig3D::auto_shared((1, 1, 1), (4, 4, 4)))?; + for mode in 0..10 { + let input: Vec = (0..256) + .map(|i| { + if mode == 2 { + if i == 0 { 2 } else { 1 } + } else { + i % 7 + 1 + } + }) + .collect(); + let source = DeviceBuffer::from_host(&stream, &input)?; + let target = DeviceBuffer::from_host(&stream, &[0_i32; 256])?; + let (_, target) = module + .scan_operators_async_owned(&plan, source, target, mode) + .sync()?; + let mut actual = vec![0_i32; 256]; + target.copy_to_host(&stream, &mut actual)?; + stream.synchronize()?; + let inclusive = matches!(mode, 1 | 3 | 4 | 6 | 7); + let mut prefix = match mode { + 2 => 1, + 3 => i32::MAX, + 4 => i32::MIN, + 5 => 255, + 8 | 9 => 3, + _ => 0, + }; + for (&value, &got) in input.iter().zip(&actual) { + let next = match mode { + 2 => prefix * value, + 3 => prefix.min(value), + 4 => prefix.max(value), + 5 => prefix & value, + 6 => prefix | value, + 7 => prefix ^ value, + _ => prefix + value, + }; + assert_eq!(got, if inclusive { next } else { prefix }, "mode {mode}"); + prefix = next; + } + cases += 1; + } + macro_rules! storage_cases { + ($prepare:ident, $launch:ident) => {{ + let input: Vec = (0..256).map(|i| (i % 8 == 0) as i32).collect(); + let plan = module.$prepare(LaunchConfig1D::auto_shared(1, 64))?; + for valid in [255, 256] { + let source = DeviceBuffer::from_host(&stream, &input)?; + let target = DeviceBuffer::from_host(&stream, &[0_i32; 256])?; + let (_, target) = module.$launch(&plan, source, target, valid).sync()?; + let mut actual = vec![0_i32; 256]; + target.copy_to_host(&stream, &mut actual)?; + stream.synchronize()?; + let mut prefix = 5; + for (i, (&value, &got)) in input.iter().zip(&actual).enumerate() { + assert_eq!(got, prefix); + if i < valid { + prefix += value; + } + } + cases += 1; + } + }}; + } + storage_cases!(prepare_scan_manual, scan_manual_async_owned); + storage_cases!(prepare_scan_auto, scan_auto_async_owned); + let plan = module.prepare_scan_warp(LaunchConfig1D::auto_shared(3, 64))?; + for seeded in [false, true] { + let source = DeviceBuffer::from_host(&stream, &[1_i32; 192])?; + let target = DeviceBuffer::from_host(&stream, &[0_i32; 192])?; + let (_, target) = module + .scan_warp_async_owned(&plan, source, target, seeded) + .sync()?; + let mut actual = vec![0_i32; 192]; + target.copy_to_host(&stream, &mut actual)?; + stream.synchronize()?; + for (i, &got) in actual.iter().enumerate() { + let seed = if seeded { (i % 64 / 8) as i32 + 5 } else { 0 }; + assert_eq!(got, seed + (i % 8) as i32); + } + cases += 1; + } + macro_rules! scalar_case { + ($prepare:ident, $launch:ident, $width:literal) => {{ + let plan = module.$prepare(LaunchConfig1D::auto_shared(3, 64))?; + let source = DeviceBuffer::from_host(&stream, &[1_i32; 192])?; + let target = DeviceBuffer::from_host(&stream, &[0_i32; 192])?; + let (_, target) = module.$launch(&plan, source, target).sync()?; + let mut actual = vec![0_i32; 192]; + target.copy_to_host(&stream, &mut actual)?; + stream.synchronize()?; + for (i, &got) in actual.iter().enumerate() { + assert_eq!(got, (i % $width) as i32); + } + cases += 1; + }}; + } + scalar_case!( + prepare_scan_physical_warp, + scan_physical_warp_async_owned, + 32 + ); + scalar_case!(prepare_scan_block_scalar, scan_block_scalar_async_owned, 64); + println!("PASS: {cases} public Scan cases"); + Ok(()) +} 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..3d67b2380b39 --- /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=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "libnvvm-sys", + "nvjitlink-sys", + "serde", + "sha2", + "thiserror 2.0.21", +] + +[[package]] +name = "cuda-async" +version = "0.5.0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "anyhow", + "cuda-bindings", + "cuda-core", + "dashmap", + "futures", + "half", + "once_cell", + "rustc-hash 2.1.3", + "thiserror 1.0.69", +] + +[[package]] +name = "cuda-bindings" +version = "0.5.0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "bindgen", + "libloading", + "prettyplease", + "proc-macro2", + "quote", + "syn 2.0.119", +] + +[[package]] +name = "cuda-coop" +version = "0.0.0" +dependencies = [ + "cuda-coop-codegen", + "cuda-device", +] + +[[package]] +name = "cuda-coop-codegen" +version = "0.0.0" +dependencies = [ + "libloading", + "serde", + "serde_json", + "sha2", +] + +[[package]] +name = "cuda-core" +version = "0.5.0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "anyhow", + "cuda-bindings", + "cuda-core-derive", + "half", + "oxide-artifacts", +] + +[[package]] +name = "cuda-core-derive" +version = "0.5.0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", +] + +[[package]] +name = "cuda-device" +version = "0.2.1" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "cuda-macros", +] + +[[package]] +name = "cuda-host" +version = "0.2.1" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "cuda-artifact-finalizer", + "cuda-async", + "cuda-core", + "cuda-device", + "cuda-macros", + "half", + "libc", + "oxide-artifacts", + "ptx-parse", + "sha2", + "thiserror 2.0.21", +] + +[[package]] +name = "cuda-macros" +version = "0.2.1" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "oxide-tile-export", + "proc-macro2", + "quote", + "reserved-oxide-symbols", + "serde", + "serde_json", + "syn 2.0.119", +] + +[[package]] +name = "cuda-target-spec" +version = "0.1.0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" + +[[package]] +name = "dashmap" +version = "6.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e6361d5c062261c78a176addb82d4c821ae42bed6089de0e12603cd25de2059c" +dependencies = [ + "cfg-if", + "crossbeam-utils", + "hashbrown 0.14.5", + "lock_api", + "once_cell", + "parking_lot_core", +] + +[[package]] +name = "digest" +version = "0.10.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9ed9a281f7bc9b7576e61468ba615a66a5c8cfdff42420a70aa82701a3b1e292" +dependencies = [ + "block-buffer", + "crypto-common", +] + +[[package]] +name = "either" +version = "1.19.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0e9c71c2167ca323c882b99918929403426e2373ea17242ff5653e0d5e1058be" + +[[package]] +name = "equivalent" +version = "1.0.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "00d174d5400e5e8fd687ad1049e2f578285fa914201b1af7e8b112a4546bd826" + +[[package]] +name = "errno" +version = "0.3.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "39cab71617ae0d63f51a36d69f866391735b51691dbda63cf6f96d042b63efeb" +dependencies = [ + "libc", + "windows-sys 0.61.2", +] + +[[package]] +name = "foldhash" +version = "0.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d9c4f5dac5e15c24eb999c26181a6ca40b39fe946cbe4c263c7209467bc83af2" + +[[package]] +name = "futures" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9a31d2a3fbaaeb2af2368bbdd904aa8e812d3c04a1ee10d3171f52d556e5d0a3" +dependencies = [ + "futures-channel", + "futures-core", + "futures-executor", + "futures-io", + "futures-sink", + "futures-task", + "futures-util", +] + +[[package]] +name = "futures-channel" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b1f9e3d69d39e4862ffed03ed071a76f9a13ba1d9109d355b0f0aa6b15e393c4" +dependencies = [ + "futures-core", + "futures-sink", +] + +[[package]] +name = "futures-core" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "92d699e522242e69e3003b94ecc1f960f3a5e015aa7c5d7486e65ad01dd94f5e" + +[[package]] +name = "futures-executor" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "031b47cf1a3c6cc8bc2fc76cd437f521619387907d469316e7c0bc278f1f5432" +dependencies = [ + "futures-core", + "futures-task", + "futures-util", +] + +[[package]] +name = "futures-io" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "53c0fa8157de1303bfffdaa1cc2a673bfffb60102f76b0ef4441659124373fed" + +[[package]] +name = "futures-macro" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9fb9654ba8355388abeb8dcb4fc62f511300867002afc858860463bdd9fe0c44" +dependencies = [ + "proc-macro2", + "quote", + "syn 3.0.9", +] + +[[package]] +name = "futures-sink" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1944426bf7d03f1d14f708785e4b33efd750b36d48a157b836b3efc15ede8e1d" + +[[package]] +name = "futures-task" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cd417de3d1d015fc3bfd2b1ea46dfc7bab72ef86f1cc7cc9c78e728b34a6d1fd" + +[[package]] +name = "futures-util" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0d50a92467f8ba5dd6e3ee5d4bd04d73ab2e4e1c44474a0674821dfce14b79bc" +dependencies = [ + "futures-channel", + "futures-core", + "futures-io", + "futures-macro", + "futures-sink", + "futures-task", + "memchr", + "pin-project-lite", + "slab", +] + +[[package]] +name = "generic-array" +version = "0.14.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "85649ca51fd72272d7821adaf274ad91c288277713d9c18820d8499a7ff69e9a" +dependencies = [ + "typenum", + "version_check", +] + +[[package]] +name = "glob" +version = "0.3.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e4eba85ea1d0a966a983acd07deee566e67395d2d96b6fb39e62b5a833f1eb0b" + +[[package]] +name = "half" +version = "2.7.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" +dependencies = [ + "cfg-if", + "crunchy", + "zerocopy", +] + +[[package]] +name = "hashbrown" +version = "0.14.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e5274423e17b7c9fc20b6e7e208532f9b19825d82dfd615708b70edd83df41f1" + +[[package]] +name = "hashbrown" +version = "0.15.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9229cfe53dfd69f0609a49f65461bd93001ea1ef889cd5529dd176593f5338a1" +dependencies = [ + "foldhash", +] + +[[package]] +name = "hashbrown" +version = "0.17.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ed5909b6e89a2db4456e54cd5f673791d7eca6732202bbf2a9cc504fe2f9b84a" + +[[package]] +name = "home" +version = "0.5.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cc627f471c528ff0c4a49e1d5e60450c8f6461dd6d10ba9dcd3a61d3dff7728d" +dependencies = [ + "windows-sys 0.61.2", +] + +[[package]] +name = "indexmap" +version = "2.14.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cc4e190f5d26ca7051642629da2c52fc03bde85a03197c99408dcd291734c855" +dependencies = [ + "equivalent", + "hashbrown 0.17.1", +] + +[[package]] +name = "itertools" +version = "0.12.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ba291022dbbd398a455acf126c1e341954079855bc60dfdda641363bd6922569" +dependencies = [ + "either", +] + +[[package]] +name = "itoa" +version = "1.0.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8f42a60cbdf9a97f5d2305f08a87dc4e09308d1276d28c869c684d7777685682" + +[[package]] +name = "lazy_static" +version = "1.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "20870f649af7073d53e38067b2a84312175d56ea15217e1b15bc83506ec50afb" + +[[package]] +name = "lazycell" +version = "1.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "830d08ce1d1d941e6b30645f1a0eb5643013d835ce3779a5fc208261dbe10f55" + +[[package]] +name = "libc" +version = "0.2.191" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "65cd87c9a6d60536580b316ad985726104730a5151c99845f2490428ef91240a" + +[[package]] +name = "libloading" +version = "0.8.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d7c4b02199fee7c5d21a5ae7d8cfa79a6ef5bb2fc834d6e9058e89c825efdc55" +dependencies = [ + "cfg-if", + "windows-link", +] + +[[package]] +name = "libnvvm-sys" +version = "0.2.1" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "cuda-target-spec", + "libloading", + "thiserror 2.0.21", +] + +[[package]] +name = "linux-raw-sys" +version = "0.4.15" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d26c52dbd32dccf2d10cac7725f8eae5296885fb5703b261f7d0a0739ec807ab" + +[[package]] +name = "lock_api" +version = "0.4.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "224399e74b87b5f3557511d98dff8b14089b3dadafcab6bb93eab67d3aace965" +dependencies = [ + "scopeguard", +] + +[[package]] +name = "log" +version = "0.4.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9f8bd3e56ce4dfc153cf470fffbfa98c7620958b312ca5c3a4b8d5181fd13c6" + +[[package]] +name = "memchr" +version = "2.8.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cf8baf1c55e62ffcace7a9f06f4bd9cd3f0c4beb022d3b367256b91b87513d98" + +[[package]] +name = "minimal-lexical" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "68354c5c6bd36d73ff3feceb05efa59b6acb7626617f4962be322a825e61f79a" + +[[package]] +name = "nom" +version = "7.1.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d273983c5a657a70a3e8f2a01329822f3b8c8172b73826411a55751e404a0a4a" +dependencies = [ + "memchr", + "minimal-lexical", +] + +[[package]] +name = "nvjitlink-sys" +version = "0.2.1" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "libloading", + "thiserror 2.0.21", +] + +[[package]] +name = "object" +version = "0.36.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "62948e14d923ea95ea2c7c86c71013138b66525b86bdc08d2dcc262bdb497b87" +dependencies = [ + "crc32fast", + "hashbrown 0.15.5", + "indexmap", + "memchr", +] + +[[package]] +name = "once_cell" +version = "1.21.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9f7c3e4beb33f85d45ae3e3a1792185706c8e16d043238c593331cc7cd313b50" + +[[package]] +name = "oxide-artifacts" +version = "0.2.1" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "object", +] + +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" +dependencies = [ + "serde", + "serde_json", + "thiserror 2.0.21", +] + +[[package]] +name = "parking_lot_core" +version = "0.9.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2621685985a2ebf1c516881c026032ac7deafcda1a2c9b7850dc81e3dfcb64c1" +dependencies = [ + "cfg-if", + "libc", + "redox_syscall", + "smallvec", + "windows-link", +] + +[[package]] +name = "pin-project-lite" +version = "0.2.17" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a89322df9ebe1c1578d689c92318e070967d1042b512afbe49518723f4e6d5cd" + +[[package]] +name = "prettyplease" +version = "0.2.37" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "479ca8adacdd7ce8f1fb39ce9ecccbfe93a3f1344b3d0d97f20bc0196208f62b" +dependencies = [ + "proc-macro2", + "syn 2.0.119", +] + +[[package]] +name = "proc-macro2" +version = "1.0.107" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "985e7ec9bb745e6ce6535b544d84d6cd6f7ad8bd711c398938ae983b91a766d9" +dependencies = [ + "unicode-ident", +] + +[[package]] +name = "ptx-parse" +version = "0.1.0" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" + +[[package]] +name = "quote" +version = "1.0.47" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1fbf4db142a473a8d80c26bbf18454ed458bf8d26c8219c331daecfdbd079001" +dependencies = [ + "proc-macro2", +] + +[[package]] +name = "redox_syscall" +version = "0.5.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ed2bf2547551a7053d6fdfafda3f938979645c44812fbfcda098faae3f1a362d" +dependencies = [ + "bitflags", +] + +[[package]] +name = "regex" +version = "1.13.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f020237b6c8eed93db2e2cb53c00c60a8e1bc73da7d073199a1180401450218d" +dependencies = [ + "aho-corasick", + "memchr", + "regex-automata", + "regex-syntax", +] + +[[package]] +name = "regex-automata" +version = "0.4.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ad8553b9b26413251cbf30e620595c7a41b3887f03da04579c0e6b0d6a06b4b2" +dependencies = [ + "aho-corasick", + "memchr", + "regex-syntax", +] + +[[package]] +name = "regex-syntax" +version = "0.8.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d6f6ff9a378485b298a5286656da665ba74413d36db0979633275d2e708145d4" + +[[package]] +name = "reserved-oxide-symbols" +version = "0.2.1" +source = "git+https://github.com/tpn/cuda-rust.git?rev=96caed46c776eb49226137eac849e795b83b7ce0#96caed46c776eb49226137eac849e795b83b7ce0" + +[[package]] +name = "rustc-hash" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "08d43f7aa6b08d49f382cde6a7982047c3426db949b1424bc4b7ec9ae12c6ce2" + +[[package]] +name = "rustc-hash" +version = "2.1.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6b1e7f9a428571be2dc5bc0505c13fb6bf936822b894ec87abf8a08a4e51742d" + +[[package]] +name = "rustix" +version = "0.38.44" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fdb5bc1ae2baa591800df16c9ca78619bf65c0488b41b96ccec5d11220d8c154" +dependencies = [ + "bitflags", + "errno", + "libc", + "linux-raw-sys", + "windows-sys 0.59.0", +] + +[[package]] +name = "scopeguard" +version = "1.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "94143f37725109f92c262ed2cf5e59bce7498c01bcc1502d7b9afe439a4e9f49" + +[[package]] +name = "serde" +version = "1.0.229" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4148590afebada386688f18773da617792bf2ef03ffc1e4cbd2b1d45b023e0ba" +dependencies = [ + "serde_core", + "serde_derive", +] + +[[package]] +name = "serde_core" +version = "1.0.229" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "67dca2c9c51e58a4791a4b1ed58308b39c64224d349a935ab5039aa360942a48" +dependencies = [ + "serde_derive", +] + +[[package]] +name = "serde_derive" +version = "1.0.229" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e7a5d71263a5a7d47b41f6b3f06ba276f10cc18b0931f1799f710578e2309348" +dependencies = [ + "proc-macro2", + "quote", + "syn 3.0.9", +] + +[[package]] +name = "serde_json" +version = "1.0.154" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e7e9cc8b1b85264074fbcc02a88680c4096b1e47df8f739dceb03bf482f04bd6" +dependencies = [ + "itoa", + "memchr", + "serde", + "serde_core", + "zmij", +] + +[[package]] +name = "sha2" +version = "0.10.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a7507d819769d01a365ab707794a4084392c824f54a7a6a7862f8c3d0892b283" +dependencies = [ + "cfg-if", + "cpufeatures", + "digest", +] + +[[package]] +name = "shlex" +version = "1.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0fda2ff0d084019ba4d7c6f371c95d8fd75ce3524c3cb8fb653a3023f6323e64" + +[[package]] +name = "slab" +version = "0.4.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0c790de23124f9ab44544d7ac05d60440adc586479ce501c1d6d7da3cd8c9cf5" + +[[package]] +name = "smallvec" +version = "1.16.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5b3dc8af474f516a851ff4bd12db780f948b9250ad37211e4eec0bccea54e01b" + +[[package]] +name = "syn" +version = "2.0.119" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "872831b642d1a07999a962a351ed35b955ea2cfc8f3862091e2a240a84f17297" +dependencies = [ + "proc-macro2", + "quote", + "unicode-ident", +] + +[[package]] +name = "syn" +version = "3.0.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d78c8dee4c7bf0e14673097256fed6142ce9d3b85a408189d07482442145823b" +dependencies = [ + "proc-macro2", + "quote", + "unicode-ident", +] + +[[package]] +name = "thiserror" +version = "1.0.69" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b6aaf5339b578ea85b50e080feb250a3e8ae8cfcdff9a461c9ec2904bc923f52" +dependencies = [ + "thiserror-impl 1.0.69", +] + +[[package]] +name = "thiserror" +version = "2.0.21" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "09e52cb86a36cede5cb101bf8908837b3e4c6e5e59fe7fd85c23fb56200d189e" +dependencies = [ + "thiserror-impl 2.0.21", +] + +[[package]] +name = "thiserror-impl" +version = "1.0.69" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4fee6c4efc90059e10f81e6d42c60a18f76588c3d74cb83a0b242a2b6c7504c1" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", +] + +[[package]] +name = "thiserror-impl" +version = "2.0.21" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fe5197923287db20a58125f0bc85c062f7f2c892de97b18c356f9efb14b28524" +dependencies = [ + "proc-macro2", + "quote", + "syn 3.0.9", +] + +[[package]] +name = "typenum" +version = "1.20.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b6f5e870be6c3b371b77fe0ee0bafb859fa4964b4404c27de1d380043c4dda20" + +[[package]] +name = "unicode-ident" +version = "1.0.27" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a2c754d6c33795a1c324727428e5a7dedb5b06195f9890bdbcba760d3e246563" + +[[package]] +name = "version_check" +version = "0.9.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0b928f33d975fc6ad9f86c8f283853ad26bdd5b10b7f1542aa2fa15e2289105a" + +[[package]] +name = "which" +version = "4.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "87ba24419a2078cd2b0f2ede2691b6c66d8e47836da3b6db8265ebad47afbfc7" +dependencies = [ + "either", + "home", + "once_cell", + "rustix", +] + +[[package]] +name = "windows-link" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f0805222e57f7521d6a62e36fa9163bc891acd422f971defe97d64e70d0a4fe5" + +[[package]] +name = "windows-sys" +version = "0.59.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1e38bc4d79ed67fd075bcc251a1c39b32a1776bbe92e5bef1f0bf1f8c531853b" +dependencies = [ + "windows-targets", +] + +[[package]] +name = "windows-sys" +version = "0.61.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ae137229bcbd6cdf0f7b80a31df61766145077ddf49416a728b02cb3921ff3fc" +dependencies = [ + "windows-link", +] + +[[package]] +name = "windows-targets" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9b724f72796e036ab90c1021d4780d4d3d648aca59e491e6b98e725b84e99973" +dependencies = [ + "windows_aarch64_gnullvm", + "windows_aarch64_msvc", + "windows_i686_gnu", + "windows_i686_gnullvm", + "windows_i686_msvc", + "windows_x86_64_gnu", + "windows_x86_64_gnullvm", + "windows_x86_64_msvc", +] + +[[package]] +name = "windows_aarch64_gnullvm" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "32a4622180e7a0ec044bb555404c800bc9fd9ec262ec147edd5989ccd0c02cd3" + +[[package]] +name = "windows_aarch64_msvc" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "09ec2a7bb152e2252b53fa7803150007879548bc709c039df7627cabbd05d469" + +[[package]] +name = "windows_i686_gnu" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8e9b5ad5ab802e97eb8e295ac6720e509ee4c243f69d781394014ebfe8bbfa0b" + +[[package]] +name = "windows_i686_gnullvm" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0eee52d38c090b3caa76c563b86c3a4bd71ef1a819287c19d586d7334ae8ed66" + +[[package]] +name = "windows_i686_msvc" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "240948bc05c5e7c6dabba28bf89d89ffce3e303022809e73deaefe4f6ec56c66" + +[[package]] +name = "windows_x86_64_gnu" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "147a5c80aabfbf0c7d901cb5895d1de30ef2907eb21fbbab29ca94c5b08b1a78" + +[[package]] +name = "windows_x86_64_gnullvm" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "24d5b23dc417412679681396f2b49f3de8c1473deb516bd34410872eff51ed0d" + +[[package]] +name = "windows_x86_64_msvc" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "589f6da84c646204747d1270a2a5661ea66ed1cced2631d546fdfb155959f9ec" + +[[package]] +name = "zerocopy" +version = "0.8.63" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e5fe1f8f1b06191a00962174c61aa5005e0bb391a6d80d07e24d115c01a92ed8" +dependencies = [ + "zerocopy-derive", +] + +[[package]] +name = "zerocopy-derive" +version = "0.8.63" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "863ad3ac83293fb4d740aedbfdc9240dd8d1a50c1099acd76ce80ce7c7230c7f" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", +] + +[[package]] +name = "zmij" +version = "1.0.23" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "29666d0abbfad1e3dc4dcf6144730dd3a3ab225bbbdac83319345b1b44ccfc1b" diff --git a/rust/cuda-coop/tests/shuffle/Cargo.toml b/rust/cuda-coop/tests/shuffle/Cargo.toml new file mode 100644 index 000000000000..b33b1084556a --- /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 = "96caed46c776eb49226137eac849e795b83b7ce0" } +cuda-host = { git = "https://github.com/tpn/cuda-rust.git", rev = "96caed46c776eb49226137eac849e795b83b7ce0", features = ["async"] } +cuda-core = { git = "https://github.com/tpn/cuda-rust.git", rev = "96caed46c776eb49226137eac849e795b83b7ce0" } +cuda-async = { git = "https://github.com/tpn/cuda-rust.git", rev = "96caed46c776eb49226137eac849e795b83b7ce0" } diff --git a/rust/cuda-coop/tests/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"