[Rust] Support uniform i32 captures in Block Reduce

.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/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.

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.