Add Rust cooperative Exchange and Shuffle

.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

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;

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 <cub/util_device.cuh>
 #include <cub/util_macro.cuh>
 #include <cub/warp/warp_reduce.cuh>
 #include <cub/warp/warp_reduce_batched.cuh>
@@ -9,6 +10,8 @@
 #include <cuda/functional>
 #include <cuda/iterator>
 #include <cuda/ptx>
+#include <cuda/std/array>
+#include <cuda/std/cstddef>
 #include <cuda/std/functional>
 #include <cuda/std/limits>
 #include <cuda/std/mdspan>
@@ -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<int, 3, warp_size, false>;
+  __shared__ warp_reduce_batched_t::TempStorage temp_storage[block_size / warp_size];
+  const int tid             = static_cast<int>(threadIdx.x);
+  const int logical_warp_id = tid / warp_size;
+  const cuda::std::array<int, 3> inputs{tid, tid + 1, tid + 2};
+  int assignments = 0;
+  cuda::std::mdspan<int, cuda::std::extents<int, 1>, 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<int> 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<int> 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 <WarpReduceBatchedMode Mode,
           bool SyncPhysicalWarp,
           int Batches,

docs/index.rst

@@ -7,6 +7,7 @@ CUDA Core Compute Libraries
 
    cccl/index
    python/index
+   rust/index
    contributors/index
    maintainers/index
    infrastructure/index
@@ -18,6 +19,8 @@ make CUDA C++ and Python more delightful.
 
 - :doc:`Python Libraries <python/index>`
 
+- :doc:`Rust Development <rust/index>`
+
 - :doc:`Contributor Docs <contributors/index>`
 
 - :doc:`Maintainer Docs <maintainers/index>`

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.

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::<coop::Warp>()`` or
+``count_at::<coop::Warp>()``. 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::<coop::Grid>()`` 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::<i16>()`` and ``count_as::<i16>()`` select one of the eight fixed-width
+integer types; ``rank_at_as::<i16, coop::Warp>()`` 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::<N>()`` 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.