Support multidimensional Rust Load and Store blocks

.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

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

@@ -0,0 +1,99 @@
+.. 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
+``nightly-2026-10-11`` toolchain. The device dependency is pinned in Cargo; there
+is no backend-selection feature. 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. Cluster queries currently require sm_90 or
+newer; an ordinary, unclustered launch has one block per cluster. 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. Synchronization,
+temporary storage and collective primitives are separate work.
+
+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.

docs/rust/index.rst

@@ -0,0 +1,151 @@
+.. Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
+.. SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+Rust 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, and Direct/Transpose block
+Load/Store; see :doc:`payload`, :doc:`groups`, and :doc:`load_store`. Its private ``cuda-coop-codegen`` crate
+specializes CUB block operations and returns
+linkable device LTO-IR with exact temporary-storage requirements. It supports
+compiler integration development. Additional collective algorithms and public
+scratch storage remain under development. The crate is unpublished and its interface can change.
+
+Device code uses the pinned Oxide dependency and ``nightly-2026-10-11``.
+The host provider supports Rust 1.85 independently; validate it with
+``cargo +1.85.0 test -p cuda-coop-codegen --locked`` from ``rust/``.
+
+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 describes a complete block with exact ``[X,Y,Z]`` 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. 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 blocked order
+     - 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
+
+Both Direct and Transpose algorithms are supported. Direct accesses no shared
+scratch. Transpose Store redistributes blocked payloads across lanes into striped
+order, retaining initialized values in every slot, and uses the exact CUB storage
+layout. Compiler integrations must enforce collective participation, memory
+ownership, valid bounds, and synchronization before scratch reuse.
+
+The separate validity-aware Direct and Transpose recipes support all ten numeric payload
+types: ``i8``, ``u8``, ``i16``, ``u16``, ``i32``, ``u32``, ``i64``, ``u64``,
+``f32`` and ``f64``. ``Request::new_typed`` selects the exact type. These recipes
+take opaque value storage and a packed validity bitmap. Full and
+default-filled Load initialize every slot. Masked Load without a default copies
+only the loaded prefix and resets all other flags to zero. Full and partial
+Store require initialized slots in their write domain. Direct preserves all
+payload bytes and validity bits; Transpose permutes both into striped order,
+including slots outside its write prefix. A failing Store traps; it is not transactional
+across lanes. Partial Store forms destination pointers only within its valid
+prefix, including when the allocation is shorter than a complete tile.
+
+The public Direct and Transpose API uses
+the same ``ThreadData`` representation without exposing an initialized slice or
+a public native layout.
+
+Compile a provider
+------------------
+
+Use the pinned toolchain in ``rust/rust-toolchain.toml``. 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 Direct/Transpose forms at
+item counts 1 and 4, 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.
+
+.. toctree::
+   :maxdepth: 1
+
+   payload
+   groups
+   load_store